When a linear model is too limited, torch.nn provides the building blocks for a neural network. Its layers and other components let you assemble a model; training still needs data and an optimizer. section 1.5 section 6.0
These building blocks are called modules. A module derives from nn.Module and can hold nn.Parameter tensors whose values are optimized during training. It can also hold submodules and track their parameters. section 6.2
A few useful pieces of the namespace are:
- nn.Linear applies an affine transformation using weight and bias parameters—like a simple linear model. section 6.2
- nn.Conv1d, nn.Conv2d, and nn.Conv3d provide convolutions for time series, images, and volumes or videos, respectively. section 8.2
- Keep submodules as module attributes so their parameters can be found. For a list or dictionary of submodules, use nn.ModuleList or nn.ModuleDict. section 6.2
A convolution for an RGB image
Use `conv = nn.Conv2d(3, 16, kernel_size=3)` for three input channels, 16 output channels, and a 3 × 3 kernel. Its weight shape is `[16, 3, 3, 3]`, and its bias shape is `[16]`. section 8.2
# In[9]: img, _ = cifar2[0] output = conv(img.unsqueeze(0)) img.unsqueeze(0).shape, output.shape # Out[9]: (torch.Size([1, 3, 32, 32]), torch.Size([1, 16, 30, 30]))
The listing shows a single image gaining a batch dimension before convolution: input shape `[1, 3, 32, 32]` becomes output shape `[1, 16, 30, 30]`. section 8.2
