PyTorch实现断点继续训练（notebook源码下载）

1. 固定随机种子 和 cudnn.benchmark = true

import os
import torch
import torchvision
import math
import random
import numpy as np

# 固定CPU和CUDA的随机种子，使结果可以复现
torch.manual_seed(0)
# 如果为True，每次返回的卷积算法将是确定的，即默认算法
torch.backends.cudnn.deterministic = True
# 如果为True，则使cuDNN对多个卷积算法进行基准测试并选择最快的。
torch.backends.cudnn.benchmark = False

2. 初始化模型，设置优化器和学习率调整算法

model = torchvision.models.resnet18()
model.fc = torch.nn.Linear(512, 2)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.95, nesterov=True)
# 余弦退火算法周期
steps = 10
# 学习率缩小的比率
lrf = 0.01
my_lambda = lambda x: ((1 - math.cos(x * math.pi / steps)) / 2) * (lrf - 1) + 1
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=my_lambda)

3. 模型训练过程

start_epoch = 0
inputs = torch.rand((2,3,32,32))

for epoch in range(start_epoch+1,20):
    output = model(inputs)
    loss = torch.nn.functional.cross_entropy(output, torch.tensor([0,1], dtype=torch.long))
    print(loss)
    loss.backward()
    optimizer.zero_grad()
    optimizer.step()
    scheduler.step()

4. 保存模型权重、优化器参数、以及参数跟新策略

checkpoint = {
    "model": model.state_dict(),
    'optimizer': optimizer.state_dict(),
    "epoch": epoch,
    'lr_schedule': scheduler.state_dict()
}
if not os.path.isdir("./model_parameter"):
    os.mkdir("./model_parameter")
torch.save(checkpoint, './model_parameter/check.pth')

5. 加载模型权重、优化器参数、以及参数跟新策略

path_checkpoint = "./model_parameter/check.pth"  # 断点路径

checkpoint = torch.load(path_checkpoint)  # 加载断点
model.load_state_dict(checkpoint['model'])  # 加载模型可学习参数
optimizer.load_state_dict(checkpoint['optimizer'])  # 加载优化器参数
start_epoch = checkpoint['epoch']  # 设置开始的epoch
scheduler.load_state_dict(checkpoint['lr_schedule'])#加载lr_scheduler

6. 整体过程如下：

import os
import torch
import torchvision
import math
import random
import numpy as np

# 固定CPU和CUDA的随机种子，使结果可以复现
torch.manual_seed(0)
# 如果为True，每次返回的卷积算法将是确定的，即默认算法
torch.backends.cudnn.deterministic = True
# 如果为True，则使cuDNN对多个卷积算法进行基准测试并选择最快的。
torch.backends.cudnn.benchmark = False

model = torchvision.models.resnet18()
model.fc = torch.nn.Linear(512, 2)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.95, nesterov=True)
# 余弦退火算法周期
steps = 100
# 学习率缩小的比率
lrf = 0.01
my_lambda = lambda x: ((1 - math.cos(x * math.pi / steps)) / 2) * (lrf - 1) + 1
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=my_lambda)

start_epoch = 0
resume = True

if resume:

    path_checkpoint = "./model_parameter/check.pth"  # 断点路径

    checkpoint = torch.load(path_checkpoint)  # 加载断点
    model.load_state_dict(checkpoint['model'])  # 加载模型可学习参数
    optimizer.load_state_dict(checkpoint['optimizer'])  # 加载优化器参数
    start_epoch = checkpoint['epoch']  # 设置开始的epoch
    scheduler.load_state_dict(checkpoint['lr_schedule'])#加载lr_scheduler

inputs = torch.rand((2,3,32,32))

for epoch in range(start_epoch+1,start_epoch+10):
    output = model(inputs)
    loss = torch.nn.functional.cross_entropy(output, torch.tensor([0,1], dtype=torch.long))
    print(loss)
    loss.backward()
    optimizer.zero_grad()
    optimizer.step()
    scheduler.step()

    checkpoint = {
        "model": model.state_dict(),
        'optimizer': optimizer.state_dict(),
        "epoch": epoch,
        'lr_schedule': scheduler.state_dict()
    }
    if not os.path.isdir("./model_parameter"):
        os.mkdir("./model_parameter")
    torch.save(checkpoint, './model_parameter/check.pth')