I want to use BackboneFinetuning from PyTorch Lightning with smp

[JayThibs]

Last time I implemented it was with a HuggingFace model as the backbone. From my understanding, you need to instantiate self.backbone with the model weights, and then pass the data through self.backbone and the head separately in forward(). This is what I did with hugginface:

...in init()...
self.backbone = AutoModel.from_pretrained(backbone_name, return_dict=True)
self.classifier = nn.Linear(self.backbone.config.hidden_size, num_classes)

...

def forward(self, input_ids, attention_mask, labels=None):
    output = self.backbone(input_ids, attention_mask=attention_mask)
    output = self.classifier(output.pooler_output)
    output = torch.sigmoid(output)
    ...

So, something similar to the answer to this StackOverflow question. I would instantiate like this (?):

import torchvision.models as models

class ImagenetTransferLearning(LightningModule):
    def __init__(self):
        super().__init__()

        # init a pretrained resnet
        backbone_tmp = models.resnet50(pretrained=True)
        num_filters = backbone_tmp.fc.in_features
        layers = list(backbone_tmp.children())[:-1]
        self.backbone = nn.Sequential(*layers)

        # use the pretrained model to classify cifar-10 (10 image classes)
        num_target_classes = 10
        self.classifier = nn.Linear(num_filters, num_target_classes)

What I'm struggling with is how to do this with the segmentation_models.pytorch package since this is how you instantiate a model with that package:

import segmentation_models.pytorch as smp

model = smp.Unet(
    encoder_name=self.backbone,
    encoder_weights=self.weights,
    in_channels=2,
    classes=2,
)

Is there any way I can instantiate the backbone separately from the Unet/architecture? And have the data flow through the backbone during forward()?

Thanks!