PyTorch模型必须转为TorchScript(.pt/.ts)才能被Triton加载,需用torch.jit.script()或trace()导出并验证;目录结构须严格为models/模型名/版本号/model.pt+config.pbtxt;config中platform设为"pytorch_libtorch",dims不含batch维,max_batch_size=0;若需Python backend,则改platform为"python",提供model.py和requirements.txt;客户端输入须dtype=np.float32、C-contiguous、shape匹配config。

PyTorch模型必须先转成 TorchScript 才能被Triton加载
Triton不直接支持 torch.nn.Module 原生对象,也不读取 .pth 权重文件。它只认序列化后的 TorchScript 模型(.pt 或 .ts),且要求是 scripted 或 traced 形式,不能含 Python 控制流(如动态 for、if 分支未被 trace 覆盖)。
常见错误现象:RuntimeError: Unable to cast Python instance to C++ type 或 Triton 启动时报 Failed to load model,大概率是模型没正确 script 化。
- 优先用
torch.jit.script():能保留条件逻辑,适合含if、for的模型(如带 early-exit 或 dynamic depth 的结构) - 若报
Tracing failed,再尝试torch.jit.trace(),但需确保输入 shape 固定、控制流路径唯一 - 务必用
model.eval()和torch.no_grad()导出,否则 dropout/batchnorm 行为异常 - 导出后用
model(*example_input)验证前向结果与原模型一致
模型仓库目录结构必须严格符合Triton规范
Triton 通过文件系统结构识别模型版本、配置和后端类型。任意层级错位都会导致 Model not found 或 Invalid model configuration。
最小可行结构(以模型名 resnet50 为例):
立即学习“Python免费学习笔记(深入)”;
models/
└── resnet50/
├── 1/
│ └── model.pt
└── config.pbtxt
1/ 是版本号目录,必须是纯数字;model.pt 必须放在版本目录下;config.pbtxt 必须在模型名目录下(即与 1/ 同级)。
-
config.pbtxt中platform: "pytorch_libtorch"是固定写法,不是pytorch或torchscript -
max_batch_size设为 0 表示支持变长 batch(Triton 自动 batching 时必需) -
input和output的dims必须与 TorchScript 模型实际接受/返回的 tensor shape 一致(注意是否含 batch 维) - 若模型输入是
[3, 224, 224],config 里应写dims: [3, 224, 224],而非[1, 3, 224, 224](batch 维由 Triton 管理)
Python backend需显式声明依赖并禁用自动优化
Triton 的 PyTorch backend(pytorch_libtorch)默认启用 torchscript 模式,但某些算子(如自定义 torch.nn.Module 子类、非标准 op)可能触发 JIT fallback 失败。此时需切到 Python backend,并手动管理推理逻辑。
关键点:
- 在
config.pbtxt中改用platform: "python" - 在模型目录下新增
model.py,实现initialize()和execute()方法,加载model.pt并调用 -
model.py所需的第三方包(如torchvision)必须通过requirements.txt声明,且 Triton 启动时加--pip-install参数安装 - 禁用
torch.compile()或torch.backends.cudnn.benchmark = True—— 它们在多线程 execute 下可能引发状态冲突
客户端请求必须匹配模型 I/O 的 memory layout 和 data type
Triton 默认按 row-major、C-contiguous 内存布局解析输入,且对 dtype 敏感。PyTorch 导出时常用 float32,但客户端若传 float64 或非 contiguous array,会直接报 Invalid argument: input 'input0' is expected to have datatype TYPE_FP32 或静默截断。
- 客户端用
numpy.array(..., dtype=np.float32, order='C')构造输入,避免np.float64或np.uint8(除非 config 显式声明) - 图像预处理后务必调用
np.ascontiguousarray(),尤其经过cv2.resize或torchvision.transforms后易丢失 contiguous 属性 - batch 维必须由客户端组织好(如 4 张图就传 shape
(4, 3, 224, 224)),Triton 不会帮你堆叠 - 使用
tritonclient.http.InferenceServerClient时,infer()前检查inputs[0].shape和inputs[0].datatype是否与 config 一致


















