The problem in brief
Problem 8 of the 2026 USA AI Olympiad Round 1 asks you to get your hands on a real, large network: torchvision's ImageNet-pretrained ResNet-50. You must count its learnable parameters, give the shape of an intermediate output, count the convolution parameters of one residual block by hand, and finally reuse the network, cut after block 4 of stage 3, as a frozen feature extractor for a new 5-class classification task. Everything is checked in solution.py. The counts only depend on the architecture, so we build the model without downloading weights; the pretrained model gives the same numbers.
How you are scored
5 + 5 + 5 + 10 = 25 points. Reasoning is required for 8.3, and in 8.3 you may not call numel() on parameters. No official answers are published with the problem PDF.
| part | answer |
|---|---|
| 8.1 | 25 557 032 learnable parameters |
| 8.2 | (B, 512, 28, 28) |
| 8.3 | 1 114 112 convolution parameters in layer3[2] |
| 8.4 | the backbone up to layer3[4], frozen, then global average pooling and Linear(1024, 5) |
Question by question
8.1 Total number of learnable parameters
total = sum(p.numel() for p in model.parameters() if p.requires_grad) # 25557032
Where they are: the stem convolution has 9 408, stage 1 has 0.22 M, stage 2 has 1.22 M, stage 3 has 7.10 M, stage 4 has 14.96 M, and the final fully connected layer has . Two things to notice. Most of the weights sit in the deep stages, where there are many channels, even though they work on tiny or grids. And the BatchNorm (53 120 values) are learnable, while the BatchNorm running means and variances are buffers, not parameters, and are not counted.
8.2 Shape after layer2
Follow the spatial size: the stem convolution has stride 2 (), max-pooling has stride 2 (), layer1 keeps 56, and the first block of layer2 has stride 2 (). For channels, every bottleneck block outputs its inner width: layer1 has inner width 64 and outputs 256, layer2 has inner width 128 and outputs 512. So the output is , which a forward pass confirms.
8.3 Convolution parameters of layer3[2], by hand
layer3[2] is a bottleneck block that is not the first of its stage, so it has no downsampling shortcut: input and output both have 1024 channels and the shortcut is the identity. Its three convolutions have no bias (a BatchNorm follows each of them, and its would make a bias redundant). A convolution with inputs, outputs and a kernel has weights:
Total: . In the code we read the shapes from the module attributes (in_channels, out_channels, kernel_size, bias), not from numel(), and the cross-check with numel() agrees.
This is the point of the "bottleneck" design: the convolutions shrink 1024 channels to 256 before the expensive convolution. A plain convolution from 1024 to 1024 channels alone would cost M parameters, 8.5 times more.
8.4 A frozen truncated backbone for 5 classes
"Up to the output of layer3[4]" means: stem, layer1, layer2, and the first five blocks (indices 0 to 4) of layer3, which has six. That output has 1024 channels on a grid for a 224-pixel input. Then pool and add a new linear head:
class FrozenBackboneClassifier(nn.Module):
def __init__(self, base, num_classes=5):
super().__init__()
self.backbone = nn.Sequential(base.conv1, base.bn1, base.relu, base.maxpool,
base.layer1, base.layer2, base.layer3[:5])
for p in self.backbone.parameters():
p.requires_grad = False # frozen
self.pool = nn.AdaptiveAvgPool2d(1)
self.fc = nn.Linear(1024, num_classes) # the only trainable part
def train(self, mode=True): # keep BatchNorm statistics frozen too
super().train(mode); self.backbone.eval(); return self
def forward(self, x):
return self.fc(torch.flatten(self.pool(self.backbone(x)), 1))
Built on resnet50(weights="IMAGENET1K_V1") (or pretrained=True in older torchvision), it outputs . Only parameters are trainable, and 7 426 112 are frozen. Two details separate a good answer from a merely working one:
requires_grad = Falseis not the whole story. BatchNorm layers update their running mean and variance whenever the module is in training mode, even when their weights are frozen. Overridingtrain()so that the backbone always stays ineval()mode keeps the pretrained statistics intact.nn.AdaptiveAvgPool2d(1)makes the head independent of the input resolution.
Pass only the trainable parameters to the optimiser, e.g. torch.optim.Adam(clf.fc.parameters()).
What to remember
- Count parameters with
sum(p.numel() for p in model.parameters()); BatchNorm running statistics are buffers, not parameters. - A convolution has weights (plus biases if any). ResNet convolutions have no bias because BatchNorm follows them.
- Spatial size halves at each stride-2 stage: .
- Bottleneck blocks (, , ) make deep networks affordable.
- Freezing a backbone means
requires_grad = Falseand keeping its BatchNorm layers in eval mode.