保存模型参数

state_dict就是一个简单的Python字典，它将模型中的可训练参数（比如weights和biases，batchnorm的running_mean、torch.optim参数等）通过将模型每层与层的参数张量之间一一映射，实现保存、更新、变化和再存储。

首先定义模型的保存位置和文件名（以.pth结尾）
再save模型权重到固定路径

PATH = './my_net.pth'
torch.save(model.state_dict(), PATH)
复制

加载所保存的模型权重

需要初始化之前所定义的网络

new_model = Net()
复制

再使用load_state_dict方法将权重加载进网络

new_model.load_state_dict(torch.load(PATH))
复制

保存测试中最优的参数

代码如下：

best_model_wts = copy.deepcopy(model.state_dict())
best_acc = 0.0

train_loss = []
train_acc = []
test_loss = []
test_acc = []

for epoch in range(epochs):
    epoch_loss, epoch_acc, epoch_test_loss, epoch_test_acc = fit(epoch,
                                                                 model,
                                                                 train_dl,
                                                                 test_dl)
    train_loss.append(epoch_loss)
    train_acc.append(epoch_acc)
    test_loss.append(epoch_test_loss)
    test_acc.append(epoch_test_acc)
    
    #判断是否是最优结果
    if epoch_test_acc > best_acc:
        best_acc = epoch_acc
        best_model_wts = copy.deepcopy(model.state_dict())        
#将最优结果保存
model.load_state_dict(best_model_wts)
model.eval()
复制

完整模型的保存和加载

直接使用torch.save方法对模型进行保存

PATH = './my_whole_model.pth'

torch.save(model, PATH)

new_model2 = torch.load(PATH)
new_model2.eval()
复制