怎么将 PyTorch 模型导出为 TorchScript（.pt 或 .pb 文件

将 PyTorch 模型导出为 TorchScript（.pt 或 .pb 文件），通常用于跨平台部署或在非Python环境中运行。以下是基本步骤：
导出为.pt（保存为脚本模式）

使用PyTorch的torch.jit模块，首先确保模型是可序列化的（即，没有non_blocking、volatile属性，也不是nn.DataParallel实例等）。

import torch

# 假设你已经有了一个名为model的PyTorch模型
model.eval()

# 将模型放到脚本模式
traced_script_module = torch.jit.trace(model, example_input)

# 保存为.pt文件
traced_script_module.save("model.pt")
登录后复制
这里的example_input是你希望模型在导出后接受的实际输入数据，以便生成对应的计算图。
导出为.pb（保存为二进制文件）

如果需要更低级别的二进制格式，PyTorch不直接支持.pb，但可以借助其他库，比如onnx先转换为ONNX，然后转为.pb。注意这一步可能受限于模型架构的兼容性。

import torch.onnx

# 先将模型转换为ONNX格式
torch.onnx.export(model, example_input, "model.onnx")

# 然后使用ONNX-TensorRT或ONNX-to-TF转换为.pb文件
# (这里是一个假设的命令，实际可能需要安装额外的库)
!python -m tf2onnx.convert --input_model=model.onnx --output_model=model.pb --opset=11
登录后复制
请注意，tf2onnx和onnx-tf等工具可能需要单独安装，并且不是所有PyTorch模型都能无损地转换到ONNX。