Pytorch学习笔记（9）模型的保存与加载、模型微调、GPU使用

 

目录

 一、模型的保存与加载

1.1 序列化与反序列化

1.2 保存加载模型基本用法

1.2.1 保存模型

1.2.2 加载模型

1.3 模型的断点续训练

 二、模型微调 

2.1 Transfer Learning ＆ Model Finetune

2.2 Finetune的实例

三、GPU的使用

3.1  CPU VS GPU

3.2 数据迁移至GPU

3.3 多GPU并行运算

---

 前期回顾：

Pytorch学习笔记（1）：基本概念、安装、张量操作、逻辑回归

Pytorch学习笔记（2）：数据读取机制（DataLoader与Dataset）

Pytorch学习笔记（3）：图像的预处理（transforms）

Pytorch学习笔记（4）：模型创建（Module）、模型容器（Containers）、AlexNet构建

Pytorch学习笔记（5）：torch.nn---网络层介绍（卷积层、池化层、线性层、激活函数层）

Pytorch学习笔记（6）：模型的权值初始化与损失函数_pytorch 模型权重初始化

Pytorch学习笔记（7）：优化器、学习率及调整策略、动量

Pytorch学习笔记（8）：正则化（L1、L2、Dropout）与归一化（BN、LN、IN、GN）

---

  一、模型的保存与加载

神经网络训练后我们需要将模型进行保存，要用的时候将保存的模型进行加载，那么如何保存和加载模型呢？下面就从三个方面来理解一下。

1.1 序列化与反序列化

序列化与反序列化主要讲的是硬盘与内存之间的数据转换关系

模型在内存中是以一个对象的形式被存储的，且不具备长久存储的功能；而在硬盘中，模型是以二进制数序列进行存储。

​

• 序列化是指将内存中的某一个对象保存在硬盘当中，以二进制序列的方式存储下来
• 反序列化是指将存储的这些二进制数反序列化的放到内存当中作为一个对象 

序列化和反序列化的目的就是将我们的模型长久的保存。

pytorch中序列化和反序列化的方法： 

（1）torch.save

功能：序列化

主要参数：

• obj：对象，想要保存的数据，模型、张量等
• f：输出路径

（2）torch.load

功能：反序列化

主要参数：

• f：文件路径
• map_location：指定存放位置，cpu or gpu

---

1.2 保存加载模型基本用法

1.2.1 保存模型

 Pytorch的模型保存有两种方法：

（1）保存整个模型

保存整个网络模型（网络结构+权重参数）

torch.save(model, 'net.pkl')

直接加载整个网络模型（可能比较耗时）

model = torch.load('net.pkl')

这种方法保存的是整个模型架构，比较费时占用内存，所以官方比较推荐的是第二种方法。

（2）只保存模型参数

只保存模型的权重参数（速度快，占内存少）

torch.save(model.state_dict(), 'net_params.pkl')

因为我们只保存了模型的参数，所以需要先定义一个网络对象，然后再加载模型参数

# 构建一个网络结构
model = ClassNet()
# 将模型参数加载到新模型中
state_dict = torch.load('net_params.pkl')
model.load_state_dict(state_dict)

这个第二种方法只保存了模型上可学习的参数，等建立一个新的相同的网络结构时将这些参数加载到网络中即可。所以推荐使用第二种。

---

1.2.2 加载模型

上面保存加载的 net.pkl 其实一个字典，通常包含如下内容：

1. 网络结构：输入尺寸、输出尺寸以及隐藏层信息，以便能够在加载时重建模型。
2. 模型的权重参数：包含各网络层训练后的可学习参数，可以在模型实例上调用 state_dict() 方法来获取，比如前面介绍只保存模型权重参数时用到的 model.state_dict()。
3. 优化器参数：有时保存模型的参数需要稍后接着训练，那么就必须保存优化器的状态和所其使用的超参数，也是在优化器实例上调用 state_dict() 方法来获取这些参数。
4. 其他信息：有时我们需要保存一些其他的信息，比如 epoch，batch_size 等超参数。

（1）加载整个网络

net_load = torch.load（model_path）

（2）加载模型的参数，并应用新模型

state_dict_load = torch.load(model_path)

net_new.load_state_dict(state_dict_load)

---

1.3 模型的断点续训练

在训练过程中，可能由于某种意外原因如断点等导致训练终止，这时需要重新开始训练。断点续练是在训练过程中每隔一定次数的 epoch 就保存模型的参数和优化器的参数，这样如果意外（比如停电啦）终止训练了，下次就可以重新加载最新的模型参数和优化器的参数，在这个基础上继续训练。

需要保存的数据：checkpoint

  

由上图我们可以看到模型训练的五个步骤： 数据 -> 模型 -> 损失函数 -> 优化器 -> 迭代训练这五个步骤中数据和损失函数是没法改变的，而在迭代训练的过程中模型的一些可学习参数和优化器中的一些缓存是会变的，所以需要保留这些信息，另外还需要保留迭代的次数，如下：

 应用：模拟意外中断，然后续训练

下面的代码中，每隔 5 个 epoch 就保存一次，保存的是一个 dict，包括模型参数、优化器的参数、epoch。然后在 epoch 大于 5 时，就break模拟训练意外终止。关键代码如下：

 if (epoch+1) % checkpoint_interval == 0:

        checkpoint = {"model_state_dict": net.state_dict(),
                      "optimizer_state_dict": optimizer.state_dict(),
                      "epoch": epoch}
        path_checkpoint = "./checkpoint_{}_epoch.pkl".format(epoch)
        torch.save(checkpoint, path_checkpoint)

在 epoch 大于 5 时，就break模拟训练意外终止

    if epoch > 5:
        print("训练意外中断...")
        break

 断点续训练的恢复代码如下：

path_checkpoint = "./checkpoint_4_epoch.pkl"
checkpoint = torch.load(path_checkpoint)

net.load_state_dict(checkpoint['model_state_dict'])

optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

start_epoch = checkpoint['epoch']

scheduler.last_epoch = start_epoch

需要注意的是，还要设置scheduler.last_epoch参数为保存的 epoch。模型训练的起始 epoch 也要修改为保存的 epoch。 

---

 二、模型微调 

2.1 Transfer Learning ＆ Model Finetune

Transfer Learning：机器学习分支，研究源域（source domain）的知识如何应用到目标域（target domain）

将在其他任务中学习到的知识，应用到新任务中。

Model Finetune：模型微调

应用时一般根据任务修改最后一个全连接层的输出神经元个数，分几类就几个神经元

模型微调的步骤：

1、获取预训练模型参数（源任务当中学习到的知识）；
2、加载模型（load_state_dict）将学习到的知识放到新的模型；
3、修改输出层， 以适应新的任务；

模型微调的训练方法：

1、固定预训练的参数 (requires_grad=False; lr=0)
2、Features Extractor 较小学习率 (params_group)

---

2.2 Finetune的实例

下面就通过一个例子，看看如何使用模型的finetune：

我们来使用训练好的ResNet-18进行蚂蚁和蜜蜂二分类： 

数据集：链接  ResNet-18: 链接

该数据集训练数据有120张，验证数据有70张，训练数据太少，所以用模型重新训练可能达不到想要的效果，这里用迁移学习，用已经训练好的ResNet-18来完成这个二分类任务。

import os
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import torch.optim as optim
from matplotlib import pyplot as plt
 
import sys
hello_pytorch_DIR = os.path.abspath(os.path.dirname(__file__)+os.path.sep+".."+os.path.sep+"..")
sys.path.append(hello_pytorch_DIR)
 
from tools.my_dataset import AntsDataset
from tools.common_tools import set_seed
import torchvision.models as models
import torchvision
BASEDIR = os.path.dirname(os.path.abspath(__file__))
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("use device :{}".format(device))
 
set_seed(1)  # 设置随机种子
label_name = {"ants": 0, "bees": 1}
 
# 参数设置
MAX_EPOCH = 25
BATCH_SIZE = 16
LR = 0.001
log_interval = 10
val_interval = 1
classes = 2
start_epoch = -1
lr_decay_step = 7
 
 
# ============================ step 1/5 数据 ============================
data_dir = os.path.abspath(os.path.join(BASEDIR, "..", "..", "data", "07-02-数据-模型finetune"))
if not os.path.exists(data_dir):
    raise Exception("\n{} 不存在，请下载 07-02-数据-模型finetune.zip  放到\n{} 下，并解压即可".format(
        data_dir, os.path.dirname(data_dir)))
 
train_dir = os.path.join(data_dir, "train")
valid_dir = os.path.join(data_dir, "val")
 
norm_mean = [0.485, 0.456, 0.406]
norm_std = [0.229, 0.224, 0.225]
 
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(norm_mean, norm_std),
])
 
valid_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(norm_mean, norm_std),
])
 
# 构建MyDataset实例
train_data = AntsDataset(data_dir=train_dir, transform=train_transform)
valid_data = AntsDataset(data_dir=valid_dir, transform=valid_transform)
 
# 构建DataLoder
train_loader = DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)
valid_loader = DataLoader(dataset=valid_data, batch_size=BATCH_SIZE)
 
# ============================ step 2/5 模型 ============================
 
# 1/3 构建模型
resnet18_ft = models.resnet18()
 
# 2/3 加载参数
# flag = 0
flag = 1
if flag:
    path_pretrained_model = os.path.join(BASEDIR, "..", "..", "data", "finetune_resnet18-5c106cde.pth")
    if not os.path.exists(path_pretrained_model):
        raise Exception("\n{} 不存在，请下载 07-02-数据-模型finetune.zip\n放到 {}下，并解压即可".format(
            path_pretrained_model, os.path.dirname(path_pretrained_model)))
    #加载与训练权重
    state_dict_load = torch.load(path_pretrained_model)
    resnet18_ft.load_state_dict(state_dict_load)
 
# 法1 : 冻结卷积层
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层
num_ftrs = resnet18_ft.fc.in_features
resnet18_ft.fc = nn.Linear(num_ftrs, classes)
 
 
resnet18_ft.to(device)
# ============================ step 3/5 损失函数 ============================
criterion = nn.CrossEntropyLoss()                                                   # 选择损失函数
 
# ============================ step 4/5 优化器 ============================
# 法2 : conv 小学习率
# flag = 0
flag = 1
if flag:
    fc_params_id = list(map(id, resnet18_ft.fc.parameters()))     # 返回的是parameters的 内存地址
    base_params = filter(lambda p: id(p) not in fc_params_id, resnet18_ft.parameters())
    optimizer = optim.SGD([
        {'params': base_params, 'lr': LR*0},   # 0
        {'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)     # 设置学习率下降策略
 
 
# ============================ step 5/5 训练 ============================
train_curve = list()
valid_curve = list()
 
for epoch in range(start_epoch + 1, MAX_EPOCH):
 
    loss_mean = 0.
    correct = 0.
    total = 0.
 
    resnet18_ft.train()
    for i, data in enumerate(train_loader):
 
        # forward
        inputs, labels = data
        inputs, labels = inputs.to(device), labels.to(device)
        outputs = resnet18_ft(inputs)
 
        # backward
        optimizer.zero_grad()
        loss = criterion(outputs, labels)
        loss.backward()
 
        # update weights
        optimizer.step()
 
        # 统计分类情况
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).squeeze().cpu().sum().numpy()
 
        # 打印训练信息
        loss_mean += loss.item()
        train_curve.append(loss.item())
        if (i+1) % log_interval == 0:
            loss_mean = loss_mean / log_interval
            print("Training:Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}".format(
                epoch, MAX_EPOCH, i+1, len(train_loader), loss_mean, correct / total))
            loss_mean = 0.
 
            # if flag_m1:
            print("epoch:{} conv1.weights[0, 0, ...] :\n {}".format(epoch, resnet18_ft.conv1.weight[0, 0, ...]))
 
    scheduler.step()  # 更新学习率
 
    # validate the model
    if (epoch+1) % val_interval == 0:
 
        correct_val = 0.
        total_val = 0.
        loss_val = 0.
        resnet18_ft.eval()
        with torch.no_grad():
            for j, data in enumerate(valid_loader):
                inputs, labels = data
                inputs, labels = inputs.to(device), labels.to(device)
 
                outputs = resnet18_ft(inputs)
                loss = criterion(outputs, labels)
 
                _, predicted = torch.max(outputs.data, 1)
                total_val += labels.size(0)
                correct_val += (predicted == labels).squeeze().cpu().sum().numpy()
 
                loss_val += loss.item()
 
            loss_val_mean = loss_val/len(valid_loader)
            valid_curve.append(loss_val_mean)
            print("Valid:\t Epoch[{:0>3}/{:0>3}] Iteration[{:0>3}/{:0>3}] Loss: {:.4f} Acc:{:.2%}".format(
                epoch, MAX_EPOCH, j+1, len(valid_loader), loss_val_mean, correct_val / total_val))
 
train_x = range(len(train_curve))
train_y = train_curve
 
train_iters = len(train_loader)
valid_x = np.arange(1, len(valid_curve)+1) * train_iters*val_interval # 由于valid中记录的是epochloss，需要对记录点进行转换到iterations
valid_y = valid_curve
 
plt.plot(train_x, train_y, label='Train')
plt.plot(valid_x, valid_y, label='Valid')
 
plt.legend(loc='upper right')
plt.ylabel('loss value')
plt.xlabel('Iteration')
plt.show()

---

三、GPU的使用

3.1  CPU VS GPU

CPU(Central Processing Unit，中央处理器)：主要包括控制器和运算器

GPU(Graphics Processing Unit，图形处理单元)：处理统一的，无依赖的大规模数据运算

 

---

3.2 数据迁移至GPU

数据迁移：通过to函数实现数据迁移

首先， 这个数据主要有两种： Tensor和Module

• CPU -> GPU： data.to(“cpu”)
• GPU -> CPU： data.to(“cuda”)

to函数： 转换数据类型/设备

（1）tensor.to(*args, **kwargs)

x = torch.ones((3,3))
x = x.to(torch.float64)    # 转换数据类型

x = torch.ones((3,3))
x = x.to("cuda")    # 设备转移

（2）module.to(*args, **kwargs) 

linear = nn.Linear(2,2)
linear.to(torch.double)  # 这样模型里面的可学习参数的数据类型变成float64

gpu1 = torch.device("cuda")
linear.to(gpu1)    # 把模型从CPU迁移到GPU

如果模型在GPU上， 那么数据也必须在GPU上才能正常运行。也就是说数据和模型必须在相同的设备上。

torch.cuda常用的方法：

• torch.cuda.device_count():   计算当前可见可用的GPU数
• torch.cuda.get_device_name():   获取GPU名称
• torch.cuda.manual_seed():   为当前GPU设置随机种子
• torch.cuda.manual_seed_all():   为所有可见可用GPU设置随机种子
• torch.cuda.set_device():   设置主GPU（默认GPU）为哪一个物理GPU（不推荐）

---

3.3 多GPU并行运算

多GPU的运行机制：
多 GPU 并行运算， 简单的说有多块 GPU，比如 4 块， 而这里面有个主 GPU， 当拿到了样本数据之后，比如主 GPU 拿到了 16 个样本， 那么它会经过 16/4=4 的运算，把数据分成 4 份， 自己留一份，然后把那 3 份分发到另外 3 块 GPU 上进行运算， 等其他的 GPU 运算完了之后， 主 GPU 再把结果收回来负责整合。多 GPU 并行运算可以大大节省时间。所以， 多 GPU 并行运算的三步：分发 -> 并行计算 -> 收回结果整合。

pytorch中多GPU并行运算机制的实现：
torch.nn.DataParallel:   包装模型，实现分发并行机制。

主要参数：

• module:   需要包装分发的模型
• device_ids:   可分发的gpu, 默认分发到所有的可见可用GPU， 通常这个参数不管它，而是在环境变量中管这个。
• output_device:   结果输出设备， 通常是输出到主GPU

下面看一下多GPU的使用：

---