（四）模型训练保存与加载

当保存和加载模型时，需要熟悉三个核心功能：

1. torch.save ：将序列化对象保存到磁盘。此函数使用 Python 的 pickle 模块进行序列化。使
   用此函数可以保存如模型、 tensor 、字典等各种对象。
2. torch. load ：使用 pickle 的 unpickling 功能将 pickle 对象文件反序列化到内存。此功能还可
   以有助于设备加载数据。
3. torch.nn.Module.load_state_dict ：使用反序列化函数 state_dict 来加载模型的参数字典。

state_dict 是一个Python字典，将每一层映射成它的参数张量。注意只有带有可学习参数的层（卷积层、全连接层等），以及注册的缓存（batchnorm的运行平均值）在state_dict 中才有记录。state_dict同样包含优化器对象，存储了优化器的状态，所使用到的超参数。

保存/加载 整个模型

保存：

torch.save(model, PATH)
1

加载：

# 模型类必须在别的地方定义
model = torch.load(PATH)
model.eval()
1
2
3

这种保存/加载模型的过程使用了最直观的语法，所用代码量少。这使用Python的pickle保存所有模块。这种方法的缺点是，保存模型的时候，序列化的数据被绑定到了特定的类和确切的目录。这是因为pickle不保存模型类本身，而是保存这个类的路径，并且在加载的时候会使用。因此，当在其他项目里使用或者重构的时候，加载模型的时候会出错。

一般来说，PyTorch的模型以.pt或者.pth文件格式保存。

一定要记住在评估模式的时候调用model.eval()来固定dropout和批次归一化。否则会产生不一致的推理结果。

保存/加载 模型参数 state_dict（推荐）

保存：

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

加载：

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()
1
2
3

要注意这个细节，如果使用nn.DataParallel在一台电脑上使用了多个GPU，那么加载模型的时候也必须先进行nn.DataParallel。

保存模型的推理过程的时候，只需要保存模型训练好的参数，使用torch.save()保存state_dict，能够方便模型的加载。因此推荐使用这种方式进行模型保存。

记住一定要使用model.eval()来固定dropout和归一化层，否则每次推理会生成不同的结果。

注意，load_state_dict()需要传入字典对象，因此需要先反序列化state_dict再传入load_state_dict()

断点续训练Checkpoint

训练过程中周期性保存模型：

torch.save({
			'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'epoch': epoch,
            'loss': loss,
            }, PATH)
1
2
3
4
5
6

从断点处加载：

model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()
# - 或者 -
model.train()
1
2
3
4
5
6
7
8
9
10
11
12

在保存用于推理或者继续训练的常规检查点的时候，除了模型的state_dict之外，还必须保存其他参数。保存优化器的state_dict也非常重要，因为它包含了模型在训练时候优化器的缓存和参数。除此之外，还可以保存停止训练时epoch数，最新的模型损失，额外的torch.nn.Embedding层等。

要保存多个组件，则将它们放到一个字典中，然后使用torch.save()序列化这个字典。一般来说，使用.tar文件格式来保存这些检查点。

加载各个组件，首先初始化模型和优化器，然后使用torch.load()加载保存的字典，然后可以直接查询字典中的值来获取保存的组件。

同样，评估模型的时候一定不要忘了调用model.eval()。

模型微调finetune

给定预训练模型（Pre_trained model）,基于模型进行微调（Fine Tune）。相对于从头开始训练(Training a model from scatch)，微调为你省去大量计算资源和计算时间，提高了计算效率,甚至提高准确率。

预训练模型就是已经用数据集训练好了的模型，现在我们常用的预训练模型就是他人用常用模型，比如VGG16/19，Resnet等模型，并用大型数据集来做训练集，比如Imagenet, COCO等训练好的模型参数。

为什么要微调？
预训练模型的特点是： 用了大型数据集做训练，已经具备了提取浅层基础特征和深层抽象特征的能力。如果不做微调： （1）从头开始训练，需要大量的数据，计算时间和计算资源。 （2）存在模型不收敛，参数不够优化，准确率低，模型泛化能力低，容易过拟合等风险。

如何进行模型微调（替换fc层 + 冻结前几层 / 参数组设置小学习率）
1.通常的做法是截断预先训练好的网络的最后一层（softmax层），并用与我们自己的问题相关的新的softmax层替换它。例如，ImageNet上预先训练好的网络带有1000个类别的softmax图层。如果我们的任务是对10个类别的分类，则网络的新softmax层将由10个类别组成，而不是1000个类别。然后，我们在网络上运行预先训练的权重。确保执行交叉验证，以便网络能够很好地推广。

2.使用较小的学习率(params_group浅层小lr, 深层大lr)来训练网络。由于我们预计预先训练的权重相对于随机初始化的权重已经相当不错，我们不想过快地扭曲它们太多。通常的做法是使初始学习率比用于从头开始训练（Training from scratch）的初始学习率小10倍。

3.如果数据集数量过少，我们进来只训练最后一层，如果数据集数量中等，冻结预训练网络的前几层的权重(requires_grad=False;或lr=0;)也是一种常见做法。这是因为前几个图层捕捉了与我们的新问题相关的通用特征，如曲线和边。我们希望保持这些权重不变。相反，我们会让网络专注于学习后续深层中特定于数据集的特征。

# ============================ step 2/5 模型 ============================

# 1/3 构建模型
resnet18_ft = models.resnet18()

# 2/3 加载参数
# flag = 0
flag = 1
if flag:
    path_pretrained_model = os.path.join(BASEDIR, "finetune_data", "resnet18-5c106cde.pth")#指明预训练模型参数地址
    state_dict_load = torch.load(path_pretrained_model)#加载模型参数
    resnet18_ft.load_state_dict(state_dict_load)#将模型参数加载到当前模型中

# 法1 : 冻结卷积层（替换fc层 + requires_grad=False;）
flag_m1 = 0
# flag_m1 = 1
if flag_m1:
    for param in resnet18_ft.parameters():
        param.requires_grad = False
    print("conv1.weights[0, 0, ...]:\n {}".format(resnet18_ft.conv1.weight[0, 0, ...]))


# 3/3 替换fc层 -更改为2分类
num_ftrs = resnet18_ft.fc.in_features

from torchsummary import summary
# print(summary(resnet18_ft, input_size=(3, 112, 112), device="cpu"))

# print(num_ftrs)   512
resnet18_ft.fc = nn.Linear(num_ftrs, classes)#替换层

# print(summary(resnet18_ft, input_size=(3, 112, 112), device="cpu"))


resnet18_ft.to(device)
# ============================ step 3/5 损失函数 ============================
criterion = nn.CrossEntropyLoss()                                                   # 选择损失函数

# ============================ step 4/5 优化器 ============================
# 法2 : conv 小学习率（替换fc层 + 参数组设置不同lr）
flag = 0
# flag = 1
if flag:
    fc_params_id = list(map(id, resnet18_ft.fc.parameters()))     # 返回的是parameters的 内存地址,得到最好fc层的参数地址
    base_params = filter(lambda p: id(p) not in fc_params_id, resnet18_ft.parameters())  # 得到除了fc层的前面层的参数地址
    optimizer = optim.SGD([
        {'params': base_params, 'lr': LR*0.1},   # 0或LR*0.1： 浅层小学习率
        {'params': resnet18_ft.fc.parameters(), 'lr': LR}], momentum=0.9)  # 深层大学习率

else:
    optimizer = optim.SGD(resnet18_ft.parameters(), lr=LR, momentum=0.9)               # 选择优化器

scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=lr_decay_step, gamma=0.1)     # 设置学习率下降策略
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
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53

GPU的使用

查看gpu型号与数量

from torch import cuda
print(cuda.get_device_name())
print(cuda.device_count())
1
2
3

Tensor和Module的对象可以使用.to('cpu')或.to('cuda')转移在cpu和gpu之间
tensor不执行inpalce，所以to()方法的返回值必须再赋给原对象。tensor=tensor.to("cuda")
module可以执行inplace，不需要再覆盖原对象。module.to("cuda")
.cuda()方法已经弃用

通常使用这个设置device(自动判断)：

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
x = x.to(device)
1
2

多gpu并行分发加速

TensorBoard

用于数据可视化的工具

SummaryWriter()

这个函数用于创建一个tensorboard文件，其中常用参数有：
log_dir：tensorboard文件的存放路径
flush_secs：表示写入tensorboard文件的时间间隔
调用方式如下：

writer = SummaryWriter(log_dir='logs',flush_secs=60)
1

writer.add_graph()

这个函数用于在tensorboard中创建Graphs，Graphs中存放了网络结构，其中常用参数有：

model：pytorch模型
input_to_model：pytorch模型的输入