For each layer of my model, I have a list of tensors corresponding to permutations for some experiments I’m running. These permutations are randomly generated at initialization. I realized that since I stored them in a regular python list, they were not saved when I saved and loaded the model stat dict. I would like to save and load them though. Should I use nn.ParameterList() or nn.ModuleList()? These permutation tensors are not meant to be trained.
Thanks.
It depends what exactly you are trying to store. nn.ParameterList
is used to save nn.Parameter
s, i.e. trainable tensors, while nn.ModuleList
is used to store nn.Module
s, which themselves could contain other modules or parameters.
Thank you for the reply. The tensors I want to store are integer tensors, so when I try to save them in a parameter list I get “RuntimeError: Only Tensors of floating point and complex dtype can require gradients”.
shuffles = []
for head in range(heads-1):
perm = list(range(dim)) + random.sample(range(dim, l*dim), (l-1)*dim)
shuffles.append(torch.tensor(perm, requires_grad = False))
self.shuffles = nn.ParameterList([shuffle for shuffle in shuffles])
This is the code I’m using that’s giving me an error.
For non-floating point tensors you should register them as a buffer as seen here:
class MyModel(nn.Module):
def __init__(self):
super().__init__()
self.register_buffer("my_buffer", torch.randint(0, 10, (10,)))
m = MyModel()
print(m.state_dict())
# OrderedDict([('my_buffer', tensor([6, 0, 8, 4, 6, 3, 0, 1, 6, 7]))])
1 Like