Daniel Voigt Godoy - Deep Learning with PyTorch Step-by-Step A Beginner’s Guide-leanpub

BatchNorm2d
The difference between the one-dimension and the two-dimension batch
normalization is actually quite simple: The former standardizes features (columns),
while the latter standardizes channels (pixels).
This is easier to see in code:
torch.manual_seed(39)
dummy_images = torch.rand((200, 3, 10, 10))
dummy_labels = torch.randint(2, (200, 1))
dummy_dataset = TensorDataset(dummy_images, dummy_labels)
dummy_loader = DataLoader(
dummy_dataset, batch_size=64, shuffle=True
)
iterator = iter(dummy_loader)
batch1 = next(iterator)
batch1[0].shape
Output
torch.Size([64, 3, 10, 10])
The code above creates a dummy dataset of 200 colored (three-channel) images of
size 10x10 pixels and then retrieves the first mini-batch. The mini-batch has the
expected NCHW shape.
The batch normalization is done over the C dimension, so it will compute statistics
using the remaining dimensions—N, H, and W (axis=[0, 2, 3])—representing all
pixels of a given channel from every image in the mini-batch.
The nn.BatchNorm2d layer has the same arguments as its one-dimension
counterpart, but its num_features argument must match the number of channels
of the input instead:
544 | Chapter 7: Transfer Learning