Pytorch 存储模型的同时应该存储Adam或SGD状态

最新推荐文章于 2024-07-20 21:00:33 发布

青盏

最新推荐文章于 2024-07-20 21:00:33 发布

阅读量3.6k

点赞数 6

分类专栏： DL tools

本文链接：https://blog.csdn.net/qq_16234613/article/details/81192191

版权

DL tools 专栏收录该内容

54 篇文章 1 订阅

订阅专栏

完整存储模型方便再次训练

save_checkpoint({
    'epoch': epoch + 1,
    'arch': args.arch,
    'state_dict': model.state_dict(),
    'best_prec1': best_prec1,
    'optimizer' : optimizer.state_dict(),
}, is_best)


def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
    torch.save(state, filename)
    if is_best:
        shutil.copyfile(filename, 'model_best.pth.tar')

if args.resume:
    if os.path.isfile(args.resume):
        print("=> loading checkpoint '{}'".format(args.resume))
        checkpoint = torch.load(args.resume)
        args.start_epoch = checkpoint['epoch']
        best_prec1 = checkpoint['best_prec1']
        model.load_state_dict(checkpoint['state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        print("=> loaded checkpoint '{}' (epoch {})"
              .format(args.resume, checkpoint['epoch']))
    else:
        print("=> no checkpoint found at '{}'".format(args.resume))