Размер батча на инференсе в PyTorch: влияет ли он на результаты

Разбираем, влияет ли размер батча на результаты инференса в PyTorch, можно ли запускать весь датасет одним батчем и как выбрать bsz без потери точности.

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

Пример рассматриваемого паттерна инференса

Ниже показан подход, который загружает все объекты в один батч и выполняет один прямой проход:

def infer_once(self, inputs):
    ds_loader = DataLoader(
        inputs,
        batch_size=inputs.shape[0],
        shuffle=False,
    )
    minibatch = next(iter(ds_loader))
    with torch.no_grad():
        outputs = self.net(minibatch)
    return outputs

Что на самом деле меняется с размером батча

Обучение и инференс используют батчи по разным причинам. Во время обучения батчи приближают распределение данных в целом, чтобы градиент, вычисленный обратным распространением, хорошо совпадал с истинным градиентом. Поэтому размер батча взаимодействует с другими настройками оптимизации и часто подбирается. Во время инференса цель — просто применить уже выученные параметры к новым объектам. Батчирование здесь — в основном про параллелизм и пропускную способность: группировка примеров позволяет эффективнее задействовать ускоритель, но не меняет функцию, которую модель вычисляет для каждого элемента.

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

Практический вывод и как организовать инференс

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

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

def infer_batched(self, inputs, bsz=None):
    size = inputs.shape[0] if bsz is None else bsz
    ds_loader = DataLoader(
        inputs,
        batch_size=size,
        shuffle=False,
    )
    preds = []
    with torch.no_grad():
        for part in ds_loader:
            preds.append(self.net(part))
    return torch.cat(preds, dim=0)

Влияет ли изменение размера батча на результаты?

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

# Пример sanity‑проверки
scores_1 = model.infer_batched(X, bsz=1)
scores_2 = model.infer_batched(X, bsz=X.shape[0])
ok = torch.allclose(scores_1, scores_2)

Почему это важно

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

Вывод

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