怎么将 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。