11_23_ft.py

import pytorch_lightning as pl

class TransferLearningTemplate(pl.LightningModule):
  def __init__(self):
    # basically a feature extractor
    self.pretrained_model = SomePretrainedModel(load_weights=True)
    
    # a model that uses features to do something you care about
    self.finetune_model = nn.Linear(dim, num_classes)
    
  def forward(self, x):
    features = self.pretrained_model(x)
    out = self.finetune_model(features)

  def training_step(self, batch, batch_num):
    x, y = batch
    
    features = self.forward(x)
    return some_loss(features, y)