Loading Pretrained Weights

I pretrained a custom model with targets 402. Now I need to transfer those weights to a model with 206 targets. How can I do that? This is what I am doing right now model.load_state_dict(torch.load(f"FOLD{fold}_.pth"), strict=False)
But it is not working, showing a size mismatch error.

@ptrblck do you have any idea?

if you mean a subset of targets, something like that:

d = torch.load(f"FOLD{fold}_.pth")
for suffix in (".weight",".bias"):
  key = PREFIX + suffix

PREFIX is output linear layer’s key. map_402_to_206 should reflect how you select a subset, simplest case would be a subrange: d[key] = d[key][:206]

if you want to reuse hidden layers on a different set of targets, delete above keys instead and use strict=False

1 Like

Thanks for your response, I actually want to use for a different set of targets and I am using strict =False but it is throwing a size mismatch error.

    for i in range(len(model_source.layers[:-1])):
    return model_dest

This is the Keras equivalent of what I want to do.

Did you follow @googlebot’s suggestion and deleted the key(s) before using strict=False?
The strict=False argument ignores unexpected or missing keys, not shape mismatch errors.

How can I delete the keys?

You can delete key-value pairs from an OrderedDict using del as seen here:

d = OrderedDict()
d['a'] = 1
d['b'] = 2
del d['a']

However, @googlebot already posted a solution where you would copy some of the pretrained parameters into the state_dict.

Thanks, bur what would PREFIX be?

you can just explore python directory, e.g. list(d.keys()), it will contain [sub]module names as assigned in __init__, and deleting is something like: del d[“output_layer.weight”]