【深度学习】详解如何保存训练好的最优模型

目录

 

保存最佳模型的两种方式：

第一种：

第二种：

完整代码展示： 

---

保存最佳模型的两种方式：

第一种：

'''保存模型权重'''
torch.save(model.state_dict(),'model_parameter.pt')

他的读取model方式为：

'''1.读取权重'''
model=cnn().to(device)
model.load_state_dict(torch.load("model_parameter.pt"))

详细如何读取可以看我的这篇博客： Pytorch——模型的读取_AI_dataloads的博客-CSDN博客

第二种：

'''保存完整模型'''
torch.save(model,'best_model.pt')

他的读取model方式为：

'''2.读取完整模型'''
model=torch.load('best_model.pt')

详细如何读取模型可以看我的这篇博客： Pytorch——模型的读取_AI_dataloads的博客-CSDN博客

完整代码展示： 

1. 首先我们通过定义一个变量 best_acc，用于追踪最佳准确率。
2. 然后在我们的test函数中 声明 best_acc 为全局变量，以便在函数内部修改使用它。
3. 通过计算出的准确率correct，与best_acc比较，如果correct>best_acc:，则将此时的模型通过torch.save()保存下来。
4. 通过多轮训练与测试，找到准确率最高的模型，并保存下来生成pt文件，这样我们的模型也就保存下来了。
5. 后续如何使用我们已保存在pt文件可以查看我的这篇博客：

# best_acc=0：定义一个变量 best_acc，用于追踪最佳准确率。
best_acc=0
# 定义一个函数，用于测试模型在给定数据集上的性能。
def test(dataloader, model, loss_fn):
    # 声明 best_acc 为全局变量，以便在函数内部修改它。
    global best_acc
    size = len(dataloader.dataset)  # 获取总数据
    num = len(dataloader)
    model.eval()  # 进入测试模式
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model.forward(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
    test_loss /= num
    # 计算准确率。
    correct /= size
    acc_s.append(correct)
    loss_s.append(test_loss)
    '''保存最好的模型'''
    if correct>best_acc:
        best_acc=correct
        print(model.state_dict().keys())
        '''保存模型权重'''
        torch.save(model.state_dict(),'model_parameter.pt')
        '''保存完整模型'''
        torch.save(model,'best_model.pt')
    print(f'Test result:\nAccuracy:{correct * 100:.4f}%,Avg loss :{test_loss}')