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

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保留验证集上最好的模型)