美文网首页
pytorch保存模型等相关参数,利用torch.save(),

pytorch保存模型等相关参数,利用torch.save(),

作者: _曹杰 | 来源:发表于2018-11-28 09:12 被阅读0次

    本文分为两部分,第一部分讲如何保存模型参数,优化器参数等等,第二部分则讲如何读取。

    假设网络为model = Net(), optimizer = optim.Adam(model.parameters(), lr=args.lr), 假设在某个epoch,我们要保存模型参数,优化器参数以及epoch


    一、


    1. 先建立一个字典,保存三个参数:

    state = {‘net':model.state_dict(), 'optimizer':optimizer.state_dict(), 'epoch':epoch}

    2.调用torch.save():

    torch.save(state, dir)

    其中dir表示保存文件的绝对路径+保存文件名,如'/home/qinying/Desktop/modelpara.pth'


    二、


    当你想恢复某一阶段的训练(或者进行测试)时,那么就可以读取之前保存的网络模型参数等。

    checkpoint = torch.load(dir)

    model.load_state_dict(checkpoint['net'])

    optimizer.load_state_dict(checkpoint['optimizer'])

    start_epoch = checkpoint['epoch'] + 1

    相关文章

      网友评论

          本文标题:pytorch保存模型等相关参数,利用torch.save(),

          本文链接:https://www.haomeiwen.com/subject/phooqqtx.html