Loading Pretrained Weights

1

Hi,
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.

2

@ptrblck do you have any idea?

3

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
  d[key]=map_402_to_206(d[key])
model.load_state_dict(d)

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
4

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.

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

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

6

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.

7

How can I delete the keys?

8

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

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

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

9

Thanks, bur what would PREFIX be?

10

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”]