How to run transfer learning in PyTorch

Transfer learning lets a PyTorch image classifier reuse features learned from a large source dataset while training a smaller head for project-specific classes. It is useful when the target dataset is too small to justify training every layer of a convolutional network from scratch.

In torchvision, pretrained model builders accept a weights enum such as ResNet18_Weights.DEFAULT. The same weights object provides the preprocessing transform, so the training batches use the resize, crop, and normalization expected by the source model.

A fixed-feature-extractor pass freezes the pretrained backbone, replaces the final fully connected layer, and sends only the new head parameters to the optimizer. A successful smoke run should show the new class count, a batch of logits shaped for those classes, no change in an early backbone weight, and a changed classifier head after one optimizer step.

1. Create a transfer-learning smoke script.

   transfer_learning_run.py

   import torch
   from torch import nn
   from torch.utils.data import DataLoader
   from torchvision.datasets import FakeData
   from torchvision.models import ResNet18_Weights, resnet18
    
    
   torch.manual_seed(29)
    
   class_names = ["invoice", "receipt", "statement"]
   weights = ResNet18_Weights.DEFAULT
   preprocess = weights.transforms()
    
   train_data = FakeData(
       size=8,
       image_size=(3, 224, 224),
       num_classes=len(class_names),
       transform=preprocess,
       random_offset=11,
   )
   train_loader = DataLoader(train_data, batch_size=4, shuffle=False)
    
   model = resnet18(weights=weights)
    
   for parameter in model.parameters():
       parameter.requires_grad = False
    
   in_features = model.fc.in_features
   model.fc = nn.Linear(in_features, len(class_names))
    
   trainable_parameters = [
       name for name, parameter in model.named_parameters() if parameter.requires_grad
   ]
    
   optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.01, momentum=0.9)
   loss_fn = nn.CrossEntropyLoss()
    
   images, targets = next(iter(train_loader))
   frozen_probe = model.conv1.weight.detach().clone()
   head_before = model.fc.weight.detach().clone()
    
    
   def set_frozen_backbone_eval(module):
       for child_name, child_module in module.named_children():
           if child_name != "fc":
               child_module.eval()
    
    
   model.train()
   set_frozen_backbone_eval(model)
   optimizer.zero_grad(set_to_none=True)
   logits = model(images)
   loss = loss_fn(logits, targets)
   loss.backward()
   optimizer.step()
    
   predicted_indexes = logits.argmax(dim=1).tolist()
   predicted_labels = [class_names[index] for index in predicted_indexes]
    
   print(f"torch_version={torch.__version__}")
   print(f"torchvision_weights={weights}")
   print(f"class_count={len(class_names)}")
   print(f"batch_shape={tuple(images.shape)}")
   print(f"target_shape={tuple(targets.shape)}")
   print(f"trainable_parameters={','.join(trainable_parameters)}")
   print(f"backbone_eval={not any(child.training for name, child in model.named_children() if name != 'fc')}")
   print(f"head_training={model.fc.training}")
   print(f"logits_shape={tuple(logits.shape)}")
   print(f"loss={loss.item():.4f}")
   print(f"conv1_grad={model.conv1.weight.grad}")
   print(f"conv1_changed={not torch.equal(frozen_probe, model.conv1.weight.detach())}")
   print(f"fc_changed={not torch.equal(head_before, model.fc.weight.detach())}")
   print(f"predicted_labels={','.join(predicted_labels)}")

   FakeData supplies a tiny image-shaped dataset for the smoke run. Replace it with ImageFolder or a project Dataset after the model wiring is proven.
   Related: How to create a custom Dataset in PyTorch

2. Run the transfer-learning smoke script.

   $ python transfer_learning_run.py
   torch_version=2.12.1+cpu
   torchvision_weights=ResNet18_Weights.IMAGENET1K_V1
   class_count=3
   batch_shape=(4, 3, 224, 224)
   target_shape=(4,)
   trainable_parameters=fc.weight,fc.bias
   backbone_eval=True
   head_training=True
   logits_shape=(4, 3)
   loss=1.1993
   conv1_grad=None
   conv1_changed=False
   fc_changed=True
   predicted_labels=statement,statement,statement,statement

   On a first run, torchvision may print a pretrained-weight download line before the training output. The weight file is cached by PyTorch for later runs.

3. Check the freeze and training signals in the output.

   trainable_parameters=fc.weight,fc.bias confirms that only the replacement classifier head reached the optimizer. conv1_grad=None and conv1_changed=False confirm that an early backbone layer stayed frozen, while fc_changed=True confirms that the head changed after the optimizer step.
   Related: How to freeze model layers in PyTorch

4. Build the real image dataset with the pretrained weight transform.

   from torchvision.datasets import ImageFolder
   from torchvision.models import ResNet18_Weights
    
   weights = ResNet18_Weights.DEFAULT
   preprocess = weights.transforms()
    
   train_data = ImageFolder("data/document-images/train", transform=preprocess)

   The transform tied to weights keeps project images aligned with the input size and normalization used by the pretrained ResNet18 weights.

5. Replace the classifier head with the project class count.

   from torch import nn
   from torchvision.models import ResNet18_Weights, resnet18
    
   class_names = train_data.classes
   model = resnet18(weights=ResNet18_Weights.DEFAULT)
    
   for parameter in model.parameters():
       parameter.requires_grad = False
    
   model.fc = nn.Linear(model.fc.in_features, len(class_names))

   Parameters in the new Linear layer keep requires_grad=True by default. Rebuild the head whenever the project class count changes.

6. Train only the replacement head in the project loop.

   optimizer = torch.optim.SGD(model.fc.parameters(), lr=0.01, momentum=0.9)
   loss_fn = nn.CrossEntropyLoss()
    
   model.train()
   for name, child in model.named_children():
       if name != "fc":
           child.eval()
    
   for images, targets in train_loader:
       optimizer.zero_grad(set_to_none=True)
       logits = model(images)
       loss = loss_fn(logits, targets)
       loss.backward()
       optimizer.step()

   Keeping frozen backbone children in evaluation mode prevents BatchNorm running statistics from changing during a fixed-feature-extractor pass.

7. Save the adapted model state dictionary after training.

   torch.save(
       {
           "model_state_dict": model.state_dict(),
           "class_names": class_names,
           "weights": ResNet18_Weights.DEFAULT.name,
       },
       "document-resnet18-transfer.pth",
   )

   Store the class order and weight enum with the state dictionary so inference code can rebuild the same head shape and preprocessing path.
   Related: How to save and load a PyTorch model
   Related: How to run inference in PyTorch

8. Remove the smoke script after the project pattern is confirmed.

   $ rm transfer_learning_run.py