register_buffer && nn.Parameter

torch.nn 评论

本文的大部分例子来源于知乎Link

在pytorch模型中保存模型参数的方式如下

1
torch.save(model.state_dict(), path)

模型保存的是model.state_dict()返回对象，是一个OrderDict，他的key与value分别是模型需要保存的参数名字和值。下面介绍parameter和buffer的一些用法特点

Paramter Buffer
是否被更新 是 否
返回 model.parameters() model.buffers()
是否注册到模型中 是 是
是否随模型保存 是 是

其中parameter可以被optimizer更新，我们在优化模型参数的时候，一般都会写SGD(model.parameters(), xxx)，此外parameter与buffer均可以在保存模型参数时被保存到OrderDict中。下面分别介绍中这两个参数类型的区别以及如何构建

register_buffer

register_buffer(name, tensor, persistent=True)

  • persistent: 是否将这个参数作为module的state_dict

buffer的创建需要构建一个tensor，然后将这个tensor注册到buffer中，如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.l1 = nn.Linear(2, 2)
        buffer = torch.randn(2, 3)  # tensor
        self.register_buffer('my_buffer', buffer)

    def forward(self, x):
        pass

model = MyModel()
for param in model.parameters():
    print(param)
for buffer in model.buffers():
    print(buffer)
print(model.state_dict())

'''
# model.parameters()
Parameter containing:
tensor([[ 0.0184, -0.3397],
        [ 0.1823, -0.2097]], requires_grad=True)
Parameter containing:
tensor([0.5309, 0.4586], requires_grad=True)

# model.buffers()
tensor([[-0.0885,  0.2578, -0.1473],
        [-0.1926,  0.2726, -0.5541]])
      
# model.state_dict()
OrderedDict([('my_buffer', tensor([[-0.0885,  0.2578, -0.1473],
        [-0.1926,  0.2726, -0.5541]])), ('l1.weight', tensor([[ 0.0184, -0.3397],
        [ 0.1823, -0.2097]])), ('l1.bias', tensor([0.5309, 0.4586]))])
'''

模型中一共有两种类型的参数

  1. 一个是linear操作，其中linear的weight和bias会随着model.parameters输出，并且参数可以被optimizer优化
  2. 一个是buffer，buffer类型的参数会随着model.buffers()输出，不能被optimizer优化
  3. 在模型保存时，model.state_dict可以将两种参数都保存

Parameter

nn.Parameter(data=None, requires_grad=True)

parameter类型的变量也会自动注册到模型中，具有梯度，可以被optimizer进行优化。可以将其理解为和linear等参数相同的参数类型，如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.l1 = nn.Linear(2, 2)
        self.param = nn.Parameter(torch.randn(3, 3))  # 模型的成员变量

    def forward(self, x):
        # 可以通过 self.param 和 self.my_buffer 访问
        pass

model = MyModel()
for param in model.parameters():
    print(param)
print("----------------")
print(model.state_dict())

'''
# model.parameters()
Parameter containing:
tensor([[-0.4412, -2.0199,  0.7088],
        [ 0.6840,  1.0006,  0.1266],
        [ 0.9492, -0.0404, -0.6280]], requires_grad=True)
Parameter containing:
tensor([[ 0.6309, -0.1017],
        [ 0.1819,  0.3834]], requires_grad=True)
Parameter containing:
tensor([0.5768, 0.1148], requires_grad=True)

# model.state_dict()
OrderedDict([('param', tensor([[-0.4412, -2.0199,  0.7088],
        [ 0.6840,  1.0006,  0.1266],
        [ 0.9492, -0.0404, -0.6280]])), ('l1.weight', tensor([[ 0.6309, -0.1017],
        [ 0.1819,  0.3834]])), ('l1.bias', tensor([0.5768, 0.1148]))])
'''

这里模型只有一种类型的参数了，即model.parameter，没有model.buffer类型

  1. 一个是linear操作，其中linear的weight和bias会随着model.parameters()输出，并且参数可以被optimizer优化
  2. 一个是parameter类型，参数会随着model.parameters()输出，参数可以被optimizer优化
  3. 在模型保存时，model.state_dict会保存上述参数

一些疑问

  1. 为什么不将参数都设为nn.Parameter，只是把不需要修改的参数设置为requires_grad=False?

如果不想将参数进行optimizer的更新，设置为buffer类型的话会给人更直观的感觉，表达更清晰。当然如果设置为nn.Parameter并且grad设为false也可以

  1. 为什么不直接将不需要进行更改的参数变量设为普通tensor变量？

下面通过一个例子来说明，为什么必须注册为parameters或者buffer

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.my_tensor = torch.randn(1)  # 参数直接作为模型类成员变量
        self.register_buffer('my_buffer', torch.randn(1))  # 参数注册为 buffer
        self.my_param = nn.Parameter(torch.randn(1))

    def forward(self, x):
            return x

model = MyModel()
print(model.state_dict())
model.cuda()
print(model.my_tensor)
print(model.my_buffer)

'''
OrderedDict([('my_param', tensor([-1.0101])), ('my_buffer', tensor([-0.5266]))])
tensor([-0.2454])
tensor([-0.5266], device='cuda:0')
'''

如上，如果在模型中仅设置一个普通tensor的话，他并不会成为模型的一部分，也不会随模型移动到cuda中

harryteaustc

computer vision