Loading the pretrained network on different number of labels

vision
1

Hi,
I trained a network using transfer learning to classify 2 labels.
Then I tried to use the trained weights from the above case to retrain another problem to classify 3 labels.

Now, I am getting a shape mismatch error, obviously, because the endpoint classes are different.
Is it possible to load the weights in the way I am thinking,
How to resolve the error?

def load_checkpoint(checkpoint_path, model):
    """To Load model."""
    checkpoint = torch.load(checkpoint_path, map_location='cuda:0' if torch.cuda.is_available() else 'cpu') # noqa
    model.load_state_dict(checkpoint['state_dict'])
    return model

# To load the architecture of the model to be loaded
model_transfer1 = EfficientNet.from_pretrained('efficientnet-b2',weights_path='/home/rxs1576/latest_scripts/Project_QA/EfficientNetPytorch/efficientnet-b2-8bb594d6.pth') # noqa
n_inputs = model_transfer1._fc.in_features
model_transfer1._fc = nn.Linear(n_inputs, 3)

# Loading my retrained model
model_transfer = load_checkpoint(checkpoint_path='My retrained model',model = model_transfer1)

# freezing the entire network
for name, parameter in model_transfer.named_parameters():
    parameter.requires_grad = False


# %%
# Defreezing the end points of the network so I can train them again.
update_params_name = ['_fc.weight', '_fc.bias', '_conv_head.weight']
for name, parameter in model_transfer.named_parameters():
    if name in update_params_name:
        parameter.requires_grad = True
2

If the goal is to finetune the network after changing the last layer, can you do the modification after loading the weights? For example, can you move the line

model_transfer1._fc = nn.Linear(n_inputs, 3)

to after

model_transfer = load_checkpoint(checkpoint_path='My retrained model',model = model_transfer1)

?

1 Like
3

Thanks, I have to do the step twice, first to remove the 1000 labels(ImageNet dataset) to 2(My trained network) and then to 3 (my desired finetuning network).