Используются ли статистики из батча или сохраненные параметры в батч-нормализации на инференсе

Что проверяют

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

Как отвечать

Короткий ответ: на инференсе используются сохранённые статистики, а не статистики батча. Дальше разверните механику:

• обучение: слой нормализует вход средним и дисперсией текущего мини-батча по каждому каналу и параллельно обновляет скользящие оценки running_mean и running_var экспоненциальным сглаживанием;
• инференс: вместо статистик батча берутся накопленные running_mean и running_var, поэтому ответ модели на объект не зависит от того, с какими объектами он попал в батч;
• обучаемые параметры: масштаб gamma и сдвиг beta (в PyTorch это weight и bias) применяются в обоих режимах, путать их со статистиками нельзя;
• переключение: в PyTorch это model.eval() и model.train(), в Keras аргумент training=False при вызове. У PyTorch по умолчанию momentum=0.1, у Keras momentum=0.99, и смысл у них противоположный: в Keras это вес старого значения.

Хорошо добавить инженерную сторону. В режиме инференса слой становится обычным линейным преобразованием, поэтому при экспорте в ONNX или TensorRT его сливают с предыдущей свёрткой, и отдельной операции не остаётся. Если распределение данных в проде заметно ушло от обучающего, скользящие статистики можно пересчитать на свежих данных, не трогая веса.

Где спотыкаются

Говорят, что сохранённые статистики посчитаны точно по всему обучающему набору. На деле это скользящее среднее по батчам последних шагов обучения, и после резкой смены данных оно отстаёт.

Забывают model.eval() перед валидацией или в сервисе. Предсказание начинает зависеть от состава батча, а BatchNorm1d на батче из одного вектора в режиме обучения просто падает с ошибкой.

При дообучении с батчами из двух-четырёх объектов оставляют батч-нормализацию в режиме обучения. Статистики становятся шумными, и лучше заморозить слои BN или перейти на GroupNorm и LayerNorm, которым размер батча не важен.