Transfer Learning for Images

Step 1: Select a Pre-Trained Model

The choice of base model determines the quality of transferred features, the computational cost of training and inference, and the practical accuracy ceiling for your task. Models pre-trained on ImageNet-1K (1000 classes, 1.2 million images) have been the default starting point for a decade. Larger pre-training datasets like ImageNet-21K (14 million images, 21,841 classes) and LAION-5B (5 billion image-text pairs) produce features that transfer even better because they encode a broader understanding of the visual world.

For classification tasks, ResNet-50 remains a solid baseline that runs efficiently on any hardware. EfficientNet-B0 through B7 offer a range of accuracy-efficiency trade-offs, with B3 and B4 being the most common choices that balance accuracy and training speed. Vision Transformers (ViT-Base, ViT-Large) pre-trained with self-supervised methods like DINOv2 currently produce the strongest general-purpose features but require more compute for training and inference. ConvNeXt models modernize the CNN architecture with transformer-inspired design choices, achieving ViT-level accuracy while maintaining the deployment efficiency of convolutional networks.

The pre-training method matters as much as the architecture. Supervised pre-training on ImageNet produces features optimized for the 1000 ImageNet categories. Self-supervised pre-training (DINO, MAE, MoCo) produces features that are more general because they are not biased toward any specific set of categories. CLIP pre-training on image-text pairs produces features that encode semantic concepts aligned with natural language, enabling zero-shot classification without any fine-tuning. For domains very different from natural photographs (medical images, satellite imagery, microscopy), models pre-trained on domain-specific datasets often outperform ImageNet-pre-trained models because the visual features are more relevant.

Practical considerations include model size (which determines GPU memory requirements during training and storage requirements for deployment), inference speed (critical for real-time applications), and availability of pre-trained weights. PyTorch's torchvision.models and the timm (PyTorch Image Models) library together provide pre-trained weights for over 1,000 model configurations, all accessible with a single function call. Hugging Face hosts thousands more, including domain-specific models for medical imaging, remote sensing, and scientific applications.

Step 2: Replace the Classification Head

A pre-trained model's final layer (the classification head) maps the learned features to the original training categories. A ResNet-50 pre-trained on ImageNet has a final linear layer with 1000 output units, one for each ImageNet class. This layer is specific to ImageNet and must be replaced with a new one matching your number of target classes. If you are classifying 5 types of manufacturing defects, you replace the 1000-unit layer with a 5-unit layer. If you are doing binary classification (defective vs acceptable), you use a 2-unit layer or a single output with sigmoid activation.

In PyTorch, replacing the head for a ResNet is a single line: model.fc = nn.Linear(2048, num_classes). For EfficientNet: model.classifier[-1] = nn.Linear(1280, num_classes). For ViT: model.head = nn.Linear(768, num_classes). The exact attribute name varies by architecture, so checking the model definition or printing model.named_modules() reveals where the classification head is. The new layer is initialized with random weights, while all other layers retain their pre-trained weights. This means the model immediately produces meaningful feature representations for your images, even before any training on your dataset, because only the final mapping from features to predictions needs to be learned.

For tasks beyond simple classification, more substantial head modifications are needed. Object detection requires replacing the classification head with a detection head that outputs bounding box coordinates and class scores. Segmentation requires replacing it with a decoder that upsamples feature maps back to the input resolution. Libraries like torchvision provide pre-built model constructors (e.g., torchvision.models.detection.fasterrcnn_resnet50_fpn) that attach task-specific heads to pre-trained backbones, handling the architectural changes automatically.

Step 3: Choose Feature Extraction or Fine-Tuning

The two main transfer learning strategies differ in which parts of the model are updated during training. Feature extraction freezes all pre-trained layers and trains only the new classification head. The pre-trained model serves as a fixed feature extractor, transforming each image into a feature vector that the new head learns to classify. Fine-tuning updates all (or some) of the pre-trained layers along with the new head, allowing the model to adapt its feature representations to your specific domain.

Feature extraction is faster, requires less data, and is less prone to overfitting. Because only the head's weights are updated (typically a few thousand parameters), training completes in minutes rather than hours and can succeed with as few as 50 to 100 images per class. The risk is that the pre-trained features may not capture the visual distinctions that matter for your task, particularly if your domain is very different from the pre-training data. Medical microscopy images, for instance, contain textures and structures that have no equivalent in ImageNet, so frozen ImageNet features may not discriminate well between your target classes.

Fine-tuning adapts the features to your domain, typically achieving higher accuracy than feature extraction when sufficient training data is available (roughly 1,000+ images per class). The risk is overfitting: with millions of trainable parameters and a small dataset, the model can memorize the training images rather than learning generalizable patterns. Regularization techniques including dropout, weight decay, data augmentation, and early stopping are essential to prevent this. Fine-tuning requires a lower learning rate than training from scratch (typically 10x to 100x lower) to avoid destroying the useful patterns encoded in the pre-trained weights.

A middle ground is gradual unfreezing, where you initially freeze all pre-trained layers and train only the head for a few epochs, then progressively unfreeze deeper layers and continue training. This approach lets the head learn a reasonable mapping from features to predictions before the features themselves start changing, providing a more stable training trajectory. The ULMFiT paper (2018, originally for NLP) popularized this approach, and it has proven equally effective for vision tasks. A typical schedule unfreezes the last block after 5 epochs, the second-to-last block after 10 epochs, and so on.

Step 4: Set Learning Rates and Train

Learning rate is the most important hyperparameter for transfer learning. A learning rate that is too high destroys the pre-trained features in the first few gradient updates, erasing the benefit of transfer learning entirely. A learning rate that is too low makes training painfully slow and may trap the model in a suboptimal solution. The standard starting point for fine-tuning is 1e-4 to 1e-3 for the new head and 1e-5 to 1e-4 for the pre-trained layers. If using feature extraction (frozen backbone), the head can use a higher learning rate of 1e-3 to 1e-2 since it is the only thing being trained.

Discriminative learning rates (also called layer-wise learning rate decay) assign different learning rates to different layer groups, with deeper (earlier) layers receiving lower rates than shallower (later) layers. The rationale is that early layers learn universal features (edges, textures) that should change minimally, while later layers learn more task-specific features that benefit from larger updates. A typical setup multiplies the learning rate by a decay factor (0.1 to 0.5) for each layer group moving from the classification head toward the input. A model with 4 layer groups might use learning rates of 1e-3, 1e-4, 1e-5, and 1e-6 from head to stem.

Learning rate schedulers reduce the learning rate during training, allowing the model to make large initial updates and then fine-tune with smaller steps. Cosine annealing, which smoothly decreases the learning rate following a cosine curve from the initial value to near zero over the training run, is the most popular schedule for transfer learning. Warmup, which starts with a very low learning rate and linearly increases it over the first 5 to 10% of training steps, prevents the initial large gradients from destabilizing the pre-trained weights.

Training duration for transfer learning is much shorter than training from scratch. Feature extraction typically converges in 10 to 30 epochs. Fine-tuning typically requires 20 to 100 epochs depending on dataset size and domain similarity. Early stopping, which halts training when validation accuracy stops improving (with a patience of 5 to 10 epochs), prevents overfitting and eliminates the need to guess the optimal number of epochs. Always monitor both training and validation metrics throughout training to diagnose issues: if training accuracy is high but validation accuracy is low, the model is overfitting and needs stronger regularization or less aggressive fine-tuning.

Step 5: Evaluate and Iterate

Evaluate the fine-tuned model on a held-out test set that was not used during training or for hyperparameter selection. Report accuracy, precision, recall, and F1 score for each class. A confusion matrix reveals which classes the model confuses, guiding further data collection or augmentation. Per-class accuracy matters more than overall accuracy when classes are imbalanced: a model might achieve 95% overall accuracy while failing completely on a rare class that represents only 2% of the test set.

If accuracy is unsatisfactory, diagnose whether the issue is underfitting (model has not learned enough) or overfitting (model has memorized training data). Underfitting manifests as low training accuracy and can be addressed by unfreezing more layers, increasing model capacity, training longer, or using a higher learning rate. Overfitting manifests as high training accuracy with much lower validation accuracy and can be addressed by stronger data augmentation, more aggressive dropout, weight decay, or using feature extraction instead of full fine-tuning.

For difficult tasks with limited data, combining transfer learning with data augmentation produces the best results. Start with a strong pre-trained model (DINOv2 or CLIP), apply aggressive augmentation (horizontal flip, random crop, color jitter, CutMix), use gradual unfreezing with discriminative learning rates, and train with cosine annealing and warmup. This recipe achieves strong results across an enormous range of vision tasks, from classifying satellite images of crop types to identifying species in wildlife camera traps to grading the severity of retinal disease in fundus photographs.