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')
网友评论