美文网首页
Pytorch保留验证集上最好的模型

Pytorch保留验证集上最好的模型

作者: 向着光噜噜 | 来源:发表于2021-04-14 20:45 被阅读0次

    1.在训练过程中,需要保存模型来供测试使用,以前采用隔几个epoch就保存模型:

    if epoch % 50 == 0:

      torch.save(net.state_dict(),'%d.pth' % (epoch))

    这样会导致保存的模型数量太多,占用硬盘空间,而且训练完成后寻找最优的模型也需要对照损失函数曲线去寻找,很不方便。

    2.但如果每次都记录下损失函数的值,只保存验证集上损失最小的时候的模型,就更方便使用,其实设置一个判断条件就行。

    ————————————————

    min_loss = 100000 # 随便设置一个比较大的数

    for epoch in range(epochs):

        train()

        val_loss = val()

        if val_loss < min_loss:

            min_loss = val_loss

            print("save model")

            torch.save(net.state_dict(),'model.pth')

    相关文章

      网友评论

          本文标题:Pytorch保留验证集上最好的模型

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