9.5 Convolutional neural networks

The networks of the previous sections treat their inputs as an unordered list of numbers. For an image that is a serious waste, because the arrangement of the pixels carries almost all the information. The convolutional architecture is built to exploit that arrangement, and the three ideas it rests on, locality, weight sharing and stacking, are what made computer vision work.

9.5.1 Why convolution

Suppose we want to classify images and we use the networks of the previous sections. An image must first be flattened into a vector, and two consequences follow immediately, both of them fatal.

The first is the number of parameters. A modest colour image of \(224\times 224\) pixels has \(150\,528\) inputs; connecting it to a single hidden layer of a thousand units requires a hundred and fifty million weights for that layer alone. The second is worse: flattening destroys the geometry. Two pixels that touch each other end up at arbitrary positions in the vector, and the network has no way of knowing that they were neighbours. It would have to learn that from the data, at great cost.

The convolutional layer is built on three ideas that answer these objections:

  • local connectivity: a unit looks only at a small window of the image, because what makes a contour or a corner is local;

  • weight sharing: the same small set of weights is applied at every position, because a contour is a contour wherever it appears. This is where the saving in parameters comes from, and it also gives the layer its equivariance to translation;

  • stacking: the first layer detects contours, the next assembles them into motifs, the next into parts of objects. The hierarchy of features is built by depth.

9.5.2 The convolution operation

The operation is a weighted sum over a sliding window. For a two dimensional input \(X\) and a kernel \(K\) of size \(k\times k\):

\[\begin{equation} S(i,j)=\sum_{u=0}^{k-1}\sum_{v=0}^{k-1}X(i+u,\,j+v)\,K(u,v) \tag{9.19} \end{equation}\]

Three parameters govern the geometry of the output. The stride is the step by which the window advances, and a stride of two halves the size of the output. The padding adds zeros around the border, which allows the output to keep the same size as the input and prevents the borders from being seen less often than the centre. With an input of size \(n\), the output size is:

\[\begin{equation} n_{out}=\left\lfloor\frac{n+2p-k}{s}\right\rfloor+1 \tag{9.20} \end{equation}\]

The essential point is that the kernel has \(k^2\) weights whatever the size of the image. The layer that required a hundred and fifty million weights above needs, with sixty four kernels of size three, \(64\times 3\times 3\times 3=1728\) of them.

Before learning any kernel, it is worth seeing what fixed ones do, because this is what the first layer of a trained network ends up discovering by itself.

In Python:

from scipy.signal import convolve2d

digits = load_digits()
img = digits.images[17]

kernels = {
    "vertical edges":   np.array([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]]),
    "horizontal edges": np.array([[-1, -2, -1], [0, 0, 0], [1, 2, 1]]),
    "blur":             np.ones((3, 3)) / 9.0,
}

fig, axes = plt.subplots(1, 4, figsize=(9, 2.4))
axes[0].imshow(img, cmap="gray_r"); axes[0].set_title("original", fontsize=9)
#> <matplotlib.image.AxesImage object at 0x0000020EB82794C0>
#> Text(0.5, 1.0, 'original')
for ax, (name, K) in zip(axes[1:], kernels.items()):
    ax.imshow(convolve2d(img, K, mode="same"), cmap="gray_r")
    ax.set_title(name, fontsize=9)
#> <matplotlib.image.AxesImage object at 0x0000020EB3A5E180>
#> Text(0.5, 1.0, 'vertical edges')
#> <matplotlib.image.AxesImage object at 0x0000020F2C6921B0>
#> Text(0.5, 1.0, 'horizontal edges')
#> <matplotlib.image.AxesImage object at 0x0000020F2C68D040>
#> Text(0.5, 1.0, 'blur')
for ax in axes: ax.set_xticks([]); ax.set_yticks([])
#> []
#> []
#> []
#> []
#> []
#> []
#> []
#> []
plt.tight_layout(); plt.savefig("dl_kernels_py.png"); plt.clf(); plt.close()
the effect of three fixed kernels

Figure 9.22: the effect of three fixed kernels

The first two kernels are the Sobel operators. One responds to the vertical transitions of intensity, the other to the horizontal ones, and each produces a feature map that lights up where its pattern is present. The third averages and blurs. A convolutional network does not receive these kernels: it discovers comparable ones during the training, because they are useful for the task.

9.5.3 Filters, feature maps and channels

Three words must be kept apart.

A filter, or kernel, is the small matrix of weights. Applying it to the whole image produces one feature map, an image in which each pixel measures the presence of the pattern at that place. A layer holds several filters and therefore produces several feature maps, which are stacked into channels.

An image is itself already multi-channel, three for a colour image. A filter therefore has a depth equal to the number of input channels: a \(3\times 3\) filter applied to a three channel image has \(3\times 3\times 3=27\) weights plus a bias, and it still produces a single feature map, because it sums across the channels. A layer with \(C_{out}\) filters applied to \(C_{in}\) channels has \(C_{out}(k^2C_{in}+1)\) parameters, and that quantity does not depend on the size of the image.

9.5.4 Pooling

A pooling layer reduces the resolution of the feature maps. The most common takes the maximum over each window of size two, which divides the height and the width by two and keeps only the strongest response.

It serves three purposes. It reduces the volume of computation for the layers that follow. It gives a small invariance to translation, because moving the pattern by one pixel usually does not change the maximum of the window. And it enlarges the receptive field: after two poolings, a unit sees a region of the original image four times wider than its own window, which is how a network built from small kernels ends up looking at large structures.

Pooling has no parameters. Some modern architectures replace it with a stride of two in the convolution, which achieves the same reduction with weights that are learned.

In Python:

def maxpool(im, k=2):
    h, w = im.shape
    h2, w2 = h // k, w // k
    return im[:h2*k, :w2*k].reshape(h2, k, w2, k).max(axis=(1, 3))

edge = convolve2d(img, kernels["vertical edges"], mode="same")

fig, axes = plt.subplots(1, 3, figsize=(7, 2.5))
for ax, (m, t) in zip(axes, [(img, "input 8x8"), (edge, "feature map 8x8"),
                             (maxpool(edge), "after 2x2 max pool, 4x4")]):
    ax.imshow(m, cmap="gray_r"); ax.set_title(t, fontsize=9)
    ax.set_xticks([]); ax.set_yticks([])
#> <matplotlib.image.AxesImage object at 0x0000020F2C813FB0>
#> Text(0.5, 1.0, 'input 8x8')
#> []
#> []
#> <matplotlib.image.AxesImage object at 0x0000020EB834FEC0>
#> Text(0.5, 1.0, 'feature map 8x8')
#> []
#> []
#> <matplotlib.image.AxesImage object at 0x0000020F2C813DD0>
#> Text(0.5, 1.0, 'after 2x2 max pool, 4x4')
#> []
#> []
plt.tight_layout(); plt.savefig("dl_pool_py.png"); plt.clf(); plt.close()
max pooling

Figure 9.23: max pooling

9.5.5 A complete architecture

The classical arrangement alternates convolution, activation and pooling several times, then flattens and finishes with one or two dense layers. We build one on the handwritten digits, which are small enough to be trained here in a few seconds.

In Python:

import torch
import torch.nn as nn
from sklearn.model_selection import train_test_split

torch.manual_seed(0)
#> <torch._C.Generator object at 0x0000020E2AABD510>
X_img = digits.images[:, None, :, :].astype("float32") / 16.0   # N,1,8,8
y_all = digits.target.astype("int64")
Xtr_i, Xte_i, ytr_i, yte_i = train_test_split(X_img, y_all, test_size=.25,
                                              random_state=0, stratify=y_all)
Xtr_t = torch.tensor(Xtr_i); ytr_t = torch.tensor(ytr_i)
Xte_t = torch.tensor(Xte_i); yte_t = torch.tensor(yte_i)

class SmallCNN(nn.Module):
    def __init__(self, n_out=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, padding=1),   # 8x8 -> 8x8
            nn.ReLU(),
            nn.MaxPool2d(2),                              # 8x8 -> 4x4
            nn.Conv2d(16, 32, kernel_size=3, padding=1),  # 4x4 -> 4x4
            nn.ReLU(),
            nn.MaxPool2d(2),                              # 4x4 -> 2x2
        )
        self.head = nn.Sequential(
            nn.Flatten(),
            nn.Dropout(0.25),
            nn.Linear(32 * 2 * 2, n_out),
        )
    def forward(self, x):
        return self.head(self.features(x))

cnn = SmallCNN()
opt_c = torch.optim.Adam(cnn.parameters(), lr=2e-3)
lossf = nn.CrossEntropyLoss()

cnn_hist = []
for epoch in range(40):
    cnn.train()
    perm = torch.randperm(len(Xtr_t))
    for i in range(0, len(perm), 64):
        idx = perm[i:i+64]
        opt_c.zero_grad()
        l = lossf(cnn(Xtr_t[idx]), ytr_t[idx])
        l.backward(); opt_c.step()
    cnn.eval()
    with torch.no_grad():
        tr_acc = (cnn(Xtr_t).argmax(1) == ytr_t).float().mean().item()
        te_acc = (cnn(Xte_t).argmax(1) == yte_t).float().mean().item()
    cnn_hist.append({"epoch": epoch + 1, "train": tr_acc, "test": te_acc})
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )
#> SmallCNN(
#>   (features): Sequential(
#>     (0): Conv2d(1, 16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (1): ReLU()
#>     (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>     (3): Conv2d(16, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
#>     (4): ReLU()
#>     (5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
#>   )
#>   (head): Sequential(
#>     (0): Flatten(start_dim=1, end_dim=-1)
#>     (1): Dropout(p=0.25, inplace=False)
#>     (2): Linear(in_features=128, out_features=10, bias=True)
#>   )
#> )

cnn_df = pd.DataFrame(cnn_hist)
n_par = sum(p.numel() for p in cnn.parameters())

cnn_summary = pd.DataFrame({
    "parameters": [n_par],
    "final_train_accuracy": [round(cnn_df["train"].iloc[-1], 4)],
    "final_test_accuracy": [round(cnn_df["test"].iloc[-1], 4)],
})
Table 9.4: a small convolutional network on the handwritten digits
parameters final_train_accuracy final_test_accuracy
6090 0.9985 0.9867 
fig, axes = plt.subplots(1, 2, figsize=(10, 3.2))
axes[0].plot(cnn_df["epoch"], cnn_df["train"], label="training")
#> [<matplotlib.lines.Line2D object at 0x0000020F2C65BF80>]
axes[0].plot(cnn_df["epoch"], cnn_df["test"], label="test")
#> [<matplotlib.lines.Line2D object at 0x0000020F2C96C8F0>]
axes[0].set_xlabel("epoch"); axes[0].set_ylabel("accuracy"); axes[0].legend()
#> Text(0.5, 0, 'epoch')
#> Text(0, 0.5, 'accuracy')
#> <matplotlib.legend.Legend object at 0x0000020F2C813680>
axes[0].set_title("learning curve", fontsize=10)
#> Text(0.5, 1.0, 'learning curve')
W = cnn.features[0].weight.detach().numpy()[:, 0]      # 16 filters of 3x3
grid_img = np.zeros((4 * 4, 4 * 4))
for i in range(16):
    # do not name these r and c: reticulate exposes the R session as `r`
    row_i, col_i = divmod(i, 4)
    k = W[i]; k = (k - k.min()) / (k.ptp() + 1e-9)
    grid_img[row_i*4:row_i*4+3, col_i*4:col_i*4+3] = k
axes[1].imshow(grid_img, cmap="gray_r")
#> <matplotlib.image.AxesImage object at 0x0000020EB3D96C30>
axes[1].set_title("the 16 filters of the first layer", fontsize=10)
#> Text(0.5, 1.0, 'the 16 filters of the first layer')
axes[1].set_xticks([]); axes[1].set_yticks([])
#> []
#> []
plt.tight_layout(); plt.savefig("dl_cnn_py.png"); plt.clf(); plt.close()
learning curve and first layer filters

Figure 9.24: learning curve and first layer filters

The right panel shows the sixteen kernels that the first layer has learned. They were initialized at random and nobody told them what to look for; several of them have become oriented contrast detectors, close in spirit to the Sobel operators we applied by hand at the beginning of the section. This is the point of the whole architecture: the features are not designed, they are learned.

9.5.6 Transfer learning

Training a large network from scratch requires a great deal of data and of computation. Transfer learning avoids both by reusing a network that has already been trained on another, larger problem.

The idea rests on what we have just seen. The first layers of a convolutional network learn very general features, contours and textures, which are useful for almost any image task. Only the last layers are specific to the original classes. We can therefore keep the convolutional part, called the backbone, replace the final classifier by a new one adapted to our classes, and train only that.

Two regimes are used in practice. Feature extraction freezes the backbone entirely and trains only the new head, which is fast and works with very few examples. Fine tuning unfreezes the last blocks of the backbone as well and continues training with a small learning rate, which gives better results when the new task is far from the original one but requires more data.

We can demonstrate the principle here, without downloading anything, by splitting the digits into two disjoint tasks: a network is trained on the digits from zero to four, and its convolutional part is then reused for the digits from five to nine, which it has never seen.

In Python:

torch.manual_seed(0)
#> <torch._C.Generator object at 0x0000020E2AABD510>
# task A: the digits 0 to 4
mA = y_all <= 4
XA = torch.tensor(X_img[mA]); yA = torch.tensor(y_all[mA])
# task B: the digits 5 to 9, relabelled 0 to 4
mB = y_all >= 5
XB_all = X_img[mB]; yB_all = (y_all[mB] - 5)
XB_tr, XB_te, yB_tr, yB_te = train_test_split(XB_all, yB_all, test_size=.5,
                                              random_state=0, stratify=yB_all)
XB_tr_t = torch.tensor(XB_tr); yB_tr_t = torch.tensor(yB_tr)
XB_te_t = torch.tensor(XB_te); yB_te_t = torch.tensor(yB_te)

def fit(model, X, y, epochs=40, lr=2e-3):
    o = torch.optim.Adam([p for p in model.parameters() if p.requires_grad], lr=lr)
    for _ in range(epochs):
        model.train()
        perm = torch.randperm(len(X))
        for i in range(0, len(perm), 64):
            idx = perm[i:i+64]
            o.zero_grad(); lossf(model(X[idx]), y[idx]).backward(); o.step()
    return model

def acc(model, X, y):
    model.eval()
    with torch.no_grad():
        return (model(X).argmax(1) == y).float().mean().item()

# 1. pre-train on task A
source = fit(SmallCNN(n_out=5), XA, yA)

# 2. transfer: reuse the frozen backbone, train a new head only
transferred = SmallCNN(n_out=5)
transferred.features.load_state_dict(source.features.state_dict())
#> <All keys matched successfully>
for p in transferred.features.parameters():
    p.requires_grad = False
transferred = fit(transferred, XB_tr_t, yB_tr_t)

# 3. baseline: the same architecture trained from scratch on task B
scratch = fit(SmallCNN(n_out=5), XB_tr_t, yB_tr_t)

trainable = sum(p.numel() for p in transferred.parameters() if p.requires_grad)

# 4. the same comparison when the target task has only a few examples,
#    which is the situation transfer learning is meant for
rows_tl = []
for n_small in [20, 50, 100, len(XB_tr_t)]:
    torch.manual_seed(0)
    sub = torch.randperm(len(XB_tr_t))[:n_small]
    sc = fit(SmallCNN(n_out=5), XB_tr_t[sub], yB_tr_t[sub])
    tf = SmallCNN(n_out=5)
    tf.features.load_state_dict(source.features.state_dict())
    for prm in tf.features.parameters():
        prm.requires_grad = False
    tf = fit(tf, XB_tr_t[sub], yB_tr_t[sub])
    rows_tl.append({"training examples": n_small,
                    "from scratch": round(acc(sc, XB_te_t, yB_te_t), 4),
                    "transferred": round(acc(tf, XB_te_t, yB_te_t), 4)})
#> <torch._C.Generator object at 0x0000020E2AABD510>
#> <All keys matched successfully>
#> <torch._C.Generator object at 0x0000020E2AABD510>
#> <All keys matched successfully>
#> <torch._C.Generator object at 0x0000020E2AABD510>
#> <All keys matched successfully>
#> <torch._C.Generator object at 0x0000020E2AABD510>
#> <All keys matched successfully>

transfer_tab = pd.DataFrame({
    "model": ["trained from scratch on task B", "frozen backbone from task A + new head"],
    "trainable_parameters": [sum(p.numel() for p in scratch.parameters()), trainable],
    "test_accuracy_on_task_B": [round(acc(scratch, XB_te_t, yB_te_t), 4),
                                round(acc(transferred, XB_te_t, yB_te_t), 4)],
})
transfer_size_tab = pd.DataFrame(rows_tl)
Table 9.5: transfer between two disjoint sets of digits
model trainable_parameters test_accuracy_on_task_B
trained from scratch on task B 5445 0.9844
frozen backbone from task A + new head 645 0.9598
Table 9.6: test accuracy against the number of training examples of the target task
training examples from scratch transferred
20 0.4107 0.6808
50 0.7768 0.8906
100 0.9420 0.9107
448 0.9844 0.9621

The two columns cross. With twenty or fifty examples the frozen backbone wins clearly, because the features it brings were learned on data that the small sample cannot provide. As the sample grows the advantage shrinks and then reverses, since a network trained from scratch can shape its own features for the task. This crossing is the whole practical rule: transfer is worth it exactly when the target data are insufficient to learn good features directly.

The transferred model trains only its final layer, roughly a tenth of the parameters, and still classifies digits that were absent from the pre-training. But it does not beat the network trained from scratch, it stays slightly below it, and that result should not be hidden: here the target task has several hundred examples, which is plenty for a network this small, so freezing the backbone only removes freedom without buying anything.

The benefit of transfer appears when the target data are scarce, which is the situation the method was invented for. Repeating the comparison while shrinking the training set makes this visible.

Transfer works when the two domains resemble each other. A backbone trained on photographs transfers well to other photographs, much less well to medical images or to satellite pictures, whose statistics are different. When the gap is large, fine tuning the deeper layers becomes necessary, and beyond a certain distance training from scratch on the target domain is preferable.