How to modify layers at the end of a Segmentation network in Pytorch

Hi everyone. As the title suggests, I want to modify the last few layers of the Semantic Segmentation networks provided by Pytorch. I want to use the output coming from the last ReLu layer to my own layer. I’m really having a tough time getting it working. So I’m trying to get the network output at the last ReLu layer which is just before the output layer.

t = torchvision.models.segmentation.deeplabv3_resnet50()
wow = nn.Sequential(*list(t.modules())[:-1])
wow(torch.ones((2, 3, 224, 224)).cuda())

The above code gives me an error saying

TypeError: conv2d(): argument 'input' (position 1) must be Tensor, not collections.OrderedDict

I’m really not able to understand where I am going wrong. So what I want is the output at the last but one layer. But I do not know how to access it.

nn.Sequential containers can be used for simple feed forward models, where the layers are just connected sequentially and pass a single output to the next layer.
Often models are more complicated than this approach and cannot be easily mapped to nn.Sequential, which also seems to be the case here.
Based on the error message I assume that a particular layer outputs an OrderedDict, which is then processed in the original forward method and creates the error in your approach.

To remove some layers I would thus recommend to replace them directly with your new layers using the layer name:

model.layer = MyLayers()

or if you want to disable them assign an nn.Identity() module to the corresponding attribute.

Thank you for your reply. Yeah, when I wanted to change the individual layers, I get a different output than when I’m running the whole network together. The segmentation.py file has a list of boolean variables which are used to change strided convolutions with dilated ones, which I don’t know how to do it while accessing individual modules of the segmentation network. Also, I just want to add a layer to the end. I guess I have to write the whole network again.