Ptorch深度学习-------断点续训

最新推荐文章于 2024-06-13 10:55:54 发布

笨小蛙

最新推荐文章于 2024-06-13 10:55:54 发布

阅读量1.1k

点赞数 2

文章标签：深度学习 python pytorch

本文链接：https://blog.csdn.net/GREEN_cq/article/details/126782423

版权

1、参数设置一个标志：

resume：是否进行续训

initepoch：断点位置的epoch
2、checkpoint载入：

 resume = True
    initepoch = 0
    if resume:
        if os.path.isfile(os.path.join(args.save_model_path, 'latest_dice_loss.pth')) or os.path.join(
                args.save_model_path, 'latest_dice_loss.pth'):
            print("Resume from checkpoint...")
            checkpoint = torch.load(os.path.join(args.save_model_path, 'latest_dice_loss.pth'))
            model.module.load_state_dict(checkpoint['model_state_dict'])
            optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
            initepoch = checkpoint['epoch']
            print("====>loaded checkpoint (epoch{})".format(checkpoint['epoch']))
        else:
            print("====>no checkpoint found.")
            initepoch = 0

3、保存checkpoint包含的参数

checkpoint = {"model_state_dict": model.module.state_dict(),
                          "optimizer_state_dict": optimizer.state_dict(),
                          "epoch": epoch}
checkpoint_path = os.path.join(args.save_model_path, 'latest_dice_loss.pth') 

torch.save(checkpoint, checkpoint_path)