pytorch中,load_state_dict和torch.load的区别?

在 PyTorch 中,load_state_dicttorch.load 是两个不同的函数,用于不同的目的。

  1. torch.load:

    • 用途: 从磁盘加载一个保存的对象。这个对象可以是一个模型的整个状态字典(包含模型参数)、优化器状态字典、甚至是任意其他 Python 对象。

    • 用法 : 通常用于加载之前用 torch.save 保存的对象。

    • 示例 :

      python 复制代码
      # 保存对象
      torch.save(model.state_dict(), 'model.pth')
      torch.save(optimizer.state_dict(), 'optimizer.pth')
      
      # 加载对象
      model_state_dict = torch.load('model.pth')
      optimizer_state_dict = torch.load('optimizer.pth')
  2. load_state_dict:

    • 用途 : 将加载的状态字典(通常是模型参数)应用到一个模型实例上。这个函数通常用于将 torch.load 加载的状态字典应用到模型或优化器上。

    • 用法: 在模型或优化器实例上调用,用于将加载的状态字典设置为模型或优化器的当前状态。

    • 示例 :

      python 复制代码
      # 创建模型实例
      model = MyModel()
      
      # 加载并应用状态字典
      model.load_state_dict(torch.load('model.pth'))

总结

  • torch.load 用于从磁盘加载任意对象(通常是状态字典)。
  • load_state_dict 用于将加载的状态字典应用到模型或优化器实例上。

以下是一个完整的示例代码,演示如何保存和加载模型参数:

python 复制代码
import torch
import torch.nn as nn
import torch.optim as optim

# 定义模型
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.fc = nn.Linear(10, 1)

    def forward(self, x):
        return self.fc(x)

# 创建模型和优化器
model = MyModel()
optimizer = optim.SGD(model.parameters(), lr=0.001)

# 保存模型和优化器的状态字典
torch.save(model.state_dict(), 'model.pth')
torch.save(optimizer.state_dict(), 'optimizer.pth')

# 加载模型和优化器的状态字典
model.load_state_dict(torch.load('model.pth'))
optimizer.load_state_dict(torch.load('optimizer.pth'))

这段代码展示了如何定义一个简单的模型,保存它的状态字典,然后加载这些状态字典到新的模型和优化器实例中。

相关推荐
毋语天35 分钟前
FastAPI 进阶实战:请求体、文件上传、响应模型与数据校验
python·fastapi·api开发·数据校验·pydantic
T.i.s36 分钟前
parall scan(并行扫描)通俗理解
人工智能·深度学习
h64648564h1 小时前
CANN 昇腾训练食谱全景解读:cann-recipes-train 架构与使用指南
人工智能·深度学习
ZhengEnCi1 小时前
09a-斯坦福 CS336 作业一:BPE 分词器
python·神经网络
测试员周周2 小时前
【Appium 系列】第18节-重试与容错 — 移动端测试的稳定性保障
人工智能·python·功能测试·ui·单元测试·appium·测试用例
还是鼠鼠2 小时前
AI掘金头条新闻系统 (Toutiao News)-用户注册-创建用户
后端·python·mysql·fastapi·web
灰灰勇闯IT2 小时前
DeepSeek-R1 在 CANN 上的推理部署
pytorch·python·深度学习
努力学习_小白2 小时前
Inception V1——学习记录
pytorch·深度学习·inception v1
天才测试猿3 小时前
Jenkins+Docker自动化测试全攻略
自动化测试·软件测试·python·测试工具·docker·jenkins·测试用例
5201-3 小时前
向量数据库在 NPU 上的加速
数据库·pytorch·python