总说
我们知道,pytorch一般情况下,是将网络中的参数保存成OrderedDict形式的。这里额参数其实包括2种。一种是模型中各种module含的参数,即nn.Parameter
,我们当然可以在网络中定义其他的nn.Parameter
参数。;另外一种是buffer。前者每次optim.step
会得到更新,而不会更新后者。
例子
classmyModel(nn.Module):def__init__(self, kernel_size=3):
super(Depth_guided1, self).__init__()
self.kernel_size = kernel_size
self.back_end = torch.nn.Sequential(
torch.nn.Conv2d(3,32,3, padding=1),
torch.nn.ReLU(True),
torch.nn.Conv2d(3,64,3, padding=1),
torch.nn.ReLU(True),
torch.nn.Conv2d(64,3,3, padding=1),
torch.nn.ReLU(True),
)
mybuffer = np.arange(1,10,1)
self.mybuffer_tmp = np.randn((len(mybuffer),1,1,10), dtype='float32')
self.mybuffer_tmp = torch.from_numpy(self.mybuffer_tmp)# register preset variables as buffer# So that, in testing , we can use buffer variables.
self.register_buffer('mybuffer', self.mybuffer_tmp)# Learnable weights
self.conv_weights = nn.Parameter(torch.FloatTensor(64,10).normal_(mean=0, std=0.01))# Other codedefforward(self):
...# 这里使用 self.mybuffer!
知识点:
I. 定义Parameter
和buffer
都只需要传入Tensor
即可。也不需要将其转成gpu。这是因为,当网络进行.cuda()
时候,会自动将里面的层的参数,buffer等转换成相应的GPU上。
II.self.register_buffer
可以将tensor注册成buffer,在forward中使用self.mybuffer
, 而不是self.mybuffer_tmp
.
III. 网络存储时也会将buffer存下,当网络load模型时,会将存储的模型的buffer也进行赋值。
IV. buffer的更新在forward
中,optim.step
只能更新nn.Parameter
类型的参数。