PyTorch register_buffer

通常我們想要建立一個 tensor 時，會用以下方法：

a = torch.tensor(0.1)

假如要在一個 module 中建立一些 tensor，可能會使用如下方法：

class Layer(nn.Module):
    self.__init__(self, num_outputs, num_inputs):
        self.w = nn.Parameter(torch.Tensor(num_outputs, num_inputs))
        self.a = torch.tensor(0.1)

其中的 w 是可以被訓練的參數，而 a 是無法被訓練的。

這時我們會遇到一個問題，當我們想要將整個 module 從 CPU 移動到 GPU 時：

layer = Layer(10, 100)
layer = layer.to('cuda:0')

我們會發現 w 可以成功的移動，但是 a 卻仍然停留在 CPU 上。其中的差異就在，nn.Parameter 會將 w 宣告成一個 module 內的參數，但是用一般 torch.tensor 方法宣告的卻不會。

那麼假如我們想要在 module 內建立一個不可訓練的參數要怎麼做呢？答案是：使用 register_buffer 功能。

以上面的例子來說，a 可以改用以下方式宣告：

self.__init__(self):
    self.register_buffer('a', torch.tensor(0.1))

之後我們要使用 a 時，一樣用 self.a 就可以使用了。如此一來，當我們將 module 移動到 GPU 時，a 也會跟著一起移動。