使用 TorchScript 导出模型

使用 PyTorch Lightning 构建可扩展 AI 模型

Sergiy Tkachuk

Director, GenAI Productivity

什么是 TorchScript？

独立于 Python
适合生产部署
示例：
- 部署到移动设备

一部手机

什么是 TorchScript？

独立于 Python
适合生产部署
示例：
- 部署到移动设备
- 生产系统中的高性能推理

一部手机和一台笔记本电脑，带有代表生产环境的齿轮

将模型转换为 TorchScript

两种转换方式：

torch.jit.trace：用示例输入跟踪执行
torch.jit.script：通过分析 Python 代码编译模型

何时使用：

简单模型用 trace
有控制流（如循环）的模型用 script

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def forward(self, x):
        return x * 2

model = SimpleModel()
scripted_model = torch.jit.script(model)

保存与加载 TorchScript 模型

保存模型：
- torch.jit.save：将脚本化模型保存到文件
加载模型：
- torch.jit.load：加载模型用于推理

# Save the model
torch.jit.save(scripted_mod,"model.pt")

# Load the model
loaded_model=torch.jit.load("model.pt")

使用 TorchScript 进行推理

步骤：
- 加载 TorchScript 模型
- 传入输入以获取预测
- 输出与 PyTorch 预测一致

示例输入：

输入张量：[1.0, 2.0, 3.0]

示例输出：

输出张量：[2.0, 4.0, 6.0]

# Perform inference
input_arr = [1.0, 2.0, 3.0]
input_tensor = torch.tensor(input_arr)

output = loaded_model(input_tensor)
print(output)

TorchScript 要点速览

torch.jit.trace：适用于静态模型
torch.jit.script：处理动态控制流
torch.jit.save：保存脚本化模型
torch.jit.load：加载用于推理

Vamos praticar!

使用 PyTorch Lightning 构建可扩展 AI 模型