模型保存下来的参数有两种:需要更新的 Parameter ,与是不需要更新的 buffer 。
在模型中,利用backward反向传播,可以通讨requires_grad来得到 buffer 和 parameter 的梯度信息,但是利用 optimizer 进行更新的是 parameter , buffer 不会更新,这也是两者最重要的区别。这两种参数都存在于 model.state_dict()的OrderedDict中,也会随着模型”移动”(model.cuda())。
class my_model(nn.Module):def _init_ (self):super(my_model,self)._init_()self.conv = nn.conv2d(11,3,1,1)self.tensor =torch.randn(size=(1,1,5,5))self.register_buffer('my_buffer', self.tensor)
def forward(self,x):return self.conv(x)+self.my_buffer # 这里不再是self.tensor
之后要输出就是model.my_buffer了。
视频链接:28-masked_fill张量掩码|Dropout正则化|view|register_buffer|contiguous|多头注意力-transformer_哔哩哔哩_bilibili