Don't Mock Machine Learning Models In Unit Tests

Начните здесь Статьи Выступления Прототипы Об авторе

Don't Mock Machine Learning Models In Unit Tests

Не мокайте ML-модели в юнит-тестах

[ machinelearning engineering python ] · 6 min read

[ machinelearning engineering python ] · 6 мин чтения

I’ve been applying typical unit testing practices to machine learning code and it hasn’t been straightforward. In software, units are small, isolated pieces of logic that we can test independently and quickly. In machine learning, models are blobs of logic learned from data, and machine learning code is the logic to learn and use these derived blobs of logic. This difference makes it necessary to rethink how we unit test machine learning code.

Я пробовал применять типичные практики юнит-тестирования к ML-коду, и это оказалось непросто. В обычном софте «юниты» — это маленькие изолированные куски логики, которые можно тестировать независимо и быстро. В ML же модели — это «комки» логики, выученные из данных, а ML-код — это логика, которая обучает и использует эти производные комки логики. Из-за такого различия приходится переосмысливать, как именно стоит юнит-тестировать ML-код.

How ML code differs from regular software

Чем ML-код отличается от обычного софта

In software, we write code that contains logic; in ML, we write code that learns logic and then uses that learned logic. Software code transforms input data + handcrafted logic into expected output. We can then test these outputs against asserts. In contrast, machine learning code transforms input data + expected output into learned logic (i.e., a model).

В обычном софте мы пишем код, который содержит логику; в ML мы пишем код, который выучивает логику, а потом её использует. Софтовый код превращает входные данные + написанную вручную логику в ожидаемый выход. Этот выход потом можно проверять ассертами. В ML же код превращает входные данные + ожидаемый выход в выученную логику (то есть в модель).

Thus, in machine learning, instead of writing code that contains logic, we write code to learn logic, such as via building a decision tree or finetuning a hallucination classifier. Because the logic that acts on the input data is embedded within the model, if we want to test the learned logic, we’ll need to load the model, perform inference on some sample output, and then assert if the output matches the expected input.

Иначе говоря, в машинном обучении вместо написания кода, содержащего логику, мы пишем код, который её выучивает — например, строит дерево решений или дообучает классификатор галлюцинаций. Поскольку логика, действующая на входные данные, зашита внутрь модели, для тестирования выученной логики придётся загрузить модель, прогнать инференс на каком-то сэмпле и затем проверить, совпадает ли выход с ожидаемым.

In software, we typically mock dependencies like APIs; in ML, we want to test the actual model (sometimes). When unit testing software, it’s good practice to mock database calls, filesystem access, sending emails/push notifications, etc. However, in ML, there are scenarios where we’ll want to test against the actual model.

В обычном софте мы обычно мокаем зависимости вроде API; в ML мы хотим (иногда) тестировать саму настоящую модель. При юнит-тестировании обычного софта считается хорошим тоном мокать обращения к базе, файловой системе, отправку писем и пуш-уведомлений и т.п. В ML же есть сценарии, где мы хотим прогнать тест именно против реальной модели.

For example, we want to test that loss decreases with each batch and the model can overfit (before wasting compute on an hopeless run.) If the model is a classifier, we want to check that the inference logic is correct. For instance, two models may have different output classes: Google’s T5 NLI model classifies factual consistency with class = 1 while Meta’s BART NLI model classifies it with class = 2!

Например, мы хотим проверить, что loss падает с каждым батчем и что модель способна переобучиться (прежде чем тратить compute на безнадёжный прогон). Если модель — классификатор, нам важно убедиться, что логика инференса корректна. Скажем, у двух моделей могут быть разные выходные классы: T5 NLI от Google размечает фактическую согласованность как class = 1, а BART NLI от Meta — как class = 2!

Machine learning / language models can be large and unwieldy. Some neural networks can be in the billions of parameters, exceeding what a laptop or standard dev environment can load. And even if we have the memory for smaller models, they are slow to load and perform inference on, testing our patience as we unit test while coding.

ML- и языковые модели могут быть большими и неподъёмными. Некоторые нейросети содержат миллиарды параметров и попросту не влезут на ноутбук или в стандартное dev-окружение. Но даже если памяти под модель поменьше хватает, такие модели медленно грузятся и медленно делают инференс — терпение быстро заканчивается, когда юнит-тесты гоняются по ходу разработки.

Some guidelines for unit testing ML code & models

Несколько ориентиров для юнит-тестирования ML-кода и моделей

(These are a work in progress and my thinking’s still evolving—all feedback welcome!)

(Это рабочие наброски, моё мышление ещё эволюционирует — любой фидбэк приветствуется!)

Use small, simple data samples. Avoid loading CSVs or Parquet files as sample data. (It’s fine for integration tests and evals but not unit tests.) Define sample data directly in unit test code—so that the test is self-contained—to test key functionality such as:

Используйте маленькие, простые сэмплы данных. Не загружайте CSV или Parquet в качестве тестовых данных. (Для интеграционных тестов и evals это норм, но не для юнит-тестов.) Описывайте сэмпл прямо в коде теста — так тест остаётся self-contained — чтобы проверять ключевую функциональность, например:

Splitting into train/test tests when you have custom logic

Custom implementations, such as Cosine or Euclidean distance in Java

Preprocessing such as data augmentation or encoding

Postprocessing such as diversification or filtering recommendations

Error handling for empty or malformed input

Сплиты train/test, если у вас своя кастомная логика. Кастомные реализации — например, косинусное или евклидово расстояние на Java. Препроцессинг — например, аугментация данных или энкодинг. Постпроцессинг — например, диверсификация или фильтрация рекомендаций. Обработка ошибок при пустом или некорректном входе.

When viable, test against random or empty weights. For example, we can initialize a model configuration with random weights to test output shape and device movement (from CPU to GPU and back). Here’s an example of how to initialize a model without having to download the weights and then assert the output shape:

Где это уместно, тестируйте против случайных или пустых весов. Например, можно инициализировать конфиг модели со случайными весами, чтобы проверить форму выхода и перемещение между устройствами (с CPU на GPU и обратно). Вот пример, как инициализировать модель без скачивания весов и затем проверить форму выхода:

from transformers import AutoConfig, AutoModelForSequenceClassification model_name = "valhalla/distilbart-mnli-12-1" config = AutoConfig.from_pretrained("valhalla/distilbart-mnli-12-1") model = AutoModelForSequenceClassification.from_config(config) assert model.classification_head.out_proj.out_features == 3

The accelerate library also has an example of initializing a model with empty weights:

В библиотеке accelerate тоже есть пример инициализации модели с пустыми весами:

def test_dispatch_model_bnb(self): """Tests that `dispatch_model` quantizes int8 layers""" from huggingface_hub import hf_hub_download from transformers import AutoConfig, AutoModel, BitsAndBytesConfig from transformers.utils.bitsandbytes import replace_with_bnb_linear with init_empty_weights(): model = AutoModel.from_config(AutoConfig.from_pretrained("bigscience/bloom-560m")) quantization_config = BitsAndBytesConfig(load_in_8bit=True) model = replace_with_bnb_linear( model, modules_to_not_convert=["lm_head"], quantization_config=quantization_config ) model_path = hf_hub_download("bigscience/bloom-560m", "pytorch_model.bin") model = load_checkpoint_and_dispatch( model, checkpoint=model_path, device_map="balanced", ) assert model.h[0].self_attention.query_key_value.weight.dtype == torch.int8 assert model.h[0].self_attention.query_key_value.weight.device.index == 0 assert model.h[(-1)].self_attention.query_key_value.weight.dtype == torch.int8 assert model.h[(-1)].self_attention.query_key_value.weight.device.index == 1

Write critical tests against the actual model. If they take a while to run, mark them as slow and run only when needed (e.g., pre-commit and pre-merge). Some essentials include:

Критически важные тесты пишите против настоящей модели. Если они идут долго, помечайте их как slow и запускайте только по необходимости (например, на pre-commit и pre-merge). Что обязательно стоит проверять:

Verify training is done correctly, such as loss going down, model overfitting, and training till convergence on a small sample of data

Verify model outputs match expectation, such as 0.99 = unsafe instead of safe

Verify model server can start, take batch input, and return the expected output

Что обучение идёт корректно: loss падает, модель переобучается на маленькой выборке и обучение сходится. Что выходы модели совпадают с ожиданием — например, 0.99 = unsafe, а не safe. Что сервер инференса стартует, принимает батч на вход и возвращает ожидаемый выход.

Don’t test external libraries. We can assume that external libraries work. Thus, no need to test data loaders, tokenizers, optimizers, etc.

Не тестируйте внешние библиотеки. Можно считать, что внешние библиотеки работают. То есть нет смысла тестировать data loaders, токенизаторы, оптимизаторы и т.п.

• • •

What are your best practices for unit testing machine learning code and models? I would love to hear from you. Please reach out!

А какие у вас best practices для юнит-тестирования ML-кода и моделей? Очень интересно услышать. Пишите!

Что почитать дальше

How to Test Machine Learning Code and Systems

Writing Robust Tests for Data & Machine Learning Pipelines

Effective Testing for Machine Learning Systems

How to Trust Your Deep Learning Code

A Recipe for Training Neural Networks

How to Unit Test Deep Learning

Testing Data Science and MLOps Code

How to Test Machine Learning Code and Systems Writing Robust Tests for Data & Machine Learning Pipelines Effective Testing for Machine Learning Systems How to Trust Your Deep Learning Code A Recipe for Training Neural Networks How to Unit Test Deep Learning Testing Data Science and MLOps Code

Notable/fun comments from HackerNews

Заметные/забавные комментарии с HackerNews

The problem when mocks happen when all your unit test passes and the program fails on integration. The mocks are a pristine place where your library unit test works like a champ. Bad mocks or bad library? Or both. Developers are then sent to debug the unit test… overhead.

Проблема с моками всплывает, когда все ваши юнит-тесты зелёные, а программа падает на интеграции. Моки — это стерильное место, где ваш юнит-тест библиотеки работает на ура. Плохие моки или плохая библиотека? Или и то и другое. И вот разработчиков отправляют дебажить юнит-тест… накладные расходы.

Probablistic tests create an impossible problem: if you tighten your assertions you struggle with meaningless failing tests and normalize ignoring test failures, while if you loosen your assertions your tests aren’t really asserting anything any more. And there isn’t a happy balance: if you go somewhere in the middle, you end up having both problems.

Вероятностные тесты создают неразрешимую проблему: если ужесточить ассерты, вы тонете в бессмысленных падениях и привыкаете игнорировать упавшие тесты; а если ослабить ассерты — тесты по сути перестают что-либо проверять. И счастливой середины тут нет: если расположиться где-то посередине, у вас будут обе проблемы сразу.

How does it matter whether I inline my test data inside the unit test code, or have my unit test code load that same data from a checked-in file instead?

It makes tests self-contained and easier to reason about. As a side-effect, random tests won’t accidentally break whenever you change some seemingly unrelated csv file. As a rule of thumb, I also only assert on input/output values that are explicitly defined as part of the test body. Saves a ton of time chasing down fixture definitions.

Какая разница, инлайнить ли тестовые данные внутрь кода юнит-теста или загружать те же самые данные из закоммиченного файла? А такая, что тесты становятся self-contained и их проще понимать. Как побочный эффект, случайные тесты не будут внезапно ломаться, когда вы поменяете какой-то на вид не связанный csv. Как правило, я также делаю ассерты только на тех значениях входа/выхода, которые явно описаны в теле теста. Это экономит уйму времени на гонку за определениями фикстур.

I was expecting an article about side effects of hurting an LLM’s feelings in tests.

А я-то ожидал статью про побочные эффекты от того, что в тестах задевают чувства LLM.

If you found this useful, please cite this write-up as:

Если статья оказалась полезной, ссылайтесь на неё так:

Yan, Ziyou. (Feb 2024). Don't Mock Machine Learning Models In Unit Tests. eugeneyan.com. https://eugeneyan.com/writing/unit-testing-ml/.

or

или

@article{yan2024unit, title = {Don't Mock Machine Learning Models In Unit Tests}, author = {Yan, Ziyou}, journal = {eugeneyan.com}, year = {2024}, month = {Feb}, url = {https://eugeneyan.com/writing/unit-testing-ml/} }