pytorch保存模型运行时状态,记录点checkpoint/torch.save()--转载自 “呆萌的代Ma”

该博客介绍了如何在PyTorch中保存和加载模型及其优化器的状态。通过`torch.save`函数,可以将模型的迭代次数、状态字典和优化器的状态保存到文件,方便后续继续训练。同时,展示了如何使用`torch.load`恢复这些状态,包括模型参数和优化器的状态。示例中给出了一个简单的线性回归模型训练过程,展示了保存和加载的完整流程。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

本文目录

保存

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)
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值