note6 模型权重保存

# 保存模型参数
PATH = 'my_model.pth'
torch.save(model.state_dict(), PATH)   # 将模型中的可训练参数用字典形式保存,路径是PATH
# 加载
new_model = Net()
new_model.load_state_dict(torch.load(PATH))
new_model.to('cuda')

保存最优权重

for epoch in range(epochs):
    epoch_loss, epoch_acc, epoch_test_loss, epoch_test_acc = fit(epoch, model, train_dl, test_dl)
    # 保存最优权重
    if epoch_test_acc > best_acc:
        best_model_wts = copy.deepcopy(model.state_dict())
        best_acc = epoch_test_acc

    train_loss.append(epoch_loss)
    train_acc.append(epoch_acc)
    test_loss.append(epoch_test_loss)
    test_acc.append(epoch_test_acc)
# 加载最优
model.load_state_dict(best_model_wts)

测试时需要放到测试模式