Don't Mock Machine Learning Models In Unit Tests
Юджин Ян объясняет, почему привычные практики юнит-тестирования плохо ложатся на ML-код: в обычном софте логика пишется руками, а в ML она выучивается из данных и зашита в саму модель. Поэтому моки внешних зависимостей здесь часто бесполезны — иногда нужно тестировать именно реальную модель: проверять, что loss падает, что модель способна переобучиться на маленькой выборке, и что инференс возвращает корректные классы (например, у T5 NLI factual consistency = 1, а у BART NLI = 2). Автор предлагает несколько правил: использовать маленькие сэмплы данных прямо в коде теста, а не загружать CSV/Parquet; где возможно — тестировать на случайных или пустых весах через AutoConfig и init_empty_weights, чтобы не качать гигабайты; критические тесты гонять против настоящей модели, помечая медленные через pytest mark и запуская их на pre-commit/pre-merge; и не тестировать внешние библиотеки вроде токенизаторов и оптимизаторов. В конце приведены цитаты с HackerNews о проблемах моков и вероятностных тестов, а также дальнейшее чтение по теме.
Не мокайте ML-модели в юнит-тестах
[ machinelearning engineering python ] · 6 мин чтения
Я пробовал применять типичные практики юнит-тестирования к ML-коду, и это оказалось непросто. В обычном софте «юниты» — это маленькие изолированные куски логики, которые можно тестировать независимо и быстро. В ML же модели — это «комки» логики, выученные из данных, а ML-код — это логика, которая обучает и использует эти производные комки логики. Из-за такого различия приходится переосмысливать, как именно стоит юнит-тестировать ML-код.
Чем ML-код отличается от обычного софта
В обычном софте мы пишем код, который содержит логику; в ML мы пишем код, который выучивает логику, а потом её использует. Софтовый код превращает входные данные + написанную вручную логику в ожидаемый выход. Этот выход потом можно проверять ассертами. В ML же код превращает входные данные + ожидаемый выход в выученную логику (то есть в модель).
Иначе говоря, в машинном обучении вместо написания кода, содержащего логику, мы пишем код, который её выучивает — например, строит дерево решений или дообучает классификатор галлюцинаций. Поскольку логика, действующая на входные данные, зашита внутрь модели, для тестирования выученной логики придётся загрузить модель, прогнать инференс на каком-то сэмпле и затем проверить, совпадает ли выход с ожидаемым.
В обычном софте мы обычно мокаем зависимости вроде API; в ML мы хотим (иногда) тестировать саму настоящую модель. При юнит-тестировании обычного софта считается хорошим тоном мокать обращения к базе, файловой системе, отправку писем и пуш-уведомлений и т.п. В ML же есть сценарии, где мы хотим прогнать тест именно против реальной модели.
Например, мы хотим проверить, что loss падает с каждым батчем и что модель способна переобучиться (прежде чем тратить compute на безнадёжный прогон). Если модель — классификатор, нам важно убедиться, что логика инференса корректна. Скажем, у двух моделей могут быть разные выходные классы: T5 NLI от Google размечает фактическую согласованность как class = 1, а BART NLI от Meta — как class = 2!
ML- и языковые модели могут быть большими и неподъёмными. Некоторые нейросети содержат миллиарды параметров и попросту не влезут на ноутбук или в стандартное dev-окружение. Но даже если памяти под модель поменьше хватает, такие модели медленно грузятся и медленно делают инференс — терпение быстро заканчивается, когда юнит-тесты гоняются по ходу разработки.
Несколько ориентиров для юнит-тестирования ML-кода и моделей
(Это рабочие наброски, моё мышление ещё эволюционирует — любой фидбэк приветствуется!)
Используйте маленькие, простые сэмплы данных. Не загружайте CSV или Parquet в качестве тестовых данных. (Для интеграционных тестов и evals это норм, но не для юнит-тестов.) Описывайте сэмпл прямо в коде теста — так тест остаётся self-contained — чтобы проверять ключевую функциональность, например:
Сплиты train/test, если у вас своя кастомная логика. Кастомные реализации — например, косинусное или евклидово расстояние на Java. Препроцессинг — например, аугментация данных или энкодинг. Постпроцессинг — например, диверсификация или фильтрация рекомендаций. Обработка ошибок при пустом или некорректном входе.
Где это уместно, тестируйте против случайных или пустых весов. Например, можно инициализировать конфиг модели со случайными весами, чтобы проверить форму выхода и перемещение между устройствами (с 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
В библиотеке 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
Критически важные тесты пишите против настоящей модели. Если они идут долго, помечайте их как slow и запускайте только по необходимости (например, на pre-commit и pre-merge). Что обязательно стоит проверять:
Что обучение идёт корректно: loss падает, модель переобучается на маленькой выборке и обучение сходится. Что выходы модели совпадают с ожиданием — например, 0.99 = unsafe, а не safe. Что сервер инференса стартует, принимает батч на вход и возвращает ожидаемый выход.
Не тестируйте внешние библиотеки. Можно считать, что внешние библиотеки работают. То есть нет смысла тестировать data loaders, токенизаторы, оптимизаторы и т.п.
• • •
А какие у вас 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
Заметные/забавные комментарии с HackerNews
Проблема с моками всплывает, когда все ваши юнит-тесты зелёные, а программа падает на интеграции. Моки — это стерильное место, где ваш юнит-тест библиотеки работает на ура. Плохие моки или плохая библиотека? Или и то и другое. И вот разработчиков отправляют дебажить юнит-тест… накладные расходы.
Вероятностные тесты создают неразрешимую проблему: если ужесточить ассерты, вы тонете в бессмысленных падениях и привыкаете игнорировать упавшие тесты; а если ослабить ассерты — тесты по сути перестают что-либо проверять. И счастливой середины тут нет: если расположиться где-то посередине, у вас будут обе проблемы сразу.
Какая разница, инлайнить ли тестовые данные внутрь кода юнит-теста или загружать те же самые данные из закоммиченного файла? А такая, что тесты становятся self-contained и их проще понимать. Как побочный эффект, случайные тесты не будут внезапно ломаться, когда вы поменяете какой-то на вид не связанный csv. Как правило, я также делаю ассерты только на тех значениях входа/выхода, которые явно описаны в теле теста. Это экономит уйму времени на гонку за определениями фикстур.
А я-то ожидал статью про побочные эффекты от того, что в тестах задевают чувства LLM.
Если статья оказалась полезной, ссылайтесь на неё так:
Yan, Ziyou. (Feb 2024). Don't Mock Machine Learning Models In Unit Tests. eugeneyan.com. https://eugeneyan.com/writing/unit-testing-ml/.
или
@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/} }
К рассылке уже присоединились 11 800+ читателей, которые получают апдейты про machine learning, RecSys, LLM и инженерию.