Batch normalization standardizes activations channel by channel, helping neural networks train more stably and quickly. In 2D convolutional networks, it is especially useful for image models and generative networks because it keeps inputs to deeper layers in a more appropriate range. section 4.3 section 4.3.2
For each feature channel, the procedure follows a short sequence:
- Calculate the channel's mean and variance across the observations. section 4.3.2
- Subtract the mean from each value and divide by the standard deviation. section 4.3.2
- Produce values with approximately mean 0 and standard deviation 1, stabilizing gradients during backpropagation. section 4.3.2
Unlike an unnormalized convolutional output, whose channels can have different means and spreads, a normalized output has a more consistent distribution. This can reduce internal covariate shift and help address vanishing or exploding gradients, which otherwise cause slow, unstable, or ineffective updates. section 4.3 section 4.3.2
Normalizing a three-channel image tensor
Create a random tensor with shape [1, 3, 64, 64], pass it through Conv2d(3, 3, 3, stride=1, padding=1), and inspect each output channel. Before normalization, the example produces channel means of about -0.377, -0.391, and 0.393, with standard deviations near 0.18. Applying nn.BatchNorm2d(3) preserves the shape while moving each channel's mean close to 0 and standard deviation close to 1.
Listing 4.4 A discriminator in DCGAN
import torch.nn as nn
import torch
device = "cuda" if torch.cuda.is_available() else "cpu"
D = nn.Sequential(
nn.Conv2d(3, 64, 4, 2, 1, bias=False), #1
nn.LeakyReLU(0.2, inplace=True), #2
nn.Conv2d(64, 128, 4, 2, 1, bias=False),
nn.BatchNorm2d(128), #3
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(128, 256, 4, 2, 1, bias=False),
nn.BatchNorm2d(256),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(256, 512, 4, 2, 1, bias=False),
nn.BatchNorm2d(512),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(512, 1, 4, 1, 0, bias=False),
nn.Sigmoid(),
nn.Flatten()).to(device) #4
Use BatchNorm2d after convolutional layers in a DCGAN discriminator; the listed model applies it from the second convolution onward, before LeakyReLU. section 4.5.1
Listing 4.5 Designing a generator in DCGAN
G=nn.Sequential(
nn.ConvTranspose2d(100, 512, 4, 1, 0, bias=False), #1
nn.BatchNorm2d(512),
nn.ReLU(inplace=True),
nn.ConvTranspose2d(512, 256, 4, 2, 1, bias=False), #2
nn.BatchNorm2d(256),
nn.ReLU(inplace=True),
nn.ConvTranspose2d(256, 128, 4, 2, 1, bias=False),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.ConvTranspose2d(128, 64, 4, 2, 1, bias=False),
nn.BatchNorm2d(64),
nn.ReLU(inplace=True),
nn.ConvTranspose2d(64, 3, 4, 2, 1, bias=False), #3
nn.Tanh()).to(device) #4
Use BatchNorm2d between transposed convolution and ReLU in the generator, supporting progressively upsampled image features. section 4.5.1
