背景痛点：ChatTTS 原生部署到底卡在哪？
第一次把 ChatTTS 搬到服务器，我踩了三个大坑：
- 框架捆绑：PyTorch 1.13 + CUDA 11.7 的“黄金组合”在 Ubuntu 20.04 上跑得好好的，一到 CentOS 7 就缺 so，折腾两天才把 glibc 对齐。
- 推理延迟：同一段 8 s 音频，在 RTX-3060 上平均 2.7 s，CPU 更是飙到 12 s，线上并发一高就排队。
- 内存暴涨：每来一条请求就加载一次 Dict 和 Decoder，峰值 6 GB，Kubernetes 直接把 Pod 重启。
一句话：ChatTTS 的“训练友好”不等于“部署友好”，必须给它做一次“断舍离”。
技术选型：为什么最后选了 ONNX？
我把 TensorRT、TorchScript、ONNX 放在同一张表上跑 100 条 8 s 文本，硬件 RTX-3060 / Intel-12700H，结果如下：
| 方案 | 平均延迟 | 峰值显存 | 跨平台 | 备注 |
|---|---|---|---|---|
| TorchScript | 2.3 s | 3.8 GB | 需 libtorch | 动态 shape 支持差 |
| TensorRT 8.6 | 1.1 s | 2.9 GB | 仅限 NVIDIA | 编译 20 min，算子不支持 GELU |
| ONNX Runtime 1.17 | 1.2 s | 2.7 GB | 全平台 | 动态轴一次搞定 |
TensorRT 确实最快，但 ChatTTS 里 GELU、LayerNorm 版本多，手写 plugin 维护成本高；TorchScript 在 transformer 动态长度上总报错。ONNX 属于“90 分且不挑硬件”，于是敲定。
核心实现：30 行代码完成 PyTorch → ONNX
下面脚本在 ChatTTS v0.2 官方 checkpoint 上验证通过，Python 3.9、torch 2.1。关键点是：
- 把长度维度设成动态，避免推理时重复建图
- 用
opset_version=14保证 MultiHeadAttention 被官方支持 - 对
forward()做包装，只导出核心 TTS 链（省略 speaker embedding 的预处理，放到后处理用 numpy 算）
# export_onnx.py import torch import ChatTTS from pathlib import Path def export(): # 1. 加载官方权重 chat = ChatTTS.Chat() chat.load(compile=False) # 关掉 torch.compile，避免算子融合 model = chat.model.gpt # 只导出 GPT 部分，vocoder 用原框架 # 2. 构造伪输入 x = torch.randint(0, 512, (1, 100)) # token x_len = torch.tensor([100], dtype=torch.long) spk = torch.randn(1, 256) # speaker vector # 3. 动态轴 dynamic_axes = { "x": {0: "batch", 1: "len"}, "x_len": {0: "batch"}, "spk": {0: "batch"}, "output": {0: "batch", 1: "len"} } # 4. 导出 torch.onnx.export( model, (x, x_len, spk), "chatts_gpt.onnx", input_names=["x", "x_len", "spk"], output_names=["output"], dynamic_axes=dynamic_axes, opset_version=14, do_constant_folding=True ) print(" ONNX 已写入：chatts_gpt.onnx") if __name__ == "__main__": export()运行后得到 480 MB 的chatts_gpt.onnx，节点 742 个，OP 全部在官方支持列表。
ONNX Runtime 推理：内存池 + IOBinding 双优化
推理脚本里我习惯做三件事：
- 复存池：把
SessionOptions.enable_cpu_mem_arena打开，避免频繁 malloc - IOBinding：GPU 场景下把输入/输出 tensor 直接绑到 cuda 内存，省一次 copy
- 复用 InferenceSession：多线程环境下一个进程只建一次 Session，线程安全由 ORT 内部保证
# infer_onnx.py import onnxruntime as ort import numpy as np providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] sess_options = ort.SessionOptions() sess_options.enable_cpu_mem_arena = True sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL session = ort.InferenceSession("chatts_gpt.onnx", sess_options, providers=providers) def synthesize(tokens, speaker): # tokens: [1, L] int64 x_len = np.array([tokens.shape[1]], dtype=np.int64) audio = session.run( None, { "x": tokens, "x_len": x_len, "spk": speaker } )[0] return audio单条 8 s 音频在 RTX-3060 上延迟降到 1.2 s，CPU 线程池并发 4 路时吞吐量 3.2→9.6 条/分钟，显存稳定在 2.7 GB。
性能测试：数据说话
为了排除“实验室误差”，我在三种硬件各跑 200 条文本，取 P50/P99 延迟和最大吞吐，结果如下：
| 硬件 | 方案 | P50 延迟 | P99 延迟 | 最大吞吐 (条/分) |
|---|---|---|---|---|
| RTX-3060 | PyTorch | 2.7 s | 3.1 s | 3.2 |
| RTX-3060 | ONNX | 1.2 s | 1.4 s | 9.6 |
| Tesla-T4 | ONNX | 1.3 s | 1.5 s | 9.2 |
| Intel-12700H | ONNX-CPU | 4.8 s | 5.5 s | 2.1 |
结论：ONNX 版本在 GPU 上延迟减半、吞吐 ×3；CPU 也能接受，只是不适合实时场景。
避坑指南：把踩过的坑一次说清
算子不支持
ChatTTS 早期用torch.nn.GELU(approximate="tanh")，ONNX 默认只认erf版。解决：导出前全局替换为torch.nn.GELU(approximation="none")，再转 ONNX。多线程安全
ONNX Runtime 的 InferenceSession 是线程安全，但run()的输入 dict 必须保证每个线程独立；我曾把同一个np.array传进多线程，结果输出随机串音。解决：用array.copy()或者预分配 buffer。量化掉精度
我试把 GPT 部分用onnxruntime.quantization.quantize_dynamic()转成 INT8，模型体积 480→130 MB，但 MOS 分从 4.1 掉到 3.4。解决：只对 MatMul 权重做量化，跳过 embedding 和 LayerNorm，MOS 降到 3.9，体积 220 MB，可接受。
总结与延伸：下一步还能怎么卷？
ONNX 让 ChatTTS 脱离 PyTorch“温室”，但 480 MB 仍然不够边缘友好。我的下一步计划：
- 结构化剪枝：把 GPT 里 16 头注意力剪成 12 头，再用
onnxruntime-training做微调，预计体积 −30% - 知识蒸馏：训练一个 6 层小模型去拟合 20 层大模型，目标在 CPU 实时 ≤1 s
- 结合 NNAPI / CoreML：手机端把同一套 ONNX 转成端侧加速器，实现真正的“一套模型，多端运行”
如果你也在做语音合成落地，欢迎交换数据：同样的文本、同样的硬件，把延迟压到 1 s 以内，我们就赢了。
写完这篇笔记，我把线上服务全部切到 ONNX，机器从 8 张卡缩到 3 张卡，电费一个月省了 900 块。ChatTTS 还是那个 ChatTTS，只是换了个“壳”，就能跑得又快又省——技术债早点还，睡觉也更香。
