Fine tuning FCN_ResNet101

Hi everyone,

I want to finetune a FCN_ResNet101. I would like to change the last layer as my dataset has a different number of classes. Also, finetune only the FCN head.
For the former, is it enough to only change the num_classes argument when defining the model or I need to use something like this:

model = torchvision.models.segmentation.fcn_resnet101(pretrained=True)
model.classifier=nn.identity()
model.Conv2d = nn.Conv2d(
in_channels=256,
out_channels=nb_classes,
kernel_size=1,
stride=1
)

I took this piece of code from another thread. I am not sure if it is necessary to use nn.identity(). When I do, the last layer does not change but the last layer of the one to the last FCN!
And, how many layers must be changed so my FCN head is re_trianed?
I wrote it this way but I’m mostly confused about FCN_ResNet101 architecture.

model = torchvision.models.segmentation.fcn_resnet101(pretrained=True, progress=True, num_classes=?)
#model.classifier[4] = nn.Identity()

“”"
FCNHead(
(0): Conv2d(2048, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
(3): Dropout(p=0.1)
(4): Conv2d(512, 21, kernel_size=(1, 1), stride=(1, 1))
), FCNHead(
(0): Conv2d(1024, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(2): ReLU()
(3): Dropout(p=0.1)
(4): Conv2d(256, 21, kernel_size=(1, 1), stride=(1, 1))
)]
“”"

#setting our own number of classes

layer_list = list(model.children())[-5:]
model_small = nn.Sequential(*list(model.children()))[-5:]

for param in model_small.parameters():
param.requires_grad = False

model_small.Conv2d = nn.Conv2d( in_channels=1024,kernel_size=(3,3),stride=(1,1))
model_small.BatchNorm2d = nn.BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
model_small.ReLU = nn.ReLU()
model_small.Dropout = nn.Dropout(p=0.1)
model_small.Conv2d = nn.Conv2d(
in_channels=256,
out_channels=nb_classes,
kernel_size=1,
stride=1
)
model = model_small.to(device)
Any guidance is very much appreciated!

No, the nn.Identity module is not necessary in your use case, but can be used to extract features.
I would recommend to stick to the Fine tuning tutorial. :wink: