Pytorch中的register_buffer()Pytorch中的register_buffer1.register_buffer( )的使用随着例子边看边讲例子1使用类成员变量类成员变量并不会在我们的model.state_dict()即无法保存例子2使用类成员变量类成员变量并不会随着model.cuda()复制到gpu上例子3使用register_buffer()总结2.Parameter与Buffer2.1 model.buffers()和model.named_buffers对模型中的buffer进行访问2.2 Buffer变量可以通过backward()得到梯度信息2.3 Buffer变量不需要求梯度时可通过Parameter代替3.BN中的参数1.register_buffer( )的使用回顾模型保存torch.save(model.state_dict())model.state_dict()是一个字典里边存着我们模型各个部分的参数。在model中我们需要更新其中的参数训练结束将参数保存下来。但在某些时候我们可能希望模型中的某些参数不更新从开始到结束均保持不变但又希望参数保存下来model.state_dict()这是我们就会用到register_buffer()。随着例子边看边讲例子1使用类成员变量类成员变量并不会在我们的model.state_dict()即无法保存。成员变量(self.tensor)在前向传播中用到希望它也能保存下来但他不在我们的state_dict中。classmy_model(nn.Module):def__init__(self):super(my_model,self).__init__()self.convnn.Conv2d(1,1,3,1,1)self.tensortorch.randn(size(1,1,5,5))# 成员变量defforward(self,x):returnself.conv(x)self.tensor xtorch.randn(size(1,1,5,5))例子2使用类成员变量类成员变量并不会随着model.cuda()复制到gpu上将上一个例子中的模型复制到GPU上但成员变量并不会随着model.cuda()复制到gpu上。torch中如果有数据不在同一个“地方”进行运算程序会报错即self.tensor在 cpu 上模型和x在 “cuda:0” 上。classmy_model(nn.Module):def__init__(self):super(my_model,self).__init__()self.convnn.Conv2d(1,1,3,1,1)self.tensortorch.randn(size(1,1,5,5))# 成员变量defforward(self,x):returnself.conv(x)self.tensor xtorch.randn(size(1,1,5,5))例子3使用register_buffer()self.register_buffer(my_buffer, self.tensor)my_buffer是名字str类型self.tensor是需要进行register登记的张量。这样我们就得到了一个新的张量这个张量会保存在model.state_dict()中也就可以随着模型一起通过.cuda()复制到gpu上。classmy_model(nn.Module):def__init__(self):super(my_model,self).__init__()self.convnn.Conv2d(1,1,3,1,1)self.tensortorch.randn(size(1,1,5,5))self.register_buffer(my_buffer,self.tensor)defforward(self,x):returnself.conv(x)self.my_buffer# 这里不再是self.tensor总结成员变量不更新但是不算是模型中的参数model.state_dict()通过register_buffer()登记过的张量会自动成为模型中的参数随着模型移动gpu/cpu而移动但是不会随着梯度进行更新。2.Parameter与Buffer模型保存下来的参数有两种一种是需要更新的Parameter另一种是不需要更新的buffer。在模型中利用backward反向传播可以通过requires_grad来得到buffer和parameter的梯度信息但是利用optimizer进行更新的是parameterbuffer不会更新这也是两者最重要的区别。这两种参数都存在于model.state_dict()的OrderedDict中也会随着模型“移动”model.cuda()。2.1 model.buffers()和model.named_buffers对模型中的buffer进行访问与model.parameters()和model.named_parameters()相同只是一个是对模型中的parameter访问一个是对模型中的buffer访问。classmy_model(nn.Module):def__init__(self):super(my_model,self).__init__()self.convnn.Conv2d(1,1,3,1,1)self.tensortorch.randn(size(1,1,5,5))self.tensor2torch.randn(size(1,1))self.register_buffer(my_buffer,self.tensor)self.register_buffer(my_buffer2,self.tensor2)defforward(self,x):# ... (展开)2.2 Buffer变量可以通过backward()得到梯度信息buffer变量和parameter变量一样都可以通过backward()得到梯度信息但区别是优化器optimizer更新的parameter变量所以buffer并不会更新。classmy_model(nn.Module):def__init__(self):super(my_model,self).__init__()self.tensortorch.randn(size(2,2))self.register_buffer(my_buffer,self.tensor)defforward(self,x):returnself.my_buffer*x2.3 Buffer变量不需要求梯度时可通过Parameter代替在构造模型时候可以将某些Parameter从模型中通过.detach()方法或者直接将Parameter的requires_grad设置为False使得此变量不求梯度也可达到不更新的效果。通过nn.Parameter()将张量设置为变量同时设置requires_grad为False这个变量也会随着模型保存并且随着模型“移动”可达到与buffer相同的效果为什么要存在bufferbuffer与parameter具有“同等地位”所以将某些不需要更新的变量“拿出来”作为buffer可能更方便操作可读性也更高对Paramter的各种操作固定网络的等可能也不会“误伤到” buffer这种变量。buffer最重要的意义应该是需要得到梯度信息时不会更新因为optimizer而更新这也是parameter所不能代替的。classmy_model(nn.Module):def__init__(self):super(my_model,self).__init__()self.convnn.Conv2d(1,1,3,1,1)self.tensornn.Parameter(torch.randn(size(1,1,5,5)),requires_gradFalse)defforward(self,x):returnself.conv(x)self.tensor xtorch.randn(size(1,1,5,5))3.BN中的参数最近发现bn中的running_mean, running_var, num_batches_tracked这三个参数是buffer类型的这样既可以用state_dict()保存也不会随着optimizer更新。此外我们要注意state_dict()只会保存parameters和buffers类型的变量如果我们有变量没有转成这两种类型最后是不会被保存的classnetwork(nn.Module):def__init__(self):super(network,self).__init__()self.convnn.Conv2d(1,1,1,padding0)self.bnnn.BatchNorm2d(2)defforward(self,x):returnself.bn(self.conv(x))netnetwork()