1import torch
2import torch.nn as nn
3
4
5class AlexNet(nn.Module):
6 """Krizhevsky, Sutskever & Hinton (NeurIPS 2012)."""
7
8 def __init__(self, num_classes: int = 1000):
9 super().__init__()
10 self.features = nn.Sequential(
11 # Conv1: 96 big filters hunting edges and color blobs
12 nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=2),
13 nn.ReLU(inplace=True), # the 6x speed unlock
14 nn.MaxPool2d(kernel_size=3, stride=2), # overlapping pooling
15 # Conv2
16 nn.Conv2d(96, 256, kernel_size=5, padding=2),
17 nn.ReLU(inplace=True),
18 nn.MaxPool2d(kernel_size=3, stride=2),
19 # Conv3-5: composing edges -> textures -> object parts
20 nn.Conv2d(256, 384, kernel_size=3, padding=1),
21 nn.ReLU(inplace=True),
22 nn.Conv2d(384, 384, kernel_size=3, padding=1),
23 nn.ReLU(inplace=True),
24 nn.Conv2d(384, 256, kernel_size=3, padding=1),
25 nn.ReLU(inplace=True),
26 nn.MaxPool2d(kernel_size=3, stride=2),
27 )
28 self.classifier = nn.Sequential(
29 nn.Dropout(p=0.5), # tame the 60M params
30 nn.Linear(256 * 6 * 6, 4096),
31 nn.ReLU(inplace=True),
32 nn.Dropout(p=0.5),
33 nn.Linear(4096, 4096),
34 nn.ReLU(inplace=True),
35 nn.Linear(4096, num_classes), # 1000-way verdict
36 )
37
38 def forward(self, x: torch.Tensor) -> torch.Tensor:
39 x = self.features(x)
40 x = torch.flatten(x, 1)
41 return self.classifier(x)
42
43
44net = AlexNet()
45out = net(torch.randn(1, 3, 224, 224))
46print(out.shape) # torch.Size([1, 1000])
47print(sum(p.numel() for p in net.parameters()) / 1e6, "M params") # ~61M