Cannot load state_dict in Docker

deployment
1

I have a weird problem. I can load the state_dict in my machine which was trained with Pytorch lightning on the GPU. However if i push it to docker it does not load. The model is a Pytorch Lightning module.

On Machine

loaded_model = LitModel()
loaded_model.load_state_dict(torch.load(‘models/trained-model-2021.pt’), strict=False)

Parameter containing:
tensor([[ 0.0817, -0.2579, -0.0080, …, 0.1619, -0.0783, -0.0087],
[ 0.0866, -0.2708, 0.2528, …, 0.0203, -0.2319, 0.1011],
[ 0.0564, 0.0195, -0.2858, …, -0.1899, 0.1004, -0.1013],
…,
[ 0.2584, -0.2400, -0.0358, …, -0.1391, 0.2278, 0.0264],
[ 0.1901, -0.3056, -0.2537, …, -0.0747, 0.2703, 0.0803],
[-0.2639, -0.0381, 0.0538, …, 0.2504, 0.2152, 0.1084]],
requires_grad=True)

On Docker

device = torch.device(‘cpu’)
loaded_model = LitModel()
loaded_model.load_state_dict(torch.load(‘models/trained-model-2021.pt’, map_location=device), strict=False)

loaded_model.load_state_dict(torch.load(‘src/trained-model-2021-11-23–23:08-hs.pt’, map_location=device), strict=True)
Traceback (most recent call last):
File “”, line 1, in
File “/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py”, line 1407, in load_state_dict
self.class.name, “\n\t”.join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for LitModel:
Missing key(s) in state_dict: “layer_1.weight”, “layer_1.bias”, “layer_3.weight”, “layer_3.bias”.
Unexpected key(s) in state_dict: “model”.

Any ideas what could be wrong here.

2

Based on the error message it seems you are not loading a state_dict, but a custom dict class containing the 'model' key, so you might want to access it first.
Something like this could work:

loaded_model.load_state_dict(torch.load(‘src/trained-model-2021-11-23–23:08-hs.pt’, map_location=device)['model'])
1 Like
3

Thanks. Thank works.