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

Posted 2023-03-08

参考技术A 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')