Issue with DeepLabv3 inbuilt model and custom backbones

When creating a deeplabv3 model using torchvision inbuilt model i’ve noticed it appears to not work the same when instantiating the model with a prebuilt setup like deeplabv3_mobilenet_v3_large().

Specifically when building a deeplabv3 model like so:

def createdeeplabv3_mobilenet(outputchannels=1):
    """DeepLabv3 class with custom head
    Args:
        outputchannels (int, optional): The number of output channels
        in your dataset masks. Defaults to 1.
    Returns:
        model: Returns the DeepLabv3 model with the MobileNetV3 backbone.
    """
    model = torchvision.models.segmentation.deeplabv3_mobilenet_v3_large(pretrained=True, progress=True)
    #model = torchvision.models.segmentation.deeplabv3_resnet50(pretrained=True, progress=True)
    model.classifier = torchvision.models.segmentation.deeplabv3.DeepLabHead(960, outputchannels)
    # Set the model in training mode
    model.train()
    return model

It works perfectly fine in training and converges.

However if I try to put a different backbone on the network (even even the exact same one), like so:

def createdeeplabv3_effnetv2(outputchannels=1):
    """DeepLabv3 class with custom head
    Args:
        outputchannels (int, optional): The number of output channels
        in your dataset masks. Defaults to 1.
    Returns:
        model: Returns the DeepLabv3 model with the Efficienetv2s backbone.
    """
    #backbone = torchvision.models.efficientnet_v2_s(weights="DEFAULT").features
    #backbone.out_channels = 1280
    #classifier = torchvision.models.segmentation.deeplabv3.DeepLabHead(1280, outputchannels)
    #model = torchvision.models.segmentation.DeepLabV3(backbone=backbone, classifier=classifier)
    model = torchvision.models.segmentation.deeplabv3_mobilenet_v3_large(pretrained=True, progress=True)
    model.backbone = torchvision.models.mobilenet_v3_large(weights="DEFAULT").features
    model.classifier = torchvision.models.segmentation.deeplabv3.DeepLabHead(1280, outputchannels)
    # Set the model in training mode
    print(model)
    model.train()
    return model

This error is returned when I try to train:

File “/home/justin/Documents/Code/computer-vision-utils/deeplearning/DeepLabV3Plus/main.py”, line 90, in main
model_trained, history = trainer.train_segmentation(model, dataloaders, criterion, optimizer, DEVICE, args.Checkpoint, scheduler=scheduler, num_epochs=EPOCHS, project=project)
File “/home/justin/Documents/Code/computer-vision-utils/deeplearning/DeepLabV3Plus/trainer.py”, line 70, in train_segmentation
preds = model(inputs.to(device))
File “/home/justin/miniconda3/envs/pytorch_env/lib/python3.9/site-packages/torch/nn/modules/module.py”, line 1130, in _call_impl
return forward_call(*input, **kwargs)
File “/home/justin/miniconda3/envs/pytorch_env/lib/python3.9/site-packages/torchvision/models/segmentation/_utils.py”, line 26, in forward
x = features[“out”]
TypeError: new(): invalid data type ‘str’

This error is not returned with the base deeplabv3 models provided by trochvision (resnet or mobilenet)

When printing these models (the base deeplbav3_mobilenetv3_large and the custom one) with print(model), as far as I can tell they look exactly the same. Yet one works fine and the other doesn’t.

I have tried multiple other ways of instantiating a custom backbone i.e.:

backbone = torchvision.models.efficientnet_v2_s(weights="DEFAULT").features
classifier = torchvision.models.segmentation.deeplabv3.DeepLabHead(1280, outputchannels)
model = torchvision.models.segmentation.DeepLabV3(backbone, classifier, None)

It gives the same error.

Where might I be going wrong here?