PyTorch实现断点继续训练

梦里梦外; 2023-01-11 13:13 118阅读 0赞

> 点击上方“**机器学习与生成对抗网络**”，关注"星标"
> 
> 获取有趣、好玩的前沿干货！

作者：知乎—HUST小菜鸡

地址：https://www.zhihu.com/people/panda-9-6

## 最近在尝试用CIFAR10训练分类问题的时候，由于数据集体量比较大，训练的过程中时间比较长，有时候想给停下来，但是停下来了之后就得重新训练，之前师兄让我们学习断点继续训练及继续训练的时候注意epoch的改变等，今天上午给大致整理了一下，不全面仅供参考 ##

##     ##

**01**

**模型的保存与加载**

PyTorch中的保存（序列化，从内存到硬盘）与反序列化（加载，从硬盘到内存）

torch.save主要参数：obj：对象 、f：输出路径

torch.load 主要参数 ：f：文件路径 、map\_location：指定存放位置、 cpu or gpu

**模型的保存的两种方法：**

### 1、保存整个Module ###

torch.save(net, path)

### 2、保存模型参数 ###

state_dict = net.state_dict()torch.save(state_dict , path)

##     ##

**02**

**模型的训练过程中保存**

checkpoint = {        "net": model.state_dict(),        'optimizer':optimizer.state_dict(),        "epoch": epoch    }

将网络训练过程中的网络的权重，优化器的权重保存，以及epoch 保存，便于继续训练恢复

在训练过程中，可以根据自己的需要，每多少代，或者多少epoch保存一次网络参数，便于恢复，提高程序的鲁棒性。

checkpoint = {
            "net": model.state_dict(),
            'optimizer':optimizer.state_dict(),
            "epoch": epoch
        }
        if not os.path.isdir("./models/checkpoint"):
            os.mkdir("./models/checkpoint")
        torch.save(checkpoint, './models/checkpoint/ckpt_best_%s.pth' %(str(epoch)))

通过上述的过程可以在训练过程自动在指定位置创建文件夹，并保存断点文件

![format_png][]

**03**

**模型的断点继续训练**

if RESUME:    path_checkpoint = "./models/checkpoint/ckpt_best_1.pth"  # 断点路径    checkpoint = torch.load(path_checkpoint)  # 加载断点
        model.load_state_dict(checkpoint['net'])  # 加载模型可学习参数
        optimizer.load_state_dict(checkpoint['optimizer'])  # 加载优化器参数    start_epoch = checkpoint['epoch']  # 设置开始的epoch

指出这里的是否继续训练，及训练的checkpoint的文件位置等可以通过argparse从命令行直接读取，也可以通过log文件直接加载，也可以自己在代码中进行修改。关于argparse参照我的这一篇文章：

https://zhuanlan.zhihu.com/p/133285373

**04**

**重点在于epoch的恢复**

start_epoch = -1
    if RESUME:    path_checkpoint = "./models/checkpoint/ckpt_best_1.pth"  # 断点路径    checkpoint = torch.load(path_checkpoint)  # 加载断点
        model.load_state_dict(checkpoint['net'])  # 加载模型可学习参数
        optimizer.load_state_dict(checkpoint['optimizer'])  # 加载优化器参数    start_epoch = checkpoint['epoch']  # 设置开始的epoch
    
    
    for epoch in  range(start_epoch + 1 ,EPOCH):    # print('EPOCH:',epoch)    for step, (b_img,b_label) in enumerate(train_loader):        train_output = model(b_img)        loss = loss_func(train_output,b_label)        # losses.append(loss)        optimizer.zero_grad()        loss.backward()        optimizer.step()

通过定义start\_epoch变量来保证继续训练的时候epoch不会变化

![format_png 1][]

断点继续训练

**05**

**初始化随机数种子**

import torchimport randomimport numpy as np
    def set_random_seed(seed = 10,deterministic=False,benchmark=False):    random.seed(seed)    np.random(seed)    torch.manual_seed(seed)    torch.cuda.manual_seed_all(seed)    if deterministic:        torch.backends.cudnn.deterministic = True    if benchmark:        torch.backends.cudnn.benchmark = True

关于torch.backends.cudnn.deterministic和torch.backends.cudnn.benchmark详见

https://www.cnblogs.com/captain-dl/p/11938864.html

https://blog.csdn.net/zxyhhjs2017/article/details/91348108

benchmark用在输入尺寸一致，可以加速训练，deterministic用来固定内部随机性

###  ###

**06**

**多步长SGD继续训练**

在简单的任务中，我们使用固定步长（也就是学习率LR）进行训练，但是如果学习率lr设置的过小的话，则会导致很难收敛，如果学习率很大的时候，就会导致在最小值附近，总会错过最小值，loss产生震荡，无法收敛。所以这要求我们要对于不同的训练阶段使用不同的学习率，一方面可以加快训练的过程，另一方面可以加快网络收敛。

采用多步长 torch.optim.lr\_scheduler的多种步长设置方式来实现步长的控制，lr\_scheduler的各种使用推荐参考如下教程：

https://www.cnblogs.com/devilmaycry812839668/p/10630302.html

所以我们在保存网络中的训练的参数的过程中，还需要保存lr\_scheduler的state\_dict，然后断点继续训练的时候恢复

#这里我设置了不同的epoch对应不同的学习率衰减，在10->20->30，学习率依次衰减为原来的0.1，即一个数量级
    lr_schedule = torch.optim.lr_scheduler.MultiStepLR(optimizer,milestones=[10,20,30,40,50],gamma=0.1)
    optimizer = torch.optim.SGD(model.parameters(),lr=0.1)
    
    
    for epoch in range(start_epoch+1,80):
        optimizer.zero_grad()
        optimizer.step()
        lr_schedule.step()
    
    
        if epoch %10 ==0:
            print('epoch:',epoch)
            print('learning rate:',optimizer.state_dict()['param_groups'][0]['lr'])

lr的变化过程如下：

epoch: 10learning rate: 0.1epoch: 20learning rate: 0.010000000000000002epoch: 30learning rate: 0.0010000000000000002epoch: 40learning rate: 0.00010000000000000003epoch: 50learning rate: 1.0000000000000004e-05epoch: 60learning rate: 1.0000000000000004e-06epoch: 70learning rate: 1.0000000000000004e-06

我们在保存的时候，也需要对lr\_scheduler的state\_dict进行保存，断点继续训练的时候也需要恢复lr\_scheduler

#加载恢复if RESUME:    path_checkpoint = "./model_parameter/test/ckpt_best_50.pth"  # 断点路径    checkpoint = torch.load(path_checkpoint)  # 加载断点
        model.load_state_dict(checkpoint['net'])  # 加载模型可学习参数
        optimizer.load_state_dict(checkpoint['optimizer'])  # 加载优化器参数    start_epoch = checkpoint['epoch']  # 设置开始的epoch    lr_schedule.load_state_dict(checkpoint['lr_schedule'])#加载lr_scheduler
    #保存for epoch in range(start_epoch+1,80):
        optimizer.zero_grad()
        optimizer.step()    lr_schedule.step()
    
        if epoch %10 ==0:        print('epoch:',epoch)        print('learning rate:',optimizer.state_dict()['param_groups'][0]['lr'])        checkpoint = {            "net": model.state_dict(),            'optimizer': optimizer.state_dict(),            "epoch": epoch,            'lr_schedule': lr_schedule.state_dict()        }        if not os.path.isdir("./model_parameter/test"):            os.mkdir("./model_parameter/test")        torch.save(checkpoint, './model_parameter/test/ckpt_best_%s.pth' % (str(epoch)))

##     ##

**07**

**保存最好的结果**

每一个epoch中的每个step会有不同的结果，可以保存每一代最好的结果，用于后续的训练

## **第一次实验代码** ##

RESUME = True
    
    
    EPOCH = 40
    LR = 0.0005
    
    
    
    
    model = cifar10_cnn.CIFAR10_CNN()
    
    
    print(model)
    optimizer = torch.optim.Adam(model.parameters(),lr=LR)
    loss_func = nn.CrossEntropyLoss()
    
    
    start_epoch = -1
    
    
    
    
    if RESUME:
        path_checkpoint = "./models/checkpoint/ckpt_best_1.pth"  # 断点路径
        checkpoint = torch.load(path_checkpoint)  # 加载断点
    
    
        model.load_state_dict(checkpoint['net'])  # 加载模型可学习参数
    
    
        optimizer.load_state_dict(checkpoint['optimizer'])  # 加载优化器参数
        start_epoch = checkpoint['epoch']  # 设置开始的epoch
    
    
    
    
    
    
    for epoch in  range(start_epoch + 1 ,EPOCH):
        # print('EPOCH:',epoch)
        for step, (b_img,b_label) in enumerate(train_loader):
            train_output = model(b_img)
            loss = loss_func(train_output,b_label)
            # losses.append(loss)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    
    
            if step % 100 == 0:
                now = time.time()
                print('EPOCH:',epoch,'| step :',step,'| loss :',loss.data.numpy(),'| train time: %.4f'%(now-start_time))
    
    
        checkpoint = {
            "net": model.state_dict(),
            'optimizer':optimizer.state_dict(),
            "epoch": epoch
        }
        if not os.path.isdir("./models/checkpoint"):
            os.mkdir("./models/checkpoint")
        torch.save(checkpoint, './models/checkpoint/ckpt_best_%s.pth' %(str(epoch)))

**更新实验代码**

optimizer = torch.optim.SGD(model.parameters(),lr=0.1)
    lr_schedule = torch.optim.lr_scheduler.MultiStepLR(optimizer,milestones=[10,20,30,40,50],gamma=0.1)
    start_epoch = 9
    # print(schedule)
    
    
    
    
    if RESUME:
        path_checkpoint = "./model_parameter/test/ckpt_best_50.pth"  # 断点路径
        checkpoint = torch.load(path_checkpoint)  # 加载断点
    
    
        model.load_state_dict(checkpoint['net'])  # 加载模型可学习参数
    
    
        optimizer.load_state_dict(checkpoint['optimizer'])  # 加载优化器参数
        start_epoch = checkpoint['epoch']  # 设置开始的epoch
        lr_schedule.load_state_dict(checkpoint['lr_schedule'])
    
    
    for epoch in range(start_epoch+1,80):
    
    
        optimizer.zero_grad()
    
    
        optimizer.step()
        lr_schedule.step()
    
    
    
    
        if epoch %10 ==0:
            print('epoch:',epoch)
            print('learning rate:',optimizer.state_dict()['param_groups'][0]['lr'])
            checkpoint = {
                "net": model.state_dict(),
                'optimizer': optimizer.state_dict(),
                "epoch": epoch,
                'lr_schedule': lr_schedule.state_dict()
            }
            if not os.path.isdir("./model_parameter/test"):
                os.mkdir("./model_parameter/test")
            torch.save(checkpoint, './model_parameter/test/ckpt_best_%s.pth' % (str(epoch)))

文仅交流，不代表本公众号赞同观点或对内容负责，版权归原作者，如侵删

**猜您喜欢：**

**[超100篇！CVPR 2020最全GAN论文梳理汇总！][100_CVPR 2020_GAN]**

[**附下载 | 《Python进阶》中文版**][_ _Python]

[**附下载 | 经典《Think Python》中文版**][_ _Think Python]

[**附下载 | 《Pytorch模型训练实用教程》**][_ _Pytorch]

[**附下载 | 最新2020李沐《动手学深度学习》**][_ _2020]

**[附下载 | 《可解释的机器学习》中文版][Link 1]**

[**附下载 |《TensorFlow 2.0 深度学习算法实战》**][_TensorFlow 2.0]

**[附下载 | 超100篇！CVPR 2020最全GAN论文梳理汇总！][100_CVPR 2020_GAN]**

**[附下载 |《计算机视觉中的数学方法》分享][Link 2]**

![format_png 2][]

[format_png]: /images/20221119/b335527ef2814c16837f20fd7c673659.png
[format_png 1]: /images/20221119/44046c80a43a4db58ba1e1be87047e77.png
[100_CVPR 2020_GAN]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a94cc95e205ae880cb9518cb16a54c0abde69e5a8654f7fedf70a0cc990775265bf28a1c&idx=1&mid=2247486455&scene=21&sn=94156f53845cdf07e859e343862be745#wechat_redirect
[_ _Python]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29abb6c95e22a04d1eb019f964c8ae97b2f3af98c22f730f5d5ed2aea9fdb320b00c688967&idx=2&mid=2247485709&scene=21&sn=55a51c98ab48f393ba305c082ba0e4fe#wechat_redirect
[_ _Think Python]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a418c95e2d0e9a3de46e0e4d1b8d4947154d8639418a511a3091300aa31699932046556e&idx=2&mid=2247485091&scene=21&sn=f569d5d70c2e6c9cc567fbc5147787a4#wechat_redirect
[_ _Pytorch]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29aa78c95e236e380c9fd2ae79917d7451ccdd4803c5e6047a16f921d1b61d18efc79ada7c&idx=2&mid=2247485635&scene=21&sn=90608bce32fe353cf3c64b39e3406bc3#wechat_redirect
[_ _2020]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a4a4c95e2db2aab16bdcefa78795bbb7efbc70b533bb197d73f617558f2f8a523120c6e3&idx=2&mid=2247484959&scene=21&sn=563999cc8d48498beed00270844ceec4#wechat_redirect
[Link 1]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a4bac95e2dac758c4aac10a3aac49eb062bc0829d4ef2ebc8e6d5069d246385c5084a2b4&idx=1&mid=2247484929&scene=21&sn=523220858287ed3e6105919703c8e833#wechat_redirect
[_TensorFlow 2.0]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29a5b0c95e2ca6599d4a91d5d40900566b7525de3dd85259e7a6c173aa26aa7ea3641d1aaf&idx=2&mid=2247485195&scene=21&sn=e3387067db29d1759755c6c2ea96c870#wechat_redirect
[Link 2]: http://mp.weixin.qq.com/s?__biz=MzU5MTgzNzE0MA%3D%3D&chksm=fe29aa8fc95e23993e057fb0dbe09f2810ea07184a439c173f404a7ab657cadcde3d24efa52a&idx=2&mid=2247485492&scene=21&sn=c4fb135de02cd52f06bc61b0e79476e4#wechat_redirect
[format_png 2]: /images/20221119/e951be18f5b94eff9a8eeecee72d4570.png