Saving list of dictionaries of tensors as model parameters

Suppose I have a list of dictionaries of Tensors, which are used (but not necessarily trainable) in a model’s computation. What is the best way to save this list of dictionaries of Tensors when I save the model?

1 Like

What do you mean by the best ? fastest ? memory optimized ? most compact ? simplest ?I think it depends on your context. I guess JSON serialization or pickling would be good enough for the vast majority of use cases.

Indeed you can have many definitions of ‘best’. I do know that you can save as this list by pickle or json but then it will be separate from the model (Correct me if I am wrong). I am looking for a way to save this list as model parameters, rather than a separate object.

You can put whatever you want in the state_dict, here is an example:

import torch
net = torch.nn.Linear(2, 2)
d = net.state_dict()
d.update({"custom_stuff": 123})
torch.save(d, "net.pth")

# Load it from the serialized file

new_net = torch.nn.Linear(2, 2)
# By default, strict=True, then you need to force it to False to not raise an error
state_dict = torch.load("net.pth")
custom_stuff = state_dict["custom_stuff"] # Retrieving your stuff after loading
new_net.load_state_dict(state_dict, strict=False) 
# NOTE: you can just del state_dict["custom_stuff"] if you want, then strict can be set True

# Check that the parameters have been loaded as expected

print(net.weight.data.equal(new_net.weight.data))
print(net.bias.data.equal(new_net.bias.data))

1 Like

This looks good, Thanks :smiley: