pixel-level-contrastive-learning
pixel-level-contrastive-learning copied to clipboard
Extending the model for Semantic Segmentation task
Hi,
First of all, thanks for the repo. I am new to this area but was able to use the repo quite easily.
My downstream task is semantic segmentation and I am wondering if I can just change last layers to do the task.
In particular, this is what I did
- I have a large unlabeled dataset which I am successfully able to pass through the model (it accepts only 3 channeled images, mine are grayscale images, so I just stacked the same image for all channels. I hope that's ok!). My understanding is that, if this is done, the pretext task is completed :)
- Now, for downstream task, I am wondering how to change the last layer in 'improved-resnet.pth' and add new layers so that, I can pass new data (small quantity, semantic labeled) and finetune the model for segmentation. In particular, I don't know what layers to add for segmentation task (removing the last Linear layer can be done using children())
Any help would be greatly appreciated.
Thanks a lot
You can use it as a backbone for a segmentation model just like a regular resnet for example : Mask RCNN with resnet as backbone
Thanks Spinkoo for your comment.
I tried to add FCN head by making this model as backbone.
- First I am freezing the layers (so for downstream task these weights will not be trained)
improved_resnet = torch.load('improved-resnet.pt')
for param in improved_resnet.parameters():
param.requires_grad = False
- Then I took the fcn-classifier head, added a transpose2d conv and attached it to the end of the given resnet model
fcn_model = torch.hub.load('pytorch/vision:v0.10.0', 'fcn_resnet50', pretrained=False)
fcn_model.classifier[4] = torch.nn.ConvTranspose2d(512, 2, stride = (8,8), kernel_size = (8,8))
improved_resnet.avgpool = fcn_model.classifier
improved_resnet.fc = Identity()
Note: Because somewhere pool/fc, the dimensions are mismatching, I had to upconvolute.
This is the summary of the model when a (16, 3, 512, 512) is passed (16 batch size, 3 channels, only last layer output is placed for brevity)
│ │ └─ReLU: 3-148 [16, 2048, 16, 16] -- ├─FCNHead: 1-9 [16, 2, 128, 128] -- │ └─Conv2d: 2-17 [16, 512, 16, 16] 9,437,184 │ └─BatchNorm2d: 2-18 [16, 512, 16, 16] 1,024 │ └─ReLU: 2-19 [16, 512, 16, 16] -- │ └─Dropout: 2-20 [16, 512, 16, 16] -- │ └─ConvTranspose2d: 2-21 [16, 2, 128, 128] 65,538 ├─Identity: 1-10 [16, 32768] --
I have the following doubts:
- Is my approach to add a segmentation head on top of this model correct?
- I don't understand why it's flattening even though I made the flatten as Identity()
if my segmentor is an encoder-decoder, should i use the encoder as the encoder-online ? the decoder as the MLP ?