1import torch
2import torch.nn as nn
3
4
5class ResidualBlock(nn.Module):
6 """Basic block from 'Deep Residual Learning' (ResNet-18/34)."""
7
8 def __init__(self, channels: int):
9 super().__init__()
10 # The residual function F(x): conv -> BN -> ReLU -> conv -> BN
11 self.conv1 = nn.Conv2d(channels, channels, kernel_size=3,
12 padding=1, bias=False)
13 self.bn1 = nn.BatchNorm2d(channels)
14 self.conv2 = nn.Conv2d(channels, channels, kernel_size=3,
15 padding=1, bias=False)
16 self.bn2 = nn.BatchNorm2d(channels)
17 self.relu = nn.ReLU(inplace=True)
18
19 def forward(self, x: torch.Tensor) -> torch.Tensor:
20 identity = x # <-- save the input
21
22 out = self.relu(self.bn1(self.conv1(x)))
23 out = self.bn2(self.conv2(out)) # this is F(x)
24
25 out += identity # <-- y = F(x) + x (the whole paper)
26 return self.relu(out) # ReLU after the addition
27
28
29# --- A deep network is now just a stack of these ---
30class TinyResNet(nn.Module):
31 def __init__(self, num_blocks: int = 18, channels: int = 64):
32 super().__init__()
33 self.stem = nn.Conv2d(3, channels, kernel_size=7,
34 stride=2, padding=3, bias=False)
35 self.blocks = nn.Sequential(
36 *[ResidualBlock(channels) for _ in range(num_blocks)]
37 )
38 self.head = nn.Linear(channels, 1000)
39
40 def forward(self, x):
41 x = self.stem(x)
42 x = self.blocks(x) # 18 blocks deep, trains happily
43 x = x.mean(dim=(2, 3)) # global average pooling
44 return self.head(x)
45
46
47net = TinyResNet()
48out = net(torch.randn(1, 3, 224, 224))
49print(out.shape) # torch.Size([1, 1000])