import torch
from torch import nn
from d2l import torch as d2l
class AlexNet(d2l.Classifier):
def __init__(self, lr=0.1, num_classes=10):
super().__init__()
self.save_hyperparameters()
self.net = nn.Sequential(
nn.LazyConv2d(96, kernel_size=11, stride=4, padding=1),
nn.ReLU(), nn.MaxPool2d(kernel_size=3, stride=2),
nn.LazyConv2d(256, kernel_size=5, padding=2), nn.ReLU(),
nn.MaxPool2d(kernel_size=3, stride=2),
nn.LazyConv2d(384, kernel_size=3, padding=1), nn.ReLU(),
nn.LazyConv2d(384, kernel_size=3, padding=1), nn.ReLU(),
nn.LazyConv2d(256, kernel_size=3, padding=1), nn.ReLU(),
nn.MaxPool2d(kernel_size=3, stride=2), nn.Flatten(),
nn.LazyLinear(4096), nn.ReLU(), nn.Dropout(p=0.5),
nn.LazyLinear(4096), nn.ReLU(),nn.Dropout(p=0.5),
nn.LazyLinear(num_classes))
self.net.apply(d2l.init_cnn)
Construct a single-channel data example to observe the output shape of each layer
AlexNet().layer_summary((1, 1, 224, 224))
Conv2d output shape: torch.Size([1, 96, 54, 54]) ReLU output shape: torch.Size([1, 96, 54, 54]) MaxPool2d output shape: torch.Size([1, 96, 26, 26]) Conv2d output shape: torch.Size([1, 256, 26, 26]) ReLU output shape: torch.Size([1, 256, 26, 26]) MaxPool2d output shape: torch.Size([1, 256, 12, 12]) Conv2d output shape: torch.Size([1, 384, 12, 12]) ReLU output shape: torch.Size([1, 384, 12, 12]) Conv2d output shape: torch.Size([1, 384, 12, 12]) ReLU output shape: torch.Size([1, 384, 12, 12]) Conv2d output shape: torch.Size([1, 256, 12, 12]) ReLU output shape: torch.Size([1, 256, 12, 12]) MaxPool2d output shape: torch.Size([1, 256, 5, 5]) Flatten output shape: torch.Size([1, 6400]) Linear output shape: torch.Size([1, 4096]) ReLU output shape: torch.Size([1, 4096]) Dropout output shape: torch.Size([1, 4096]) Linear output shape: torch.Size([1, 4096]) ReLU output shape: torch.Size([1, 4096]) Dropout output shape: torch.Size([1, 4096]) Linear output shape: torch.Size([1, 10])
Fashion-MNIST images have lower resolution than ImageNet images. We upsample them to $224 \times 224$ start training AlexNet
model = AlexNet(lr=0.01)
data = d2l.FashionMNIST(batch_size=128, resize=(224, 224))
trainer = d2l.Trainer(max_epochs=10, num_gpus=1)
trainer.fit(model, data)