TorchScript模型和普通PyTorch模型

本文主要是介绍TorchScript模型和普通PyTorch模型，希望对大家解决编程问题提供一定的参考价值，需要的开发者们随着小编来一起学习吧！

文章目录

TorchScript模型和普通PyTorch模型
- 普通的PyTorch模型（基于Python的序列化）:
- - 例程:
- TorchScript模型:
- - 例程:
  - 结论:

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 nnclass 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 nnclass 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')