Нормализация батча в PyTorch

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

Создадим слой на три канала, явно включим режим обучения и выведем форму результата:

import torch
import torch.nn as nn

batch = torch.randn(2, 3, 8, 8)
norm = nn.BatchNorm2d(3)
norm.train()
result = norm(batch)
print(result.shape) # выведет torch.Size([2, 3, 8, 8])

В режиме оценки используют накопленную статистику. Вызовем метод eval и снова прогоним тот же батч:

import torch
import torch.nn as nn

batch = torch.randn(2, 3, 8, 8)
norm = nn.BatchNorm2d(3)
norm.eval()
result = norm(batch)
print(result.shape) # выведет torch.Size([2, 3, 8, 8])

Соберите пакет 4 на 8 на 16 на 16 из единиц и слой нормализации на 8 каналов. Включите режим обучения, прогоните батч и выведите размеры результата.

Подайте случайный тензор 1 на 3 на 5 на 5 через нормализацию на 3 канала в режиме оценки. Выведите размеры на выходе.

Опишите выравнивание для 16 каналов. Переключите его в режим обучения, обработайте нулевой батч 2 на 16 на 8 на 8 и выведите размеры полученной карты.