pytorch保存模型运行时状态，记录点checkpoint

最新推荐文章于 2023-08-21 18:28:12 发布

呆萌的代Ma

最新推荐文章于 2023-08-21 18:28:12 发布

阅读量752

点赞数

分类专栏： pytorch/神经网络文章标签： pytorch

本文为CSDN博主"呆萌的代Ma"原创文章，转载请注明博客链接：https://blog.csdn.net/weixin_35757704/

本文链接：https://blog.csdn.net/weixin_35757704/article/details/119634679

版权

pytorch/神经网络专栏收录该内容

60 篇文章 28 订阅

订阅专栏

本文介绍了如何使用PyTorch保存和加载训练过程中的模型和优化器状态。通过`torch.save`将模型的`state_dict`、优化器的`state_dict`以及迭代次数等关键信息保存到文件，然后利用`torch.load`加载这些信息以继续训练或评估。示例中展示了线性回归模型的训练，并演示了保存和加载的具体操作。

摘要由CSDN通过智能技术生成

保存

torch.save({
    'epoch': epoch, # 保存迭代次数
    'model_state_dict': model.state_dict(), # 模型的状态
    'optimizer_state': optimizer.state_dict(), # 优化器的状态
}, 'checkpoint')  # 路径，这里每次保存都会覆盖掉checkpoint文件

当然字典里可以保存任意的内容，路径也可以根据epoch不同而改变
还有一种写法：

torch.save((epoch, model.state_dict(), optimizer.state_dict()), 'checkpoint')

加载

checkpoint = torch.load('checkpoint') # 加载到checkpoint
epoch = checkpoint['epoch'] # 读取epoch
model.load_state_dict(checkpoint['model_state_dict']) # 加载模型状态
optimizer.load_state_dict(checkpoint['optimizer_state']) # 加载优化器的状态

案例

import torch
import torch.nn as nn
import numpy as np


class LinearRegressionModel(nn.Module):
    def __init__(self, input_shape, output_shape):
        super(LinearRegressionModel, self).__init__()
        self.linear = nn.Linear(input_shape, output_shape)

    def forward(self, x):
        out = self.linear(x)
        return out


def train_model(x_train, y_train):
    # 指定参数与损失函数
    model = LinearRegressionModel(x_train.shape[1], 1)
    epochs = 10  # 迭代10次
    learning_rate = 0.01  # 学习率
    optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate)  # 优化函数
    criterion = nn.MSELoss()  # Loss使用MSE值，目标是使MSE最小

    for epoch in range(epochs):
        epoch += 1
        optimizer.zero_grad()  # 梯度清零
        outputs = model(x_train)  # 前向传播
        loss = criterion(outputs, y_train)  # 计算损失
        loss.backward()  # 返向传播
        optimizer.step()  # 更新权重参数
        # 保存
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state': optimizer.state_dict(),
        }, 'checkpoint')  # 这里每次保存都会覆盖掉checkpoint文件
        # 加载
        checkpoint = torch.load('checkpoint')  # 加载到checkpoint
        epoch = checkpoint['epoch']  # 读取epoch
        model.load_state_dict(checkpoint['model_state_dict'])  # 加载模型状态
        optimizer.load_state_dict(checkpoint['optimizer_state'])  # 加载优化器的状态


if __name__ == '__main__':
    x_train = torch.randn(100, 4)  # 生成100个4维的随机数，作为训练集的 X
    y_train = torch.randn(100, 1)  # 作为训练集的label
    train_model(x_train, y_train)