Changing dilation rate of backbones

[vedantdalimkar]

Is there any way to change the dilation rate of the convolution layers in the last few layers (like DeepLabV3 style) through the SMP API?

[vitormbesen]

The most likely way to solve this is by utilizing the some PyTorch API and Python class manipulation. As you can see in the DeepLabV3 class, there is a self.encoder module (and in all pretty much all models of this library). You should instantiate your DeepLabV3 model, and modify the encoder.

import segmentation_models_pytorch as smp

model = smp.DeepLabV3(
    encoder_name="resnet34",  # or any other backbone
    encoder_weights="imagenet",
    in_channels=3,
    classes=1,
)

model.encoder = <modified_encoder>

For simplicity, let's assume the encoder was created with timm.create_model(). The following function modifies the dilation rate of every Conv2d module present in the resnet34 "layer4" backbone. You should be careful, as this may require you to manually change other attributes, such as kernel_size, stride, and padding.

import torch
import torch.nn as nn
import timm

encoder = timm.create_model('resnet34')

def update_dilation(module):
    for name, child in module.named_children():
        if isinstance(child, nn.Conv2d):
            child.dilation = (2, 2)
            child.kernel_size = ... 
        else:
            update_dilation(child)
    return module

module = encoder.layer4
updated_module = update_dilation(module)
encoder.layer4 = updated_module
print(encoder)  # to see the changes made

If you want to modify specific convolutions in different layers, you need to understand the backbone structure you're working with an perhaps manually access them in order to make the desired modifications.