PyTorch register_buffer|极客笔记

PyTorch register_buffer

在PyTorch中，register_buffer函数主要用于将一个tensor注册为模型的buffer，将其放在模型的参数列表之外，使其不被优化器更新。在训练过程中，我们可能需要使用一些不需要被优化的常量或者运行时不会改变的变量，这时就可以使用register_buffer来管理这些变量。本文将详细介绍register_buffer的用法和示例。

register_buffer的用法

register_buffer函数的用法非常简单，只需要调用模型的register_buffer方法，将要注册为buffer的tensor传递给它即可。注册后，这个tensor将被放在模型的_buffers列表中，不会参与梯度更新。

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.register_buffer('my_buffer', torch.tensor([1, 2, 3]))

    def forward(self, x):
        return x + self.my_buffer

model = Model()
print(model._buffers)

运行结果如下：

OrderedDict([('my_buffer', tensor([1, 2, 3]))])

通过这个示例可以看到，我们成功地将一个tensor注册为了模型的buffer，并且在_buffers列表中找到了它。

register_buffer的实际应用

register_buffer主要用于保存模型中不需要训练的一些常量或者固定参数。下面是一个实际应用的示例：

class Embedding(nn.Module):
    def __init__(self, vocab_size, embed_dim):
        super(Embedding, self).__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim)
        self.register_buffer('position_enc', self._get_positional_encoding(embed_dim))

    def _get_positional_encoding(self, embed_dim):
        pe = torch.zeros(1000, embed_dim)
        position = torch.arange(0, 1000).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, embed_dim, 2) * -(math.log(10000.0) / embed_dim))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        return pe

    def forward(self, input):
        return self.embed(input) + self.position_enc[:input.size(0), :]

model = Embedding(10000, 512)
print(model._buffers)

运行结果如下：

OrderedDict([('position_enc', tensor([[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,
          0.0000e+00,  1.0000e+00],
        [ 8.4147e-01,  5.4030e-01,  8.2186e-01,  ...,  1.0000e+00,
          1.6340e-05,  1.0000e+00],
        [ 9.0929e-01, -4.1615e-01,  9.5105e-01,  ...,  1.0000e+00,
          3.2680e-05,  1.0000e+00],
        ...,
        [-7.9385e-01,  6.1627e-01, -9.4548e-01,  ...,  1.0000e+00,
          9.9766e-01,  1.0000e+00],
        [-2.9420e-01, -9.5631e-01, -2.3924e-01,  ...,  1.0000e+00,
          9.9833e-01,  1.0000e+00],
        [ 6.5699e-01, -7.5395e-01,  7.4876e-01,  ...,  1.0000e+00,
          9.9900e-01,  1.0000e+00]]))])

在这个示例中，我们定义了一个Embedding模型，其中包含了一个self.position_enc的buffer，用来存储位置编码。这个位置编码是一个不需要训练的固定参数，因此非常适合使用register_buffer来管理。