I have created the model and save the weights using google colab. Now I have created a prediction script. The prediction script contains the model class. I am trying to load the model weights using the following method-
Saving & Loading Model Across Devices
Save on GPU, Load on CPU Save:
torch.save(model.state_dict(), PATH)
Load:
device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))
The above method should work, right? Yes.
But when I am trying to do so I have different parameters of the model in Google Colab (Prediction, runtime-None, device=CPU) and different in my local machine (prediction, device=cpu)
Model Params in Colab-
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f'The model has {count_parameters(model):,} trainable parameters')
The model has 12,490,234 trainable parameters
+-------------------------------------------------------+------------+
| Modules | Parameters |
+-------------------------------------------------------+------------+
| encoder.tok_embedding.weight | 2053376 |
| encoder.pos_embedding.weight | 25600 |
| encoder.layers.0.self_attn_layer_norm.weight | 256 |
| encoder.layers.0.self_attn_layer_norm.bias | 256 |
| encoder.layers.0.ff_layer_norm.weight | 256 |
| encoder.layers.0.ff_layer_norm.bias | 256 |
| encoder.layers.0.self_attention.fc_q.weight | 65536 |
| encoder.layers.0.self_attention.fc_q.bias | 256 |
| encoder.layers.0.self_attention.fc_k.weight | 65536 |
| encoder.layers.0.self_attention.fc_k.bias | 256 |
| encoder.layers.0.self_attention.fc_v.weight | 65536 |
| encoder.layers.0.self_attention.fc_v.bias | 256 |
| encoder.layers.0.self_attention.fc_o.weight | 65536 |
| encoder.layers.0.self_attention.fc_o.bias | 256 |
| encoder.layers.0.positionwise_feedforward.fc_1.weight | 131072 |
| encoder.layers.0.positionwise_feedforward.fc_1.bias | 512 |
| encoder.layers.0.positionwise_feedforward.fc_2.weight | 131072 |
| encoder.layers.0.positionwise_feedforward.fc_2.bias | 256 |
| encoder.layers.1.self_attn_layer_norm.weight | 256 |
| encoder.layers.1.self_attn_layer_norm.bias | 256 |
| encoder.layers.1.ff_layer_norm.weight | 256 |
| encoder.layers.1.ff_layer_norm.bias | 256 |
| encoder.layers.1.self_attention.fc_q.weight | 65536 |
| encoder.layers.1.self_attention.fc_q.bias | 256 |
| encoder.layers.1.self_attention.fc_k.weight | 65536 |
| encoder.layers.1.self_attention.fc_k.bias | 256 |
| encoder.layers.1.self_attention.fc_v.weight | 65536 |
| encoder.layers.1.self_attention.fc_v.bias | 256 |
| encoder.layers.1.self_attention.fc_o.weight | 65536 |
| encoder.layers.1.self_attention.fc_o.bias | 256 |
| encoder.layers.1.positionwise_feedforward.fc_1.weight | 131072 |
| encoder.layers.1.positionwise_feedforward.fc_1.bias | 512 |
| encoder.layers.1.positionwise_feedforward.fc_2.weight | 131072 |
| encoder.layers.1.positionwise_feedforward.fc_2.bias | 256 |
| encoder.layers.2.self_attn_layer_norm.weight | 256 |
| encoder.layers.2.self_attn_layer_norm.bias | 256 |
| encoder.layers.2.ff_layer_norm.weight | 256 |
| encoder.layers.2.ff_layer_norm.bias | 256 |
| encoder.layers.2.self_attention.fc_q.weight | 65536 |
| encoder.layers.2.self_attention.fc_q.bias | 256 |
| encoder.layers.2.self_attention.fc_k.weight | 65536 |
| encoder.layers.2.self_attention.fc_k.bias | 256 |
| encoder.layers.2.self_attention.fc_v.weight | 65536 |
| encoder.layers.2.self_attention.fc_v.bias | 256 |
| encoder.layers.2.self_attention.fc_o.weight | 65536 |
| encoder.layers.2.self_attention.fc_o.bias | 256 |
| encoder.layers.2.positionwise_feedforward.fc_1.weight | 131072 |
| encoder.layers.2.positionwise_feedforward.fc_1.bias | 512 |
| encoder.layers.2.positionwise_feedforward.fc_2.weight | 131072 |
| encoder.layers.2.positionwise_feedforward.fc_2.bias | 256 |
| decoder.tok_embedding.weight | 3209728 |
| decoder.pos_embedding.weight | 25600 |
| decoder.layers.0.self_attn_layer_norm.weight | 256 |
| decoder.layers.0.self_attn_layer_norm.bias | 256 |
| decoder.layers.0.enc_attn_layer_norm.weight | 256 |
| decoder.layers.0.enc_attn_layer_norm.bias | 256 |
| decoder.layers.0.ff_layer_norm.weight | 256 |
| decoder.layers.0.ff_layer_norm.bias | 256 |
| decoder.layers.0.self_attention.fc_q.weight | 65536 |
| decoder.layers.0.self_attention.fc_q.bias | 256 |
| decoder.layers.0.self_attention.fc_k.weight | 65536 |
| decoder.layers.0.self_attention.fc_k.bias | 256 |
| decoder.layers.0.self_attention.fc_v.weight | 65536 |
| decoder.layers.0.self_attention.fc_v.bias | 256 |
| decoder.layers.0.self_attention.fc_o.weight | 65536 |
| decoder.layers.0.self_attention.fc_o.bias | 256 |
| decoder.layers.0.encoder_attention.fc_q.weight | 65536 |
| decoder.layers.0.encoder_attention.fc_q.bias | 256 |
| decoder.layers.0.encoder_attention.fc_k.weight | 65536 |
| decoder.layers.0.encoder_attention.fc_k.bias | 256 |
| decoder.layers.0.encoder_attention.fc_v.weight | 65536 |
| decoder.layers.0.encoder_attention.fc_v.bias | 256 |
| decoder.layers.0.encoder_attention.fc_o.weight | 65536 |
| decoder.layers.0.encoder_attention.fc_o.bias | 256 |
| decoder.layers.0.positionwise_feedforward.fc_1.weight | 131072 |
| decoder.layers.0.positionwise_feedforward.fc_1.bias | 512 |
| decoder.layers.0.positionwise_feedforward.fc_2.weight | 131072 |
| decoder.layers.0.positionwise_feedforward.fc_2.bias | 256 |
| decoder.layers.1.self_attn_layer_norm.weight | 256 |
| decoder.layers.1.self_attn_layer_norm.bias | 256 |
| decoder.layers.1.enc_attn_layer_norm.weight | 256 |
| decoder.layers.1.enc_attn_layer_norm.bias | 256 |
| decoder.layers.1.ff_layer_norm.weight | 256 |
| decoder.layers.1.ff_layer_norm.bias | 256 |
| decoder.layers.1.self_attention.fc_q.weight | 65536 |
| decoder.layers.1.self_attention.fc_q.bias | 256 |
| decoder.layers.1.self_attention.fc_k.weight | 65536 |
| decoder.layers.1.self_attention.fc_k.bias | 256 |
| decoder.layers.1.self_attention.fc_v.weight | 65536 |
| decoder.layers.1.self_attention.fc_v.bias | 256 |
| decoder.layers.1.self_attention.fc_o.weight | 65536 |
| decoder.layers.1.self_attention.fc_o.bias | 256 |
| decoder.layers.1.encoder_attention.fc_q.weight | 65536 |
| decoder.layers.1.encoder_attention.fc_q.bias | 256 |
| decoder.layers.1.encoder_attention.fc_k.weight | 65536 |
| decoder.layers.1.encoder_attention.fc_k.bias | 256 |
| decoder.layers.1.encoder_attention.fc_v.weight | 65536 |
| decoder.layers.1.encoder_attention.fc_v.bias | 256 |
| decoder.layers.1.encoder_attention.fc_o.weight | 65536 |
| decoder.layers.1.encoder_attention.fc_o.bias | 256 |
| decoder.layers.1.positionwise_feedforward.fc_1.weight | 131072 |
| decoder.layers.1.positionwise_feedforward.fc_1.bias | 512 |
| decoder.layers.1.positionwise_feedforward.fc_2.weight | 131072 |
| decoder.layers.1.positionwise_feedforward.fc_2.bias | 256 |
| decoder.layers.2.self_attn_layer_norm.weight | 256 |
| decoder.layers.2.self_attn_layer_norm.bias | 256 |
| decoder.layers.2.enc_attn_layer_norm.weight | 256 |
| decoder.layers.2.enc_attn_layer_norm.bias | 256 |
| decoder.layers.2.ff_layer_norm.weight | 256 |
| decoder.layers.2.ff_layer_norm.bias | 256 |
| decoder.layers.2.self_attention.fc_q.weight | 65536 |
| decoder.layers.2.self_attention.fc_q.bias | 256 |
| decoder.layers.2.self_attention.fc_k.weight | 65536 |
| decoder.layers.2.self_attention.fc_k.bias | 256 |
| decoder.layers.2.self_attention.fc_v.weight | 65536 |
| decoder.layers.2.self_attention.fc_v.bias | 256 |
| decoder.layers.2.self_attention.fc_o.weight | 65536 |
| decoder.layers.2.self_attention.fc_o.bias | 256 |
| decoder.layers.2.encoder_attention.fc_q.weight | 65536 |
| decoder.layers.2.encoder_attention.fc_q.bias | 256 |
| decoder.layers.2.encoder_attention.fc_k.weight | 65536 |
| decoder.layers.2.encoder_attention.fc_k.bias | 256 |
| decoder.layers.2.encoder_attention.fc_v.weight | 65536 |
| decoder.layers.2.encoder_attention.fc_v.bias | 256 |
| decoder.layers.2.encoder_attention.fc_o.weight | 65536 |
| decoder.layers.2.encoder_attention.fc_o.bias | 256 |
| decoder.layers.2.positionwise_feedforward.fc_1.weight | 131072 |
| decoder.layers.2.positionwise_feedforward.fc_1.bias | 512 |
| decoder.layers.2.positionwise_feedforward.fc_2.weight | 131072 |
| decoder.layers.2.positionwise_feedforward.fc_2.bias | 256 |
| decoder.fc_out.weight | 3209728 |
| decoder.fc_out.bias | 12538 |
+-------------------------------------------------------+------------+
Total Trainable Params: 12490234
Model Params in Local-
def count_parameters(model):
return sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f'The model has {count_parameters(model):,} trainable parameters')
The model has 12,506,137 trainable parameters
+-------------------------------------------------------+------------+
| Modules | Parameters |
+-------------------------------------------------------+------------+
| encoder.tok_embedding.weight | 2053376 |
| encoder.pos_embedding.weight | 25600 |
| encoder.layers.0.self_attn_layer_norm.weight | 256 |
| encoder.layers.0.self_attn_layer_norm.bias | 256 |
| encoder.layers.0.ff_layer_norm.weight | 256 |
| encoder.layers.0.ff_layer_norm.bias | 256 |
| encoder.layers.0.self_attention.fc_q.weight | 65536 |
| encoder.layers.0.self_attention.fc_q.bias | 256 |
| encoder.layers.0.self_attention.fc_k.weight | 65536 |
| encoder.layers.0.self_attention.fc_k.bias | 256 |
| encoder.layers.0.self_attention.fc_v.weight | 65536 |
| encoder.layers.0.self_attention.fc_v.bias | 256 |
| encoder.layers.0.self_attention.fc_o.weight | 65536 |
| encoder.layers.0.self_attention.fc_o.bias | 256 |
| encoder.layers.0.positionwise_feedforward.fc_1.weight | 131072 |
| encoder.layers.0.positionwise_feedforward.fc_1.bias | 512 |
| encoder.layers.0.positionwise_feedforward.fc_2.weight | 131072 |
| encoder.layers.0.positionwise_feedforward.fc_2.bias | 256 |
| encoder.layers.1.self_attn_layer_norm.weight | 256 |
| encoder.layers.1.self_attn_layer_norm.bias | 256 |
| encoder.layers.1.ff_layer_norm.weight | 256 |
| encoder.layers.1.ff_layer_norm.bias | 256 |
| encoder.layers.1.self_attention.fc_q.weight | 65536 |
| encoder.layers.1.self_attention.fc_q.bias | 256 |
| encoder.layers.1.self_attention.fc_k.weight | 65536 |
| encoder.layers.1.self_attention.fc_k.bias | 256 |
| encoder.layers.1.self_attention.fc_v.weight | 65536 |
| encoder.layers.1.self_attention.fc_v.bias | 256 |
| encoder.layers.1.self_attention.fc_o.weight | 65536 |
| encoder.layers.1.self_attention.fc_o.bias | 256 |
| encoder.layers.1.positionwise_feedforward.fc_1.weight | 131072 |
| encoder.layers.1.positionwise_feedforward.fc_1.bias | 512 |
| encoder.layers.1.positionwise_feedforward.fc_2.weight | 131072 |
| encoder.layers.1.positionwise_feedforward.fc_2.bias | 256 |
| encoder.layers.2.self_attn_layer_norm.weight | 256 |
| encoder.layers.2.self_attn_layer_norm.bias | 256 |
| encoder.layers.2.ff_layer_norm.weight | 256 |
| encoder.layers.2.ff_layer_norm.bias | 256 |
| encoder.layers.2.self_attention.fc_q.weight | 65536 |
| encoder.layers.2.self_attention.fc_q.bias | 256 |
| encoder.layers.2.self_attention.fc_k.weight | 65536 |
| encoder.layers.2.self_attention.fc_k.bias | 256 |
| encoder.layers.2.self_attention.fc_v.weight | 65536 |
| encoder.layers.2.self_attention.fc_v.bias | 256 |
| encoder.layers.2.self_attention.fc_o.weight | 65536 |
| encoder.layers.2.self_attention.fc_o.bias | 256 |
| encoder.layers.2.positionwise_feedforward.fc_1.weight | 131072 |
| encoder.layers.2.positionwise_feedforward.fc_1.bias | 512 |
| encoder.layers.2.positionwise_feedforward.fc_2.weight | 131072 |
| encoder.layers.2.positionwise_feedforward.fc_2.bias | 256 |
| decoder.tok_embedding.weight | 3217664 |
| decoder.pos_embedding.weight | 25600 |
| decoder.layers.0.self_attn_layer_norm.weight | 256 |
| decoder.layers.0.self_attn_layer_norm.bias | 256 |
| decoder.layers.0.enc_attn_layer_norm.weight | 256 |
| decoder.layers.0.enc_attn_layer_norm.bias | 256 |
| decoder.layers.0.ff_layer_norm.weight | 256 |
| decoder.layers.0.ff_layer_norm.bias | 256 |
| decoder.layers.0.self_attention.fc_q.weight | 65536 |
| decoder.layers.0.self_attention.fc_q.bias | 256 |
| decoder.layers.0.self_attention.fc_k.weight | 65536 |
| decoder.layers.0.self_attention.fc_k.bias | 256 |
| decoder.layers.0.self_attention.fc_v.weight | 65536 |
| decoder.layers.0.self_attention.fc_v.bias | 256 |
| decoder.layers.0.self_attention.fc_o.weight | 65536 |
| decoder.layers.0.self_attention.fc_o.bias | 256 |
| decoder.layers.0.encoder_attention.fc_q.weight | 65536 |
| decoder.layers.0.encoder_attention.fc_q.bias | 256 |
| decoder.layers.0.encoder_attention.fc_k.weight | 65536 |
| decoder.layers.0.encoder_attention.fc_k.bias | 256 |
| decoder.layers.0.encoder_attention.fc_v.weight | 65536 |
| decoder.layers.0.encoder_attention.fc_v.bias | 256 |
| decoder.layers.0.encoder_attention.fc_o.weight | 65536 |
| decoder.layers.0.encoder_attention.fc_o.bias | 256 |
| decoder.layers.0.positionwise_feedforward.fc_1.weight | 131072 |
| decoder.layers.0.positionwise_feedforward.fc_1.bias | 512 |
| decoder.layers.0.positionwise_feedforward.fc_2.weight | 131072 |
| decoder.layers.0.positionwise_feedforward.fc_2.bias | 256 |
| decoder.layers.1.self_attn_layer_norm.weight | 256 |
| decoder.layers.1.self_attn_layer_norm.bias | 256 |
| decoder.layers.1.enc_attn_layer_norm.weight | 256 |
| decoder.layers.1.enc_attn_layer_norm.bias | 256 |
| decoder.layers.1.ff_layer_norm.weight | 256 |
| decoder.layers.1.ff_layer_norm.bias | 256 |
| decoder.layers.1.self_attention.fc_q.weight | 65536 |
| decoder.layers.1.self_attention.fc_q.bias | 256 |
| decoder.layers.1.self_attention.fc_k.weight | 65536 |
| decoder.layers.1.self_attention.fc_k.bias | 256 |
| decoder.layers.1.self_attention.fc_v.weight | 65536 |
| decoder.layers.1.self_attention.fc_v.bias | 256 |
| decoder.layers.1.self_attention.fc_o.weight | 65536 |
| decoder.layers.1.self_attention.fc_o.bias | 256 |
| decoder.layers.1.encoder_attention.fc_q.weight | 65536 |
| decoder.layers.1.encoder_attention.fc_q.bias | 256 |
| decoder.layers.1.encoder_attention.fc_k.weight | 65536 |
| decoder.layers.1.encoder_attention.fc_k.bias | 256 |
| decoder.layers.1.encoder_attention.fc_v.weight | 65536 |
| decoder.layers.1.encoder_attention.fc_v.bias | 256 |
| decoder.layers.1.encoder_attention.fc_o.weight | 65536 |
| decoder.layers.1.encoder_attention.fc_o.bias | 256 |
| decoder.layers.1.positionwise_feedforward.fc_1.weight | 131072 |
| decoder.layers.1.positionwise_feedforward.fc_1.bias | 512 |
| decoder.layers.1.positionwise_feedforward.fc_2.weight | 131072 |
| decoder.layers.1.positionwise_feedforward.fc_2.bias | 256 |
| decoder.layers.2.self_attn_layer_norm.weight | 256 |
| decoder.layers.2.self_attn_layer_norm.bias | 256 |
| decoder.layers.2.enc_attn_layer_norm.weight | 256 |
| decoder.layers.2.enc_attn_layer_norm.bias | 256 |
| decoder.layers.2.ff_layer_norm.weight | 256 |
| decoder.layers.2.ff_layer_norm.bias | 256 |
| decoder.layers.2.self_attention.fc_q.weight | 65536 |
| decoder.layers.2.self_attention.fc_q.bias | 256 |
| decoder.layers.2.self_attention.fc_k.weight | 65536 |
| decoder.layers.2.self_attention.fc_k.bias | 256 |
| decoder.layers.2.self_attention.fc_v.weight | 65536 |
| decoder.layers.2.self_attention.fc_v.bias | 256 |
| decoder.layers.2.self_attention.fc_o.weight | 65536 |
| decoder.layers.2.self_attention.fc_o.bias | 256 |
| decoder.layers.2.encoder_attention.fc_q.weight | 65536 |
| decoder.layers.2.encoder_attention.fc_q.bias | 256 |
| decoder.layers.2.encoder_attention.fc_k.weight | 65536 |
| decoder.layers.2.encoder_attention.fc_k.bias | 256 |
| decoder.layers.2.encoder_attention.fc_v.weight | 65536 |
| decoder.layers.2.encoder_attention.fc_v.bias | 256 |
| decoder.layers.2.encoder_attention.fc_o.weight | 65536 |
| decoder.layers.2.encoder_attention.fc_o.bias | 256 |
| decoder.layers.2.positionwise_feedforward.fc_1.weight | 131072 |
| decoder.layers.2.positionwise_feedforward.fc_1.bias | 512 |
| decoder.layers.2.positionwise_feedforward.fc_2.weight | 131072 |
| decoder.layers.2.positionwise_feedforward.fc_2.bias | 256 |
| decoder.fc_out.weight | 3217664 |
| decoder.fc_out.bias | 12569 |
+-------------------------------------------------------+------------+
Total Trainable Params: 12506137
So, that’s why I am unable to load the model. Because the model has a different parameter in local.
Even if I try to load the weights in local it gives me-
model.load_state_dict(torch.load(f"{model_name}.pt", map_location=device))
Error-
--------------------------------------------------------------------------- RuntimeError Traceback (most recent call last) <ipython-input-24-f5baac4441a5> in <module>
----> 1 model.load_state_dict(torch.load(f"{model_name}_2.pt", map_location=device))
c:\anaconda\envs\lang_trans\lib\site-packages\torch\nn\modules\module.py in load_state_dict(self, state_dict, strict)
845 if len(error_msgs) > 0:
846 raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
--> 847 self.__class__.__name__, "\n\t".join(error_msgs)))
848 return _IncompatibleKeys(missing_keys, unexpected_keys)
849
RuntimeError: Error(s) in loading state_dict for Seq2Seq: size mismatch for decoder.tok_embedding.weight: copying a param with shape torch.Size([12538, 256]) from checkpoint, the shape in current model is torch.Size([12569, 256]). size mismatch for decoder.fc_out.weight: copying a param with shape torch.Size([12538, 256]) from checkpoint, the shape in current model is torch.Size([12569, 256]). size mismatch for decoder.fc_out.bias: copying a param with shape torch.Size([12538]) from checkpoint, the shape in current model is torch.Size([12569]).
The model param of the local must be wrong because in colab (device=CPU, runtime=None) I am able to load the weights after defining model class. But in the local machine the params changes, so I am unable to load the weights. I know it’s weird, help me to find the solution.
You can check the full code of the model here-
https://gist.github.com/Dipeshpal/90c715a7b7f00845e20ef998bda35835
After this model params change.