Training deeplabv3 with custom classifier

vision
1

I attempting to create a deeplabv3 model with an effnet backbone. However I’m struggling with connecting a backbone not of the pre specified ones (i.e. torchvision.models.segmentation.deeplabv3_mobilenetv3_large etc)

I have:

def createdeeplabv3_effnet(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.
    """
    backbone = torchvision.models.efficientnet_b0(pretrained=True, progress=True)
    model = torchvision.models.segmentation.DeepLabV3(backbone, None, None)
    #model = torchvision.models.segmentation.deeplabv3_resnet50(pretrained=True, progress=True)
    model.classifier = torchvision.models.segmentation.deeplabv3.DeepLabHead(1280, outputchannels)
    model.aux_classifier = torchvision.models.segmentation.deeplabv3.DeepLabHead(1280, outputchannels)
    # Set the model in training mode
    model.train()
    return model

This returns this error:
lib/python3.9/site-packages/torchvision/models/segmentation/_utils.py", line 28, in forward
x = features[“out”]
IndexError: too many indices for tensor of dimension 2

While this:

def createdeeplabv3(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

Works perfectly fine.