A create_model() function that works with EfficientNet B0 and B2 from the Experiment Tracking section

[Rick-29]

I created a function that should work with any type of pretrained model from torchvision, I only tested it with EfficineNet B0 and B2 currently.

def create_model(architecture: torch.nn.Module,
                 weights: torchvision.models.WeightsEnum,
                 out_shape: int,
                 device: torch.device = device,
                 model: torch.nn.Module = None):
    """
    Returns a model with its weights loaded and set to the given device with the given output shape.
    
    Args:
        architecture (torch.nn.Module): The architecture of the model.
        weights (models.WeightsEnum): The pre-trained weights for the model.
        out_shape (int): The desired output shape of the model.
        device (torch.device, optional): The device where the model will be loaded. Defaults to the current device.
        model (torch.nn.Module, optional): If provided, the function will modify this model instead of creating a new one. Defaults to None.
        
    Returns:
        torch.nn.Module: The model with the modified output layer.
    """
    # If a model is provided, use it; otherwise, create a new one
    model = model or architecture(weights=weights).to(device)
    
    # Freeze the features parameters
    for param in model.features.parameters():
        param.requires_grad = False

    # Adjust the classifier head
    set_seeds()
    model.classifier[-1] = nn.Linear(model.classifier[-1].in_features, out_shape).to(device)
    
    return model

this is an example of it's use:

effnetb2_model = create_model(torchvision.models.efficientnet_b2,
                              effnetb2_weights,
                              len(class_names))

[mrdbourke]

Wow! Fantastic use case, thank you for sharing.

I really like the line model.classifier[-1] = nn.Linear(model.classifier[-1].in_features, out_shape).to(device) as well.

Good way to update the last linear layer.

One slight annoying issue is that some models outside of torchvision call the last layer different names, e.g. model.feature_extractor instead of model.classifier.

I also like the weights: torchvision.models.WeightsEnum - I didn't think of that!