Removing classification layer for resnet101-deeplabv3

2

Would you like to remove the very last layer in this model or the complete classification head (DeepLabHead)?
In the former case, you could just set model.classifier[4] to an nn.Identity layer:

model = models.segmentation.deeplabv3_resnet101()
model.classifier[4] = nn.Identity()

If you would like to change this layer for a custom one using a new number of classes, just replace it with

nn.Conv2d(
    in_channels=256,
    out_channels=nb_classes,
    kernel_size=1,
    stride=1
)
3 Likes