更新于 2026年7月25日

在正式介绍BERT模型的下游微调任务之前,先来了解一下PyTorch框架中模型保持与迁移的使用方法,以便后续更好的理解BERT下游任务中预训练模型的加载过程。

3.1 运用场景#

通常,对于模型的保存与加载会出现在以下3个场景中:

①模型推理过程;一个网络模型在完成训练后通常都需要对新样本进行推理预测,此时只需要构建模型的前向传播过程,然后载入已训练好的参数初始化网络即可。

②模型再训练过程;模型在一批数据上训练完成之后需要将其保存到本地,并且可能过了一段时间后又收集到了一批新的数据,因此这个时候就需要将之前的模型载入进行在新数据上进行增量训练(或者是在整个数据上进行全量训练)。

③模型迁移学习;这个时候就是将别人已经训练好的预模型拿过来,作为你自己网络模型参数的一部分进行初始化。例如:你在BERT模型的基础上加了几个全连接层来做分类任务,那么你就需要将原始BERT模型中的参数载入并以此来初始化你的网络中的BERT部分的权重参数。

接下来,我们就以上述3个场景为例来介绍如何利用PyTorch框架来完成上述过程。

3.2 查看网络参数#

所谓查看网络参数指的是在网络模型定义好之后,我们将整个网络中参数信息打印出来的过程。之所以要介绍这部分内容是因为在后续载入BERT与训练模型的时候,我们需要将本地预训练模型的参数名与网络中的参数名一一对应起来才能够完成参数的初始化工作。

由于BERT模型参数较多,为了简洁下面将以之前介绍的LeNet5网络模型为例来分别进行介绍。

3.2.1 查看参数#

首先需要定义好LeNet5的网络模型结构,如下代码所示:

 1 class LeNet5(nn.Module):
 2     def __init__(self, ):
 3         super(LeNet5, self).__init__()
 4         self.conv = nn.Sequential(
 5             nn.Conv2d(1, 6, 5, padding=2),
 6             nn.ReLU(),  # [n,6,24,24]
 7             nn.MaxPool2d(2, 2), 
 8             nn.Conv2d(6, 16, 5), 
 9             nn.ReLU(),
10             nn.MaxPool2d(2, 2)) 
11         self.fc = nn.Sequential(
12             nn.Flatten(),
13             nn.Linear(16  5  5, 120),
14             nn.ReLU(),
15             nn.Linear(120, 84),
16             nn.ReLU(),
17             nn.Linear(84, 10))
18     def forward(self, img):
19         output = self.conv(img)
20         output = self.fc(output)
21         return output

在定义好LeNet5网络结构之后,只要我们完成了这个类的实例化操作,那么网络中对应的权重参数也就完成了相应的初始化工作,即有了一个初始值。同时,我们可以通过如下方式来查看:

1 # Print model's state_dict
2 print("Model's state_dict:")
3 for param_tensor in model.state_dict():
4     print(param_tensor, "\t", model.state_dict()[param_tensor].size())

其输出的结果为:

1 conv.0.weight    torch.Size([6, 1, 5, 5])
2 conv.0.bias      torch.Size([6])
3 conv.3.weight    torch.Size([16, 6, 5, 5])
4 ....
5 ....

可以发现,网络模型中的参数model.state_dict()其实是以字典的形式(实质上是collections模块中的OrderedDict)在进行保存。当然,我们也可以直接输出网络中各个参数的名称:

1 print(model.state_dict().keys())
2 # odict_keys(['conv.0.weight', 'conv.0.bias', 
3 # 'conv.3.weight', 'conv.3.bias', 'fc.1.weight',   
4 # 'fc.1.bias', 'fc.3.weight', 'fc.3.bias', 'fc.5.weight', 'fc.5.bias'])

这样,我们便得到了网络中每个参数的相关信息,包括名称和维度。在后续迁移学习时就可以根据参数名称或形状来匹配相应有用的参数。

3.2.2 自定义参数前缀#

同时,这里值得注意的地方有两点:①参数名中的fc和conv前缀是根据你在上面定义nn.Sequential()时的变量名所确定的;②参数名中的数字表示每个Sequential()中网络层所在的位置。例如将网络结构定义成如下形式:

 1  class LeNet5(nn.Module):
 2     def __init__(self, ):
 3         super(LeNet5, self).__init__()
 4         self.moon = nn.Sequential( 
 5             nn.Conv2d(1, 6, 5, padding=2), 
 6             nn.ReLU(), 
 7             nn.MaxPool2d(2, 2), 
 8             nn.Conv2d(6, 16, 5),
 9             nn.ReLU(),
10             nn.MaxPool2d(2, 2),
11             nn.Flatten(),
12             nn.Linear(16  5  5, 120),
13             nn.ReLU(),
14             nn.Linear(120, 84),
15             nn.ReLU(),
16             nn.Linear(84, 10))

那么其参数名输出结果为:

1 print(model.state_dict().keys())
2 odict_keys(['moon.0.weight', 'moon.0.bias', 'moon.3.weight', 
3     'moon.3.bias', 'moon.7.weight', 'moon.7.bias', 'moon.9.weight', 
4     'moon.9.bias', 'moon.11.weight', 'moon.11.bias'])

理解了这一点对于后续我们去解析和载入一些预训练模型很有帮助。在介绍完模型参数的查看方法后,就可以进入到模型复用阶段的内容介绍了。

3.3 模型推理过程#

模型推理一般涉及到两个过程,一是训练结束后模型的保存;二是推理时模型的加载。下面我们逐一进行介绍。

3.3.1 模型保存#

在PyTorch中,对于模型的保存来说比较简单,通常来说通过如下两行代码便可以实现[11]:

1 model_save_path = os.path.join(model_save_dir, 'model.pt')
2 torch.save(model.state_dict(), model_save_path)

在指定保存的模型名称时PyTorch官方建议的后缀为.pt或者.pth(当然也不是强制,例如.bin也有后缀)。最后,只需要在合适的地方加入第2行代码即可完成模型的保存。同时,如果想要在训练过程中保存某个条件下的最优模型,那么应该通过下面方法一中的形式:

1 # 方法一
2 best_model_state = deepcopy(model.state_dict()) 
3 torch.save(best_model_state, model_save_path)
4 # 方法二
5 best_model_state = model.state_dict() 
6 torch.save(best_model_state, model_save_path)

因为方法二中best_model_state得到只是 model.state_dict() 的引用,它依旧会随着训练过程而发生改变。

3.3.2 模型复用#

在推断过程中,首先需要完成网络的初始化,然后再载入已有的模型参数来覆盖网络中的权重参数即可,示例代码如下:

 1 def inference(data_iter, device, model_save_dir='./MODEL'):
 2     model = LeNet5()  # 
 3     model.to(device)
 4     model_save_path = os.path.join(model_save_dir, 'model.pt')
 5     if os.path.exists(model_save_path):
 6         loaded_paras = torch.load(model_save_path)
 7         model.load_state_dict(loaded_paras)
 8         model.eval() # 注意不要忘记
 9     with torch.no_grad():
10         acc_sum, n = 0.0, 0
11         for x, y in data_iter:
12             x, y = x.to(device), y.to(device)
13             logits = model(x)
14             acc_sum += (logits.argmax(1) == y).float().sum().item()
15             n += len(y)
16         print("Accuracy in test data is :", acc_sum / n)

在上述代码中,第2行是实例化一个LeNet5网络模型,同时初始化现有模型的权重参数;第4~8行是载入本地已有模型,并用其来重新初始化网络中的权重参数。这样,便可以进行后续的推断工作:

1 Accuracy in test data is : 0.8851

3.4 模型再训练过程#

在介绍完模型的保存与复用之后,对于网络的追加训练就很简单了。最简便的一种方式就是在训练过程中只保存网络权重,然后在后续进行追加训练时只载入网络权重参数初始化网络进行训练即可,示例如下(完整代码参见[12]):

 1  def train(self):
 2         #......
 3         model_save_path = os.path.join(self.model_save_dir, 'model.pt')
 4         if os.path.exists(model_save_path):
 5             loaded_paras = torch.load(model_save_path)
 6             self.model.load_state_dict(loaded_paras)
 7             print("#### 成功载入已有模型,进行追加训练...")
 8         optimizer = torch.optim.Adam(self.model.parameters(), lr=self.learning_rate) 
 9        #......
10         for epoch in range(self.epochs):
11             for i, (x, y) in enumerate(train_iter):
12                 x, y = x.to(device), y.to(device)
13                 logits = self.model(x)
14                 # ......
15             print("Epochs[{}/{}]--acc {:.4}".format(epoch,self.epochs,
16                             self.evaluate(test_iter, self.model, device)))
17             torch.save(self.model.state_dict(), model_save_path)

这样,便完成了模型的追加训练。如下为相关输出结果:

1 #### 成功载入已有模型,进行追加训练...
2 Epochs[0/5]---batch[938/0]---acc 0.9062---loss 0.2926
3 Epochs[0/5]---batch[938/100]---acc 0.9375---loss 0.1598
4 ......

除此之外,也可以在保存参数的时候,将优化器参数、损失值等一同保存下来,然后在恢复模型的时候连同其它参数一起恢复,示例如下:

1 model_save_path = os.path.join(model_save_dir, 'model.pt')
2 torch.save({
3             'epoch': epoch,
4             'model_state_dict': model.state_dict(),
5             'optimizer_state_dict': optimizer.state_dict(),
6             'loss': loss,
7             ...
8             }, model_save_path)

恢复模型时的载入方式如下:

1 checkpoint = torch.load(model_save_path)
2 model.load_state_dict(checkpoint['model_state_dict'])
3 optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
4 epoch = checkpoint['epoch']
5 loss = checkpoint['loss']

3.5 模型迁移学习#

到目前为止,对于前面两种应用场景的介绍就算完成了,可以发现总体上并不复杂。但是对于第3种场景的应用来说就会略微复杂一点,而这也是后面我们载入BERT与训练模型的关键。

3.5.1 定义新模型#

假设现在有一个LeNet6网络模型,它是在LeNet5的基础最后多加了一个全连接层,其定义如下:

 1 class LeNet6(nn.Module):
 2     def __init__(self, ):
 3         super(LeNet6, self).__init__()
 4         self.conv = nn.Sequential(
 5             nn.Conv2d(1, 6, 5, padding=2),
 6             nn.ReLU(), 
 7             nn.MaxPool2d(2, 2), 
 8             nn.Conv2d(6, 16, 5),
 9             nn.ReLU(),
10             nn.MaxPool2d(2, 2))
11         self.fc = nn.Sequential(
12             nn.Flatten(),
13             nn.Linear(16  5  5, 120),
14             nn.ReLU(),
15             nn.Linear(120, 84),
16             nn.ReLU(),
17             nn.Linear(84, 64), 
18             nn.ReLU(),
19             nn.Linear(64, 10) )

接下来,我们需要将在LeNet5上训练得到的权重参数迁移到LeNet6网络中去。从上面LeNet6的定义可以发现,此时尽管只是多加了一个全连接层,但是倒数第2层参数的维度也发生了变换。因此,对于LeNet6来说只能复用LeNet5网络前面4层的权重参数。

3.5.2 读取可用参数#

所谓读取模型参数指的是将本地模型载入后输出其中相关参数信息的过程。例如对于开源的BERT中文预训练模型bert_base_chinese[10]来说,我们拿到的将是一个名为pytorch_model.bin的文件。因此首先我们需要读取这个模型并输出相关的参数信息,然后再根据相应的规则解析并将其重新初始化我们自己的网络模型中。

在这里,首先我们可以将LeNet5模型载入,然后再来查看相关参数的信息:

 1 model_save_path = os.path.join('./MODEL', 'model.pt')
 2 loaded_paras = torch.load(model_save_path)
 3 for param_tensor in loaded_paras:
 4     print(param_tensor, "\t", loaded_paras[param_tensor].size())
 5 
 6 #---- 可复用部分
 7 conv.0.weight    torch.Size([6, 1, 5, 5])
 8 conv.0.bias      torch.Size([6])
 9 conv.3.weight    torch.Size([16, 6, 5, 5])
10 conv.3.bias      torch.Size([16])
11 fc.1.weight      torch.Size([120, 400])
12 fc.1.bias    torch.Size([120])
13 fc.3.weight      torch.Size([84, 120])
14 fc.3.bias    torch.Size([84])
15 #----- 不可复用部分
16 fc.5.weight      torch.Size([10, 84])
17 fc.5.bias    torch.Size([10])

同时,对于LeNet6网络的参数信息为:

 1 model = LeNet6()
 2 for param_tensor in model.state_dict():
 3     print(param_tensor, "\t", model.state_dict()[param_tensor].size())
 4 #
 5 conv.0.weight    torch.Size([6, 1, 5, 5])
 6 conv.0.bias      torch.Size([6])
 7 conv.3.weight    torch.Size([16, 6, 5, 5])
 8 conv.3.bias      torch.Size([16])
 9 fc.1.weight      torch.Size([120, 400])
10 fc.1.bias    torch.Size([120])
11 fc.3.weight      torch.Size([84, 120])
12 fc.3.bias    torch.Size([84])
13 #------ 新加入部分
14 fc.5.weight      torch.Size([64, 84])
15 fc.5.bias    torch.Size([64])
16 fc.7.weight      torch.Size([10, 64])
17 fc.7.bias    torch.Size([10])

在理清楚了新旧模型的参数后,下面就可以将LeNet5中我们需要的参数给取出来,然后再换到LeNet6的网络中。

3.5.3 模型迁移学习#

虽然本地载入的模型参数(上面的loaded_paras)和模型初始化后的参数(上面的 model.state_dict() )都是一个字典的形式,但是我们并不能够直接改变 model.state_dict() 中的权重参数。这里需要先构造一个 state_dict 然后通过 model.load_state_dict() 方法来重新初始化网络中的参数。

同时,在这个过程中我们还需要筛选掉本地模型中不可复用的部分,具体代码如下:

 1 def para_state_dict(model, model_save_dir):
 2     state_dict = deepcopy(model.state_dict())
 3     model_save_path = os.path.join(model_save_dir, 'model.pt')
 4     if os.path.exists(model_save_path):
 5         loaded_paras = torch.load(model_save_path)
 6         for key in state_dict:  # 在新的网络模型中遍历对应参数
 7             if key in loaded_paras and state_dict[key].size() == loaded_paras[key].size():
 8                 print("成功初始化参数:", key)
 9                 state_dict[key] = loaded_paras[key]
10     return state_dict

在上述代码中,第2行的作用是先拷贝网络中(LeNet6)原有的参数;第6~9行则是用本地的模型参数(LeNet5)中可以复用的部分替换掉LeNet6中的对应部分,其中第7行就是判断可用的条件。同时需要注意的是在不同的情况下筛选的方式可能不一样,因此具体情况需要具体分析,但是整体逻辑是一样的。

最后,我们只需要在模型训练之前调用该函数,然后重新初始化LeNet6中的部分权重参数即可[12]:

1 state_dict = para_state_dict(self.model, self.model_save_dir)
2 self.model.load_state_dict(state_dict)

此时,训练时的相关信息输出如下:

 1 成功初始化参数: conv.0.weight
 2 成功初始化参数: conv.0.bias
 3 成功初始化参数: conv.3.weight
 4 成功初始化参数: conv.3.bias
 5 成功初始化参数: fc.1.weight
 6 成功初始化参数: fc.1.bias
 7 成功初始化参数: fc.3.weight
 8 成功初始化参数: fc.3.bias
 9 #### 成功载入已有模型,进行追加训练...
10 Epochs[0/5]---batch[938/0]---acc 0.1094---loss 2.512
11 Epochs[0/5]---batch[938/100]---acc 0.9375---loss 0.2141
12 Epochs[0/5]---batch[938/200]---acc 0.9219---loss 0.2729
13 ......
14 Epochs[0/5]---batch[938/900]---acc 0.8906---loss 0.2828
15 Epochs[0/5]--acc on test 0.8808

可以发现,在大约 100 个batch之后,模型的准确率就提升上来了。

您当前阅读的系列内容有完整版高清 PDF ,点击右侧了解

近200页高清 PDF + 80页配套PPT、高清无水印网络结构图,打印学习更方便!

查看详情
阅读 --

第1节 多头注意力机制原理

本文是《Attention Is All You Need》精读系列第1讲,从 RNN 串行瓶颈出发讲清 Transformer 为什么必须引入自注意力,再用图解与公式拆解 Scaled Dot-Product …

第1节 BERT原理与预训练任务

本文是 BERT 精读系列第1讲,从 BERT 与 Transformer Encoder 的关系出发,系统讲清 BERT 论文中那张易混淆的结构图背后的真实网络组成 …

第2节 BERT 从零实现过程

本文是 BERT 精读系列第2讲,配套开源仓库 moon-hotel/BertWithPretrained 系统讲解如何用 PyTorch 从零实现 BERT 模型。覆盖工程目录结构(预训练模型、cache、数据集 …