Transfer learning unable to do

[ravitejarj]

Hi @zhanghang1989
I am unable to do transfer learning on the model
i have downloaded get_deeplab_resnest101_ade and
when i changed the no of classes in ade20k.py (no of classes 8)
pre trained model is not loading ( getting error )

So I have changed the code for Transfer learning

Code changes:

1. deeplab.py
   in get_deeplab_resnest101_ade function changed from
   from
   #model = DeepLabV3(datasets[dataset.lower()].NUM_CLASS, backbone=backbone, root=root, **kwargs)
   to
   no_of_classes = 150
   model = DeepLabV3(no_of_classes, backbone=backbone, root=root, **kwargs)

so i can load pretrained model with 150 classes
then
2) In train_dist.py file

Model loading

model_ft = get_deeplab_resnest101_ade(pretrained=True)

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

model_ft.head.block = Sequential(
(Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)),
(BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)),
(ReLU(inplace=True)),
(Dropout(p=0.1, inplace=False)),
(Conv2d(256, 8, kernel_size=(1, 1), stride=(1, 1))))

for param in model_ft.head.parameters():
param.requires_grad = True

Training

python train_dist.py --dataset ade20k --model deeplab --aux --backbone resnest101 --ft --epochs 100

after successful training
i am getting 150 classes output not 8 classes(i have given 8 classes in last layer)
i need 8 classes output

can you help me with this

[zhanghang1989]

An easy solution is set strict=False when loading the pretrained model
https://pytorch.org/docs/master/generated/torch.nn.Module.html?highlight=load_state_dict#torch.nn.Module.load_state_dict

[ravitejarj]

Thank you