TorchScript模型和普通PyTorch模型

TorchScript模型和普通PyTorch模型

PyTorch提供了两种主要的模型保存和加载机制，一种是基于Python的序列化，另一种是TorchScript。

普通的PyTorch模型（基于Python的序列化）:

• 保存: 使用torch.save(model.state_dict(), 'model_path.pth')，它保存了模型的权重和参数，但不保存模型的结构。
• 加载:
  1. 首先，您需要有模型的类定义。
  2. 创建该类的一个实例。
  3. 使用model.load_state_dict(torch.load('model_path.pth'))来加载权重。
• 特点:
  • 需要Python环境和模型的原始代码来加载和运行模型。
  • 保存的文件是Python特定的，并且依赖于特定的类结构。
  • 主要用于继续训练或在Python环境中进行推断。

保存模型:

torch.save(model.state_dict(), 'model_weights.pth')

加载模型:

model = ModelClass() 
model.load_state_dict(torch.load('model_weights.pth'))
model.eval() # 设置为评估模式

例程:

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 10)

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

model = SimpleModel()
torch.save(model.state_dict(), 'simple_model_weights.pth')

# 在其他地方或时间
loaded_model = SimpleModel()
loaded_model.load_state_dict(torch.load('simple_model_weights.pth'))
loaded_model.eval()

TorchScript模型:

• TorchScript是PyTorch的一个子集，它创建了一个可以独立于Python运行的序列化模型。
• 生成方法:
  1. Tracing: 使用torch.jit.trace方法。这涉及到通过模型运行一个输入示例，从而跟踪模型的执行路径。
  2. Scripting: 使用torch.jit.script方法。这转化Python代码到TorchScript，允许更复杂的模型和控制流。
• 保存: 使用torch.jit.save(traced_model, 'model_path.pt')。
• 加载: 使用torch.jit.load('model_path.pt')。注意，加载不需要原始的模型类定义。
• 特点:
  • 可以在没有Python运行时的环境中运行，如C++。
  • 提供了一种方法，将模型从Python转移到其他平台或部署环境。
  • 包含模型的完整定义，包括结构、权重和参数。

Tracing方法:

example_input = torch.randn(1, 10)
traced_model = torch.jit.trace(model, example_input)
torch.jit.save(traced_model, 'traced_model.pt')

Scripting方法:

scripted_model = torch.jit.script(model)
torch.jit.save(scripted_model, 'scripted_model.pt')

加载模型:

loaded_model = torch.jit.load('model_path.pt')

例程:

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 10)

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

model = SimpleModel()

# Tracing
example_input = torch.randn(1, 10)
traced_model = torch.jit.trace(model, example_input)
torch.jit.save(traced_model, 'traced_simple_model.pt')

# Scripting
scripted_model = torch.jit.script(model)
torch.jit.save(scripted_model, 'scripted_simple_model.pt')

# 加载模型
loaded_model = torch.jit.load('traced_simple_model.pt')

结论:

• 普通PyTorch模型: 适用于Python环境，需要模型的原始定义来加载。它是基于Python的序列化。
• TorchScript: 为了跨平台部署和运行，例如在C++环境中。它提供了一个独立于Python的序列化模型，可以在没有Python的环境中使用。