Pre-trained Models in torchvision: ResNet, EfficientNet, and ViT

Training a deep neural network on ImageNet from scratch requires millions of labelled images and weeks of GPU compute. Pre-trained models have already learned general visual features — edges, textures, shapes, and high-level object parts — from this enormous dataset.

By reusing these weights, you benefit from the learning done on 1.2 million images without paying the training cost. This is the core idea of transfer learning: features learned on one large task transfer well to related smaller tasks. torchvision.models provides dozens of pre-trained architectures ready to download and use.

ResNet-50 (Residual Network with 50 layers) introduced skip connections that add the input of a block directly to its output: output = F(x) + x. This allows gradients to flow directly through the addition, enabling training of very deep networks without vanishing gradients.

ResNet-50 has approximately 25 million parameters and consists of: one initial 7×7 convolutional layer, max pooling, four residual blocks (layer1–layer4), and a global average pooling layer followed by a 1000-class fully connected head for ImageNet classification. The final fc layer is what we replace for custom tasks.

Before running inference, inputs must match the preprocessing used during ImageNet training: resize to at least 224×224, normalise with ImageNet mean and standard deviation. The torchvision.transforms API handles this. Always call model.eval() before inference to disable dropout and batch normalisation training mode.

The model outputs 1000 logits, one per ImageNet class. We apply softmax to get probabilities and pick the top-k classes. torchvision.models now includes category names in the weights metadata, eliminating the need for a separate labels file.

EfficientNet (2019) introduced compound scaling: systematically scaling network width, depth, and input resolution together using a single compound coefficient. Instead of arbitrarily making networks wider or deeper, EfficientNet balances all three dimensions for optimal accuracy-efficiency trade-offs.

The EfficientNet family ranges from efficientnet_b0 (5M parameters) to efficientnet_b7 (66M parameters). EfficientNet-B0 achieves comparable accuracy to ResNet-50 while using 8× fewer parameters and 6× fewer FLOPs, making it ideal for mobile and edge deployment. PyTorch provides all eight variants.

Vision Transformers (ViT) (2020) apply the transformer architecture directly to images, without any convolution. The image is split into a grid of fixed-size patches (e.g., 16×16 pixels), each flattened and projected to an embedding vector. These patch embeddings are treated like word tokens in NLP.

A CLS (classification) token is prepended to the patch sequence. After passing through multiple transformer encoder blocks with self-attention, the CLS token's output is used for classification. ViT requires large training datasets to outperform CNNs, but pre-trained ViT models from torchvision bring this power to your tasks immediately.

Choosing among these architectures depends on your constraints and task. ResNet-50 is the reliable default: well understood, strong baseline, many tutorials and implementations. EfficientNet-B0/B2 wins when inference speed and model size matter — mobile apps, real-time systems, or edge hardware.

ViT excels on large-scale tasks and benefits from self-supervised pre-training (DINO, CLIP). It requires more compute and memory than CNNs of similar accuracy. For most custom image classification tasks with moderate-sized datasets, start with EfficientNet-B2 or ResNet-50, then try ViT if you have the resources to fine-tune it.

Before modifying a pre-trained model for your task, inspect its architecture to understand which layers to replace. Use print(model) to see the layer tree, and named_modules() or named_children() to iterate programmatically.

The key insight: every torchvision model ends with a classification head sized for 1000 ImageNet classes. To adapt the model to your task with num_classes different from 1000, you replace this final layer. The feature extractor (everything before the head) retains ImageNet-learned features.

When you replace the classification head, you need to know the feature dimension output by the backbone (everything except the final layer). This dimension is the input size of your new classification head.

Common backbone output dimensions: ResNet-50 outputs 2048, EfficientNet-B0 outputs 1280, and ViT-B/16 outputs 768. These feature vectors are computed by global average pooling over the spatial feature maps, producing a single vector per image. Your replacement head takes this vector as input.

You can quickly verify that a pre-trained model's weights are correct by running it on a well-known test image and checking whether the top prediction matches the expected label. This sanity check ensures the weights loaded correctly and the preprocessing pipeline is correct.

Beyond this test, it is good practice to always check the expected input format from the weights metadata: input size (224×224 for most models), channel order (RGB, not BGR), and normalisation constants (ImageNet mean and std). Using wrong normalisation is a common bug that causes poor transfer learning results.

The typical workflow for using pre-trained models on custom tasks is: (1) load the pre-trained model with weights=...IMAGENET1K..., (2) replace the final classification head with a new nn.Linear sized for your number of classes, (3) optionally freeze the backbone weights initially, and (4) train using a lower learning rate than you would use from scratch.

The new classification head starts with random weights and needs to learn from your data. The backbone starts with excellent features and needs only minor adjustments. This is why differential learning rates — a very low rate for the backbone and a higher rate for the head — often improve convergence speed and final accuracy.

Test your understanding of pre-trained torchvision models from this lesson.

In this lesson you learned: ResNet-50 uses skip connections to train very deep networks and outputs 2048-dimensional features, EfficientNet achieves better accuracy-efficiency trade-offs through compound scaling of width, depth, and resolution, and ViT applies transformer self-attention to image patches without convolutions. Next up we learn feature extraction — freezing the pre-trained backbone and training only a new classification head on your custom dataset.