Is there a way to save a fixed pytorch tensor in a pytorch nn module as a parameter?

Admin 1
Viewed 85

So I have a fixed tensor, which are centroids that I pre-calculated. Essentially the centroids matrix is N classes by D dimensions. I want to save this in a Pytorch module at the end of training in such a way that when I save this model and load it this parameter will be there for me to use after loading the saved checkpoint. Also there should never be a gradient passing through this matrix.

pytorch python
1 Answers

You should probably register it as a buffer. Buffers will not be updated with gradient descent, but will be saved alongside the model. It's typically used for the mean and variance of the batchnorm layers, which are not learnt with backprop.

That would be like :

# replace this with your actual model
my_model = nn.Linear(10,10)
# Replace this tensor with your actual centroids
centroid_tensor = torch.randn(2,2)
# Keep this to register the centroids as a buffer
my_model.register_buffer(name="centroids", tensor=centroids_tensor)
# if you want to access your buffers:
for buf in my_model.named_buffers():
    print(buf)

As explained in the documentation link above, the registered buffer will be persistent by default, which means it will be included in the state dict of the module ; thus it will be saved and loaded alongside the model.

user_109357170
Related