Updating a register_buffer in PyTorch?

user_12184561 0
Viewed 2161

Correct way to update a register_buffer in PyTorch

I'm trying to determine the recommended way to update a register buffer which preserves the buffer's attributes. One "hacky" way to do this is shown below:

import torch
import torch.nn as nn


class SomeModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.register_buffer("A", torch.tensor(2.0))

    def forward(self, x):
        return x * self.A

    def update_buffer(self, v):
        A = getattr(self, "A")
        setattr(self, "A", torch.tensor(v, dtype=A.dtype, device=A.device))


if __name__ == "__main__":
    model = SomeModule()
    model.to(torch.float64)
    print(model(1.0).dtype)
    model.update_buffer(3.0)
    print(model(1.0).dtype)

When updating many buffers this gets a bit messy. Is there a recommended / better method to accomplish this?

pytorch python
1 Answers

Generally:

self.A = new_tensor

To copy the dtype and device from the previous tensor:

self.A = new_tensor.type_as(self.A)

If the new tensor has the same shape as the old, or you want to overwrite the current tensor's data with a new value, you can do:

self.A[:] = new_value

Examples:

>>> mod.A = torch.tensor([3, 2, 1], dtype=int, device="cuda")
>>> mod.A
tensor([3, 2, 1], device='cuda:0')
>>> mod.A = torch.Tensor([1.5, 2, 3.5]).type_as(mod.A)
>>> mod.A
tensor([1, 2, 3], device='cuda:0')
>>> mod.A[:] = torch.Tensor([3, 4, 5])
>>> mod.A
tensor([3, 4, 5], device='cuda:0')
>>> mod.A[:] = 7                      
>>> mod.A
tensor([7, 7, 7], device='cuda:0')
>>> mod.state_dict()
OrderedDict([('A', tensor([7, 7, 7], device='cuda:0'))])
user_164091230
Related