pytorch中，load_state_dict和torch.load的区别？

在 PyTorch 中，load_state_dict 和 torch.load 是两个不同的函数，用于不同的目的。

1. torch.load:

   • 用途: 从磁盘加载一个保存的对象。这个对象可以是一个模型的整个状态字典（包含模型参数）、优化器状态字典、甚至是任意其他 Python 对象。
   • 用法: 通常用于加载之前用 torch.save 保存的对象。
   • 示例:

     # 保存对象
     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')
     1
     2
     3
     4
     5
     6
     7

2. load_state_dict:

   • 用途: 将加载的状态字典（通常是模型参数）应用到一个模型实例上。这个函数通常用于将 torch.load 加载的状态字典应用到模型或优化器上。
   • 用法: 在模型或优化器实例上调用，用于将加载的状态字典设置为模型或优化器的当前状态。
   • 示例:

     # 创建模型实例
     model = MyModel()
     
     # 加载并应用状态字典
     model.load_state_dict(torch.load('model.pth'))
     1
     2
     3
     4
     5

总结

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

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

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'))
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

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