【机器学习】如何在训练过程中挑选比较好的模型保存(pytorch)常用思路代码解析

56 阅读 0 评论 37 点赞

我是靠谱客的博主安静书本，最近开发中收集的这篇文章主要介绍【机器学习】如何在训练过程中挑选比较好的模型保存(pytorch)常用思路代码解析，觉得挺不错的，现在分享给大家，希望可以做个参考。

概述

常用思路

先设置大的epoch，跑完一遍，绘制每个epoch的训练集和测试集的损失折线图，然后根据测试集的损失，选一个中间的epoch即可，测试集的损失，一定程度上反映了模型的预测性能
根据1的结果，在训练结束后保存（常用）
每个epoch保存一个模型，根据打印的最好的结果选。
根据测试集（验证集）合的实验结果保存模型（可以是精度达到多少，可以是loss降低到多少，可以是记录两次损失之差，小于某个值就break）

代码解析

train_loss_hist = []
test_loss_hist = []
# 建议20次
for epoch in tqdm(range(2)):
    cnn.train() #训练部分
    running_loss = 0.0 #batch损失累计
    for i , data in enumerate(train_loader):
        images, labels = data
        outputs = cnn(images)
        loss = criterion(outputs, labels) # 损失求和
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        running_loss += loss.item() # 一个batch的损失
        
        if (i%250 == 0):
            # 每250个  250 x 16 进行一次记录 和测试
            cnn.eval() #验证部分
            with torch.no_grad():
                for test_data in test_loader:
                    test_images, test_label = test_data
                    test_outputs = cnn(test_images)
                    test_loss = criterion(test_outputs, test_label)
            train_loss_hist.append(running_loss/250) # 250个batch的平均损失
            test_loss_hist.append(test_loss.item())
            running_loss = 0.0
            #torch.save(cnn,"model/cnn_image_model_epoch_{}_step{}.pkl".format(epoch,i))
            print("epoch:{} step:{} train Loss:{} test Loss:{}".format(epoch+1,i,loss.item(),test_loss.item()))
    torch.save(cnn,"model/cnn_image_model_epoch_{}_train Loss_{}_test_Loss_{}.pkl".format(epoch,loss.item(),test_loss.item())) #每个epoch保存一次
    print("epoch:{} train Loss:{} test Loss:{}".format(epoch+1,loss.item(),test_loss.item())) #最优一个250*batch的损失作为这一轮的损失