NTM中register_buffer的用法

用法解释

register_buffer()函数是PyTorch中nn.Module类提供的一个方法，用于注册模块的可学习缓存，即不参与模型训练但是在模型的前向计算中被使用的缓存。

具体来说，register_buffer(name, tensor)函数有两个参数：name和tensor。其中，name代表缓存的名称，tensor代表需要缓存的torch.Tensor类型的数据。调用register_buffer()函数后，指定的tensor会被自动添加到模块的可学习参数中，并且在模型进行前向计算的时候被缓存起来以加速计算过程。

需要注意的是，缓存的tensor是不参与反向传播过程的，即它们不会被梯度下降更新。因此，register_buffer()函数主要用于保存一些常量或者模型运行中固定的中间结果，这种中间结果在模型的训练和测试过程中都可以被使用到。

例子解释

import torch
class MyNN(torch.nn.Module):
    def __init__(self):
        super(MyNN, self).__init__()

        # 申请一个size为(3,)的register_buffer()内存空间
        self.register_buffer('my_buffer', torch.zeros(3))

    def forward(self, x):
        # 在my_buffer存储一些数据
        self.my_buffer[0] = 1
        self.my_buffer[1] = 2
        self.my_buffer[2] = 3

        # 对x和my_buffer进行加法运算
        out = x + self.my_buffer

        return out

# 创建MyNN类的实例
net = MyNN()
# 定义输入x
x = torch.ones(3)
# 进行前向计算
out = net(x)
print(out)

在这个例子中会输出tensor([2., 3., 4.])，说明out的三个元素分别是x和my_buffer对应元素相加的结果。其中，my_buffer的值可以在类中的其他方法中被修改或直接调用，因为它是在类的构造函数中通过register_buffer()函数申请的内存空间。