CNN：卷积、池化与图像特征图

floor((input + 2 * padding - kernel_size) / stride) + 1

N: batch size
C: channels
H: height
W: width

import torch
from torch import nn

torch.manual_seed(42)
images = torch.randn(8, 1, 28, 28)


class TinyCNN(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 8, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
            nn.Conv2d(8, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1)),
        )
        self.classifier = nn.Linear(16, 3)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        features = self.features(inputs)
        flattened = torch.flatten(features, start_dim=1)
        return self.classifier(flattened)


model = TinyCNN()
x = images
for layer in model.features:
    x = layer(x)
    print(type(layer).__name__, tuple(x.shape))
logits = model(images)
print("logits:", logits.shape)
print("parameters:", sum(p.numel() for p in model.parameters()))

edge = torch.zeros(1, 1, 7, 7)
edge[:, :, :, 3:] = 1.0
detector = nn.Conv2d(1, 1, kernel_size=3, bias=False)
with torch.no_grad():
    detector.weight.copy_(torch.tensor([[[[-1., 0., 1.], [-1., 0., 1.], [-1., 0., 1.]]]]))
print(detector(edge)[0, 0].round())