pytorch学习系列(9):只保存精度最高的模型

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

if epoch % 50 == 0:
   torch.save(net.state_dict(),'%d.pth' % (epoch))

这样会导致保存的模型数量太多,占用硬盘空间,而且训练完成后寻找最优的模型也需要对照损失函数曲线去寻找,很不方便。
但如果每次都记录下损失函数的值,只保存验证集上损失最小的时候的模型,就更方便使用,其实设置一个判断条件就行。。。

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')

后来我发现,这种方式也有不科学之处,因为在测试集上损失最小的那个epoch的模型不一定就是最好的模型,不一定具有最好的泛化能力,所以这事还是玄学,不好说。

你可能感兴趣的:(pytorch)