A custom computer vision primitive (get_pretrained_model() ) that facilities easy and smooth transfer learning (feature extraction AND fine-tuning) with ALL torchvision.models

[ChampPhil]

Before I explain, I want to give credit to @Rick-29 - who created a transfer learning function
that can returns a pre-trained model from torchvision.models with its output layer adjusted. The link to that post is here: #675

The limitations of this function are:

1. It does not support fine-tuning, and returns a model with all layers frozen except the output layer.

2. The pre-trained models within torchvision.models can have different naming conventions/organization. For example, torchvision.models.resnet18 modules' are: ['conv1', 'bn1', 'relu', 'maxpool', 'layer1', 'layer2', 'layer3', 'layer4', 'avgpool', 'fc'] and torchvision.models.efficientnet_b0 modules' are ['features', 'avgpool', 'classifier']. This drastic difference would have caused errors in @Rick-29's function and made it only work on certain pre-trained models and error out on others.

I propose a new, optimized function that

1. Supports fine tuning from X to last layer
2. Can work with EVERY pre-trained archictecture in torchvision.models
3. Provides explainable error messages using the pre-built assert functionality

def get_pretrained_model(architecture: str,
                 weights: str,
                 out_shape: int,
                 device: torch.device,
                 perform_fine_tuning: bool = False,
                 fine_tune_starting_at: int = None,
                 ):
    """
    Returns a torchvision model with its weights loaded and set to the given device with the given output shape.
    
    Args:
        architecture (str): The name of the pre-trained model
        weights (str): Name of the pre-trained weights for the model (e.g., "DEFAULT", "ImageNetV1_1K")
        out_shape (int): The desired output shape/number of classes of the model.
        device (torch.device): The device where the model will be loaded.
        perform_fine_tuning (bool): Whether or not to perform fine-tuning (where some layers of the feature extractor are unfrozen).
        fine_tuning_starting (int): Starting from FINE_TUNING_STARTING_AT block of convolutional layers (a conv block is an nn.Sequential() in model.features) the layers will be trainable. For example: If MobileNetV3 has 5 conv blocks in its feature extractor, fine_tune_starting_at=-2 means that the LAST two conv blocks are unfrozen/trainable,and the rest are not.
       
        
    Returns:
        torch.nn.Module: The model with the modified output layer.
    """
    # If a model is provided, use it; otherwise, create a new one
    classification_models = torchvision.models.list_models(module=torchvision.models)

    assert architecture in classification_models, "The specified pre-trained model is not in torchvision.models"


    weight_enum = torchvision.models.get_model_weights(architecture)
  
    try:
      weights = getattr(weight_enum, weights)
    except:
      assert False, "The specified weight name does not exist"

    model = torchvision.models.get_model(name=architecture, weights=weights)  
    module_names = list(model._modules.keys())
 

    #If the model, like resnet18, has is feature extraction head broken done into parts, get each part
    feature_extractor_blocks = []
    if len(module_names) == 3:
      feature_extractor = getattr(model, module_names[0])
    else:
      #If the feature extractor is broken up into multiple layer blocks
      for feature_extractor_block_name in module_names[:-2]:
         feature_extractor_blocks.append(getattr(model, feature_extractor_block_name))

    classifier_block = getattr(model, module_names[-1])

    
    if perform_fine_tuning:
      assert fine_tune_starting_at != None, "If perform_fine_tuning is true, than you must specify which feature extractor layers you want to fine tune with the argument FINE_TUNE_STARTING_AT"
      if fine_tune_starting_at >= 0:
        assert False, "The variable FINE_TUNING_STARTING_AT must be a negative integer"
      #If there are 4 blocks, fine_tuning_at=3 expects the last blocks to be frozen
      fine_tune_starting_at = -abs(fine_tune_starting_at)

      if len(feature_extractor_blocks) == 0:
        for param in feature_extractor[:fine_tune_starting_at].parameters():
          param.requires_grad = False

      else:
        for feature_extractor_block in feature_extractor_blocks[:fine_tune_starting_at]:
            for param in feature_extractor_block.parameters():
              param.requires_grad = False
    else:
      if len(feature_extractor_blocks) == 0:
        for param in feature_extractor.parameters():
          param.requires_grad = False
      else:
        for feature_extractor_block in feature_extractor_blocks:
            for param in feature_extractor_block.parameters():
              param.requires_grad = False
       

    # Adjust the classifier head
    torch.cuda.manual_seed(42)
    torch.manual_seed(42)

    if "__getitem__" in list(dir(classifier_block)): #That means classifier is a nn.Sequential()
      classifier_block[-1] = torch.nn.Linear(classifier_block[-1].in_features, out_shape).to(device)
    else:
      
      setattr(model, module_names[-1], torch.nn.Linear(classifier_block.in_features, out_shape).to(device))