模型权重的读取、设定、初始化与保存

发表于 更新于 分类于

本文摘自《动手学深度学习》的5.2. 参数管理、5.3. 延后初始化和5.5. 读写文件，有删改。

我们首先看一下具有单隐藏层的多层感知机。

1
2
3
4
5
6
import torch
from torch import nn

net = nn.Sequential(nn.Linear(4, 8), nn.ReLU(), nn.Linear(8, 1))
X = torch.rand(size=(2, 4))
net(X)
tensor([[-0.1887],
        [-0.0644]], grad_fn=<AddmmBackward>)

访问参数

我们从已有模型中访问参数。
当通过Sequential类定义模型时，
我们可以通过索引来访问模型的任意层。
这就像模型是一个列表一样，每层的参数都在其属性中。
如下所示，我们可以检查第二个全连接层的参数。

1
print(net[2].state_dict())
OrderedDict([('weight', tensor([[ 0.3231, -0.3373,  0.1639, -0.3125,  0.0527, -0.2957,  0.0192,  0.0039]])), ('bias', tensor([-0.2930]))])

输出的结果告诉我们一些重要的事情：
首先，这个全连接层包含两个参数，分别是该层的权重和偏置。
两者都存储为单精度浮点数（float32）。
注意，参数名称允许唯一标识每个参数，即使在包含数百个层的网络中也是如此。

获取参数

1
2
3
print(type(net[2].bias))
print(net[2].bias)
print(net[2].bias.data)
<class 'torch.nn.parameter.Parameter'>
Parameter containing:
tensor([-0.2930], requires_grad=True)
tensor([-0.2930])

参数是复合的对象，包含值、梯度和额外信息。
这就是我们需要显式参数值的原因。
除了值之外，我们还可以访问每个参数的梯度。
在上面这个网络中，由于我们还没有调用反向传播，所以参数的梯度处于初始状态。

1
net[2].weight.grad == None
True

一次性访问所有参数

当我们需要对所有参数执行操作时，逐个访问它们可能会很麻烦。
当我们处理更复杂的块（例如，嵌套块）时，情况可能会变得特别复杂，
因为我们需要递归整个树来提取每个子块的参数。
下面，我们将通过演示来比较访问第一个全连接层的参数和访问所有层。

1
2
print(*[(name, param.shape) for name, param in net[0].named_parameters()])
print(*[(name, param.shape) for name, param in net.named_parameters()])
('weight', torch.Size([8, 4])) ('bias', torch.Size([8]))
('0.weight', torch.Size([8, 4])) ('0.bias', torch.Size([8])) ('2.weight', torch.Size([1, 8])) ('2.bias', torch.Size([1]))

这为我们提供了另一种访问网络参数的方式，如下所示。

`state_dict`官网主推
1
net.state_dict()['2.bias'].data
tensor([-0.2930])

从嵌套块收集参数

让我们看看，如果我们将多个块相互嵌套，参数命名约定是如何工作的。
我们首先定义一个生成块的函数（可以说是“块工厂”），然后将这些块组合到更大的块中。

1
2
3
4
5
6
7
8
9
10
11
12
13
def block1():
    return nn.Sequential(nn.Linear(4, 8), nn.ReLU(),
                         nn.Linear(8, 4), nn.ReLU())

def block2():
    net = nn.Sequential()
    for i in range(4):
        # 在这里嵌套
        net.add_module(f'block {i}', block1())
    return net

rgnet = nn.Sequential(block2(), nn.Linear(4, 1))
rgnet(X)
tensor([[0.4196],
        [0.4195]], grad_fn=<AddmmBackward0>)

设计了网络后，我们看看它是如何工作的

1
print(rgnet)
Sequential(
  (0): Sequential(
    (block 0): Sequential(
      (0): Linear(in_features=4, out_features=8, bias=True)
      (1): ReLU()
      (2): Linear(in_features=8, out_features=4, bias=True)
      (3): ReLU()
    )
    (block 1): Sequential(
      (0): Linear(in_features=4, out_features=8, bias=True)
      (1): ReLU()
      (2): Linear(in_features=8, out_features=4, bias=True)
      (3): ReLU()
    )
    (block 2): Sequential(
      (0): Linear(in_features=4, out_features=8, bias=True)
      (1): ReLU()
      (2): Linear(in_features=8, out_features=4, bias=True)
      (3): ReLU()
    )
    (block 3): Sequential(
      (0): Linear(in_features=4, out_features=8, bias=True)
      (1): ReLU()
      (2): Linear(in_features=8, out_features=4, bias=True)
      (3): ReLU()
    )
  )
  (1): Linear(in_features=4, out_features=1, bias=True)
)

因为层是分层嵌套的，所以我们也可以像通过嵌套列表索引一样访问它们。
下面，我们访问第一个主要的块中、第二个子块的第一层的偏置项。

1
rgnet[0][1][0].bias.data
tensor([-0.2726,  0.2247, -0.3964,  0.3576, -0.2231,  0.1649, -0.1170, -0.3014])

参数初始化

知道了如何访问参数后，现在我们看看如何正确地初始化参数。
我们在 :numref:sec_numerical_stability中讨论了良好初始化的必要性。
深度学习框架提供默认随机初始化，
也允许我们创建自定义初始化方法，
满足我们通过其他规则实现初始化权重。

默认情况下，PyTorch会根据一个范围均匀地初始化权重和偏置矩阵，
这个范围是根据输入和输出维度计算出的。
PyTorch的nn.init模块提供了多种预置初始化方法。

内置初始化

让我们首先调用内置的初始化器。
下面的代码将所有权重参数初始化为标准差为0.01的高斯随机变量，
且将偏置参数设置为0。

1
2
3
4
5
6
7
8
9
def init_normal(m):
    if type(m) == nn.Linear:
        nn.init.normal_(m.weight, mean=0, std=0.01)
        nn.init.zeros_(m.bias)


net.apply(init_normal)

net[0].weight.data[0], net[0].bias.data[0]
(tensor([ 0.0005, -0.0146, -0.0056,  0.0112]), tensor(0.))

我们还可以将所有参数初始化为给定的常数，比如初始化为1。

1
2
3
4
5
6
7
8
9
def init_constant(m):
    if type(m) == nn.Linear:
        nn.init.constant_(m.weight, 1)
        nn.init.zeros_(m.bias)
        

net.apply(init_constant)

net[0].weight.data[0], net[0].bias.data[0]
(tensor([1., 1., 1., 1.]), tensor(0.))

我们还可以对某些块应用不同的初始化方法。
例如，下面我们使用Xavier初始化方法初始化第一个神经网络层，
然后将第三个神经网络层初始化为常量值42。

1
2
3
4
5
6
7
8
9
10
11
def xavier(m):
    if type(m) == nn.Linear:
        nn.init.xavier_uniform_(m.weight)
def init_42(m):
    if type(m) == nn.Linear:
        nn.init.constant_(m.weight, 42)

net[0].apply(xavier)
net[2].apply(init_42)
print(net[0].weight.data[0])
print(net[2].weight.data)
tensor([-0.4645,  0.0062, -0.5186,  0.3513])
tensor([[42., 42., 42., 42., 42., 42., 42., 42.]])

自定义初始化

有时，深度学习框架没有提供我们需要的初始化方法。
在下面的例子中，我们使用以下的分布为任意权重参数$w$定义初始化方法：

$$
\begin{aligned}
w \sim \begin{cases}
U(5, 10) & \text{ 可能性 } \frac{1}{4} \
0 & \text{ 可能性 } \frac{1}{2} \
U(-10, -5) & \text{ 可能性 } \frac{1}{4}
\end{cases}
\end{aligned}
$$

同样，我们实现了一个my_init函数来应用到net。

1
2
3
4
5
6
7
8
9
10
11
12
13
def my_init(m):
    if type(m) == nn.Linear:
        print(
            "init",
            *[(name, param.shape) for name, param in m.named_parameters()][0]
        )
        nn.init.uniform_(m.weight, -10, 10)
        # 若权重绝对值小于5则设定权重为0
        m.weight.data *= m.weight.data.abs() >= 5
        

net.apply(my_init)
net[0].weight[:2]
init weight torch.Size([8, 4])
init weight torch.Size([1, 8])





tensor([[ 9.9048,  9.5008,  0.0000,  0.0000],
        [-6.2149, -0.0000, -0.0000,  7.3783]], grad_fn=<SliceBackward>)

注意，我们始终可以直接设置参数。

1
2
3
net[0].weight.data[:] += 1
net[0].weight.data[0, 0] = 42
net[0].weight.data[0]
tensor([42.0000, 10.5008,  1.0000,  1.0000])

参数绑定

有时我们希望在多个层间共享参数：我们可以定义一个稠密层，然后使用它的参数来设置另一个层的参数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
shared = nn.Linear(8, 8)

net = nn.Sequential(
    nn.Linear(4, 8), 
    nn.ReLU(),
    shared, 
    nn.ReLU(),
    shared, 
    nn.ReLU(),
    nn.Linear(8, 1)
)

net(X)

# 检查参数是否相同
print(net[2].weight.data[0] == net[4].weight.data[0])

# 主动为net[2]赋值，net[4]也会变
net[2].weight.data[0, 0] = 100
print(net[2].weight.data[0] == net[4].weight.data[0])
tensor([True, True, True, True, True, True, True, True])
tensor([True, True, True, True, True, True, True, True])

这个例子表明第三个和第五个神经网络层的参数是绑定的。
它们不仅值相等，而且由相同的张量表示。
因此，如果我们改变其中一个参数，另一个参数也会改变。
你可能会思考：当参数绑定时，梯度会发生什么情况？
答案是由于模型参数包含梯度，因此在反向传播期间第二个隐藏层
（即第三个神经网络层）和第三个隐藏层（即第五个神经网络层）的梯度会加在一起。

模型持久化

加载和保存张量

对于单个张量，我们可以直接调用load和save函数分别读写它们。这两个函数都要求我们提供一个名称，save要求将要保存的变量作为输入。

1
2
3
4
5
6
7
8
9
10
11
import torch
from torch import nn
from torch.nn import functional as F


x = torch.arange(4)

# 将`x`保存在当前路径名为`x-file`的文件中
torch.save(x, 'x-file')

x
tensor([0, 1, 2, 3])

我们现在可以将存储在文件中的数据读回内存。

1
2
3
x2 = torch.load('x-file')

x2
tensor([0, 1, 2, 3])

我们可以存储一个张量列表，然后把它们读回内存。

1
2
3
4
5
6
7
y = torch.zeros(4)

torch.save([x, y],'x-files')

x2, y2 = torch.load('x-files')

(x2, y2)
(tensor([0, 1, 2, 3]), tensor([0., 0., 0., 0.]))

我们甚至可以写入或读取从字符串映射到张量的字典。

1
2
3
4
5
6
7
mydict = {'x': x, 'y': y}

torch.save(mydict, 'mydict')

mydict2 = torch.load('mydict')

mydict2
{'x': tensor([0, 1, 2, 3]), 'y': tensor([0., 0., 0., 0.])}

加载和保存模型参数

深度学习框架提供了内置函数来保存和加载整个网络。需要注意的一个重要细节是，这将保存模型的参数而不是保存整个模型。例如，如果我们有一个3层多层感知机，我们需要单独指定架构。因为模型本身可以包含任意代码，所以模型本身难以序列化。因此，为了恢复模型，我们需要用代码生成架构，然后从磁盘加载参数。让我们从熟悉的多层感知机开始尝试一下。

1
2
3
4
5
6
7
8
9
10
11
12
13
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden = nn.Linear(20, 256)
        self.output = nn.Linear(256, 10)

    def forward(self, x):
        return self.output(F.relu(self.hidden(x)))

    
net = MLP()
X = torch.randn(size=(2, 20))
Y = net(X)

接下来，我们将模型的参数存储在一个叫做“mlp.params”的文件中。

1
torch.save(net.state_dict(), 'mlp.params')

为了恢复模型，我们实例化了原始多层感知机模型的一个备份。 这里我们不需要随机初始化模型参数，而是直接读取文件中存储的参数。

1
2
3
clone = MLP()
clone.load_state_dict(torch.load('mlp.params'))
clone.eval()
MLP(
  (hidden): Linear(in_features=20, out_features=256, bias=True)
  (output): Linear(in_features=256, out_features=10, bias=True)
)

由于两个实例具有相同的模型参数，在输入相同的X时，
两个实例的计算结果应该相同。
让我们来验证一下。

1
2
Y_clone = clone(X)
Y_clone == Y
tensor([[True, True, True, True, True, True, True, True, True, True],
        [True, True, True, True, True, True, True, True, True, True]])

小结

  • 我们有几种方法可以访问、初始化和绑定模型参数。
  • 我们可以使用自定义初始化方法。

  • 保存与读取单个张量：

    1
    2
    torch.save(x, 'x-file')
    x2 = torch.load('x-file')

  • 保存模型

    1
    torch.save(net.state_dict(), 'mlp.params')

  • 读取模型

    1
    2
    3
    4
    5
    6
    7
    # 定义模型结构
    class MLP(nn.Module):
        ...

    clone = MLP()
    clone.load_state_dict(torch.load('mlp.params'))
    clone.eval()