pytorch 保存模型和加载模型

最新推荐文章于 2023-03-26 21:47:49 发布

原创

最新推荐文章于 2023-03-26 21:47:49 发布 · 935 阅读

0 ·

CC 4.0 BY-SA版权

这篇博客详细介绍了如何在PyTorch中保存已经训练好的模型，并且演示了如何在之后的训练过程中加载模型以继续训练。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

def save_model(save_dir, phase, name, epoch, f1score, model):
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    save_dir = os.path.join(save_dir, args.model)
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    save_dir = os.path.join(save_dir, phase)
    if not os.path.exists(save_dir):
        os.mkdir(save_dir)
    state_dict = model.state_dict()
    for key in state_dict.keys():
        state_dict[key] = state_dict[key].cpu()
    state_dict_all = {
        'state_dict': state_dict,
        'epoch': epoch,
        'f1score': f1score,
    }
    torch.save(state_dict_all, os.path.join(save_dir, '{:s}.ckpt'.format(name)))
    if 'best' in name and f1score > 0.3:
        torch.save(state_dict_all, os.path.join(save_dir, '{:s}_{:s}.ckpt'.format(name, str(epoch))))

pytorch 保存模型

pytorch 加载模型进行继续训练