"""Original reduced-width AlexNet reference with five convolutions and three FCs. 227x227 input follows the common 55x55 first-map convention. Widths are reduced to [8,16,24,24,16], with grouped conv2/4/5, LRN, overlapping pooling and dropout. The original 2012 model is much wider and trained for ImageNet classification. """ import torch import torch.nn as nn import torch.nn.functional as F # tensorviz: input=1,3,227,227 class ReducedAlexNet(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(3, 8, 11, stride=4) self.relu1 = nn.ReLU() # PyTorch divides alpha by window size; 5e-4 / 5 = paper's 1e-4. self.norm1 = nn.LocalResponseNorm(5, alpha=5e-4, beta=0.75, k=2) self.pool1 = nn.MaxPool2d(3, stride=2) self.conv2 = nn.Conv2d(8, 16, 5, padding=2, groups=2) self.relu2 = nn.ReLU() self.norm2 = nn.LocalResponseNorm(5, alpha=5e-4, beta=0.75, k=2) self.pool2 = nn.MaxPool2d(3, stride=2) self.conv3 = nn.Conv2d(16, 24, 3, padding=1) self.relu3 = nn.ReLU() self.conv4 = nn.Conv2d(24, 24, 3, padding=1, groups=2) self.relu4 = nn.ReLU() self.conv5 = nn.Conv2d(24, 16, 3, padding=1, groups=2) self.relu5 = nn.ReLU() self.pool5 = nn.MaxPool2d(3, stride=2) self.flatten = nn.Flatten(1) self.fc6 = nn.Linear(16 * 6 * 6, 64) self.relu6 = nn.ReLU() self.dropout6 = nn.Dropout(0.5) self.fc7 = nn.Linear(64, 64) self.relu7 = nn.ReLU() self.dropout7 = nn.Dropout(0.5) self.fc8 = nn.Linear(64, 10) def forward(self, x): x = self.pool1(self.norm1(self.relu1(self.conv1(x)))) x = self.pool2(self.norm2(self.relu2(self.conv2(x)))) x = self.relu3(self.conv3(x)) x = self.relu4(self.conv4(x)) x = self.pool5(self.relu5(self.conv5(x))) x = self.dropout6(self.relu6(self.fc6(self.flatten(x)))) x = self.dropout7(self.relu7(self.fc7(x))) return self.fc8(x) def verify(): torch.manual_seed(19) model = ReducedAlexNet().double().eval() with torch.no_grad(): x = torch.randn(1, 3, 227, 227, dtype=torch.float64) first = model.relu1(model.conv1(x)) assert first.shape == (1, 8, 55, 55) expected = torch.stack( [ first[:, i] / ( 2 + 1e-4 * first[:, max(0, i - 2) : min(8, i + 3)].square().sum(1) ).pow(0.75) for i in range(8) ], dim=1, ) torch.testing.assert_close(model.norm1(first), expected, rtol=1e-12, atol=1e-12) pooled = model.pool1(first) windows = first.unfold(2, 3, 2).unfold(3, 3, 2).amax((-2, -1)) torch.testing.assert_close(pooled, windows) altered = pooled.clone() altered[:, :4] += 3 torch.testing.assert_close( model.conv2(pooled)[:, 8:], model.conv2(altered)[:, 8:] ) assert not torch.allclose( model.conv2(pooled)[:, :8], model.conv2(altered)[:, :8] ) assert model(x).shape == (1, 10) torch.testing.assert_close(model(x), model(x), rtol=0, atol=0) drop = nn.Dropout(0.5) values = torch.ones(1000) torch.manual_seed(2) training = drop(values) assert set(training.tolist()) == {0.0, 2.0} drop.eval() torch.testing.assert_close(drop(values), values) return [ "The 227-pixel convention produces 55-pixel first-convolution maps and ten output logits", "LRN matches the original cross-channel sum formula, accounting for PyTorch's alpha convention", "Overlapping 3×3 stride-2 pooling matches an independent sliding-window maximum", "Grouped convolution isolates its channel groups; dropout scales survivors in training and is identity in evaluation", ] def thumbnail(values): size = min(8, values.shape[-1]) return F.adaptive_avg_pool2d(values[0, :1], (size, size))[0].tolist() def experiment(): torch.manual_seed(19) model = ReducedAlexNet().double().eval() image = torch.zeros(1, 3, 227, 227, dtype=torch.float64) image[:, :, :, 80:145] = 1 cases = [] with torch.no_grad(): c1 = model.relu1(model.conv1(image)) p1 = model.pool1(model.norm1(c1)) c2 = model.relu2(model.conv2(p1)) p2 = model.pool2(model.norm2(c2)) c5 = model.relu5( model.conv5(model.relu4(model.conv4(model.relu3(model.conv3(p2))))) ) p5 = model.pool5(c5) for identity, label, before, after, target in [ ("first", "First convolution → LRN → pool", c1, p1, "pool1"), ("second", "Grouped convolution 2 → LRN → pool", c2, p2, "conv2"), ("fifth", "Convolution 5 → final pool", c5, p5, "pool5"), ]: cases.append( { "id": identity, "label": label, "target": target, "note": "These are thumbnails of at most 8×8 cells, average-pooled from a recorded feature channel of a synthetic stripe image. The first two paths apply cross-channel LRN before pooling; the fifth pools directly. The displayed maps are untrained features and do not demonstrate ImageNet accuracy.", "matrices": [ { "label": "Before normalization / pooling · thumbnail", "values": thumbnail(before), }, { "label": "After pooling · thumbnail", "values": thumbnail(after), }, ], "vectors": [ { "label": "Feature shape · channels / height / width", "values": list(after.shape[1:]), } ], "metrics": [ { "label": "Spatial width before pool", "value": before.shape[-1], }, {"label": "Spatial width after pool", "value": after.shape[-1]}, { "label": "Values after pool · all channels", "value": after.numel(), }, ], } ) return { "kind": "matrices", "title": "Build features while reducing image resolution.", "description": "Five convolutions lead to a 16×6×6 representation and three fully connected layers. ReLU, overlapping pooling, LRN, grouped computation and dropout preserve the main architecture choices in a smaller model.", "controlLabel": "Feature stage", "cases": cases, }