最近在折腾语音合成项目团队想引入清华开源的ChatTTS。说实话刚开始用原生PyTorch部署时踩了不少坑环境配置复杂、长文本合成动不动就内存溢出、推理速度也上不去。经过一番摸索我们最终基于Docker和ONNX Runtime搞定了生产级部署性能提升明显。这里把整个实战过程包括踩过的坑和优化经验整理成笔记分享给大家。1. 部署路上的那些“坑”从环境到性能刚开始部署ChatTTS时相信不少朋友都遇到过类似的问题。我简单总结了几类典型痛点环境配置复杂ChatTTS依赖特定版本的PyTorch、TorchAudio以及其他一些音频处理库。在本地或服务器上手动配环境经常遇到CUDA版本不匹配、库冲突等问题重现环境非常困难。长文本合成内存溢出OOM这是最头疼的问题之一。当输入文本较长时模型在推理过程中占用的显存会急剧增长尤其是在使用自回归方式生成音频时很容易就把显存撑爆导致程序崩溃。推理速度慢原生PyTorch在CPU上推理速度感人即使在GPU上首次推理也因为图优化等问题不够快难以满足实时或准实时的交互需求。模型加载慢每次启动服务都要重新加载模型尤其是模型文件较大时加载耗时较长影响服务启动速度和弹性伸缩。这些问题直接影响了开发效率和线上服务的稳定性。所以我们的优化目标很明确环境标准化、推理加速、资源可控。2. 技术选型为什么是ONNX Runtime Docker面对上述问题我们评估了几个方案核心是解决推理引擎的问题。方案A坚持原生PyTorch。好处是与训练框架一致兼容性好。但缺点也很明显推理性能并非最优缺乏统一的图优化和算子融合内存管理也更依赖PyTorch自身在部署灵活性上稍逊一筹。方案B转向ONNX Runtime。这是一个高性能推理引擎支持多种硬件后端CPU GPU TensorRT等。它的优势在于高性能内置了大量图优化如算子融合、常量折叠能显著提升推理速度。跨平台一次导出ONNX模型多处部署。资源高效对内存和显存的使用通常更可控特别是与一些特定Provider如CUDA TensorRT结合时。标准化ONNX模型作为中间表示解耦了训练和部署环境。我们做了一个简单的对比测试在同一台配有V100的服务器上合成同一段5秒左右的音频引擎平均延迟 (ms)峰值显存占用 (MB)备注PyTorch (GPU)12003200包含Python解释器开销ONNX Runtime (CUDA)6502800已开启基础优化ONNX Runtime (TensorRT)4502600需要额外转换效果最好结论ONNX Runtime在延迟和资源占用上都有优势。结合Docker解决环境一致性问题这成为了我们的技术组合。3. 核心实现三步走搞定标准化部署3.1 第一步用Docker固化环境为了避免“在我机器上是好的”这种问题第一步就是用Docker把环境打包。我们基于nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04这个官方镜像构建。# Dockerfile FROM nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04 # 设置环境变量防止Python输出缓冲 ENV PYTHONUNBUFFERED1 # 安装系统依赖和Python RUN apt-get update apt-get install -y \ python3.10 \ python3-pip \ rm -rf /var/lib/apt/lists/* # 设置工作目录 WORKDIR /app # 复制依赖文件并安装 COPY requirements.txt . RUN pip3 install --no-cache-dir -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple # 复制应用代码 COPY . . # 启动命令示例 CMD [python3, app.py]requirements.txt里固定了关键版本torch2.1.0 torchaudio2.1.0 onnxruntime-gpu1.16.0 # 以及其他ChatTTS所需的包...这样无论开发、测试还是生产环境都是完全一致的。3.2 第二步将PyTorch模型转换为ONNX格式这是最关键的一步。ChatTTS的推理主要涉及两个模型一个文本编码器和一个声学模型。我们需要分别将它们导出。import torch import ChatTTS # 假设这是你clone的ChatTTS代码 import onnx # 加载原始模型 chat ChatTTS.Chat() chat.load_models() # 假设这个方法加载模型 # 获取文本编码器模型并设置为eval模式 text_encoder chat.text_encoder text_encoder.eval() # 准备示例输入 dummy input # 注意需要根据ChatTTS的实际输入结构来调整 # 假设文本编码器输入是token ids和长度 batch_size 1 seq_length 50 # 示例序列长度 dummy_input_ids torch.randint(0, 1000, (batch_size, seq_length)).long() dummy_lengths torch.tensor([seq_length]).long() # 导出文本编码器到ONNX onnx_text_encoder_path text_encoder.onnx torch.onnx.export( text_encoder, (dummy_input_ids, dummy_lengths), # 模型输入 onnx_text_encoder_path, input_names[input_ids, lengths], # 输入节点名 output_names[text_hidden], # 输出节点名 dynamic_axes{ input_ids: {0: batch_size, 1: seq_len}, # 必须设置dynamic_axes以适应可变长度输入 lengths: {0: batch_size}, text_hidden: {0: batch_size, 1: seq_len} }, opset_version14, # 使用较新的opset以获得更好支持 do_constant_foldingTrue # 常量折叠优化 ) print(f文本编码器已导出至: {onnx_text_encoder_path}) # 类似地导出声学模型 vocoder # 注意声学模型的输入通常是梅尔频谱或类似特征 # dummy_mel torch.randn(batch_size, 80, 200) # 示例梅尔频谱 # ... 导出过程类似需要根据实际模型结构调整关键参数说明dynamic_axes这个非常重要它指定了哪些输入/输出维度是动态的可变的。对于语音合成输入文本长度和输出音频长度都是可变的必须在这里声明否则导出的模型只能处理固定尺寸的输入实用性大打折扣。opset_version建议使用较新的版本如14以获得更多算子的支持和优化。do_constant_folding启用常量折叠可以优化计算图减少运行时操作。3.3 第三步使用ONNX Runtime进行推理模型导出后就可以用ONNX Runtime加载并推断了。import onnxruntime as ort import numpy as np # 配置ONNX Runtime会话选项指定GPU执行 providers [CUDAExecutionProvider] # 使用CUDA Provider sess_options ort.SessionOptions() # 可以设置一些优化选项例如启用图优化级别 sess_options.graph_optimization_level ort.GraphOptimizationLevel.ORT_ENABLE_ALL # 创建推理会话 text_encoder_session ort.InferenceSession(onnx_text_encoder_path, sess_optionssess_options, providersproviders) # 准备输入数据需要转换为numpy array input_ids_np dummy_input_ids.numpy() lengths_np dummy_lengths.numpy() # 运行推理 input_feed { input_ids: input_ids_np, lengths: lengths_np } text_hidden text_encoder_session.run([text_hidden], input_feed)[0] print(文本编码器推理完成输出形状, text_hidden.shape)对于声学模型生成音频推理过程类似。将文本编码器的输出text_hidden作为声学模型的输入最终得到音频波形。4. 性能优化让推理飞起来模型能跑起来只是第一步优化到生产级别还需要一些技巧。4.1 动态批处理Dynamic BatchingONNX Runtime本身对固定批处理支持很好。要实现动态批处理我们需要在服务端手动收集请求凑成一个批次再推理。import threading import queue import time class DynamicBatchInference: def __init__(self, onnx_model_path, max_batch_size4, max_wait_time0.05): self.session ort.InferenceSession(onnx_model_path, providers[CUDAExecutionProvider]) self.max_batch_size max_batch_size self.max_wait_time max_wait_time # 最大等待时间秒用于凑批 self.request_queue queue.Queue() self.result_dict {} self.lock threading.Lock() self.worker_thread threading.Thread(targetself._batch_worker, daemonTrue) self.worker_thread.start() def _batch_worker(self): while True: batch_inputs [] batch_req_ids [] start_time time.time() # 尝试收集一个批次的请求 while len(batch_inputs) self.max_batch_size: try: # 设置超时避免无限等待 req_id, input_data self.request_queue.get(timeoutself.max_wait_time) batch_inputs.append(input_data) batch_req_ids.append(req_id) except queue.Empty: # 超时无论凑够多少都执行一次推理 if batch_inputs: break else: continue # 继续等待 # 如果已经达到最大批次大小立即执行 if len(batch_inputs) self.max_batch_size: break if not batch_inputs: continue # 将列表输入堆叠成批次这里假设输入是相同形状的 # 实际情况中可能需要padding来处理变长序列 batched_input np.stack(batch_inputs, axis0) # 运行批次推理 (假设只有一个输入叫input一个输出叫output) # 实际输入输出名称需根据模型调整 outputs self.session.run([output], {input: batched_input})[0] # 将结果分发给各个请求 with self.lock: for req_id, single_output in zip(batch_req_ids, outputs): self.result_dict[req_id] single_output def infer(self, input_data): 外部调用接口 req_id id(input_data) # 简单生成一个请求ID self.request_queue.put((req_id, input_data)) # 轮询等待结果 while True: with self.lock: if req_id in self.result_dict: result self.result_dict.pop(req_id) return result time.sleep(0.001) # 短暂休眠避免CPU空转 # 使用示例 # batch_inferencer DynamicBatchInference(acoustic_model.onnx, max_batch_size8) # audio batch_inferencer.infer(mel_spec)4.2 显存预分配与释放最佳实践预热Warm-up在服务正式处理请求前先用一个或几个典型的小请求跑一遍推理。这会让ONNX Runtime完成初始化的图优化、内存分配等操作避免第一个真实请求的延迟过高。监控与清理定期监控GPU显存使用情况。对于长时间运行的服务如果发现显存缓慢增长可能由于碎片或某些库的内存管理可以设置一个阈值在达到后重启工作进程如果你的服务是多进程架构。ONNX Runtime本身的内存管理通常比较好但结合Python和其他库时仍需注意。使用ort.SessionOptions配置可以尝试设置session_options.enable_cpu_mem_arena False来禁用CPU内存池如果你的瓶颈在CPU内存或者调整session_options.intra_op_num_threads来控制CPU算子的并行度间接影响内存使用模式。4.3 性能测试数据RTF - Real Time FactorRTF是语音合成中常用的性能指标表示合成一段音频所需时间与音频实际时长的比值。RTF 1 表示快于实时。我们在V100 GPU上测试了不同批处理大小Batch Size下的RTF测试文本平均长度20字Batch Size平均RTF峰值显存 (GB)备注10.452.8单请求延迟最低40.183.5吞吐量优化明显80.124.8接近显存上限延迟略有增加16OOM-显存不足结论Batch Size为4或8时能在延迟和吞吐量之间取得较好平衡。具体数值需根据你的硬件和业务需求调整。5. 避坑指南那些让人头疼的细节5.1 中文标点合成异常问题合成带有中文标点如全角逗号“”、句号“。”的文本时音频可能出现不自然的停顿或音素错误。解决方案文本预处理在将文本送入模型前进行统一的标点规范化。例如将全角标点转换为半角标点或者根据TTS模型训练数据的习惯进行转换。import re def normalize_punctuation(text): # 全角转半角针对英文和数字标点 text text.replace(, ,) # 全角逗号 text text.replace(。, .) # 全角句号 text text.replace(, !) # 全角感叹号 text text.replace(, ?) # 全角问号 text text.replace(, ;) # 全角分号 text text.replace(, :) # 全角冒号 # 也可以使用更通用的映射表 # 注意中文文本中有时全角标点是合适的需根据模型训练数据决定 return text raw_text 你好世界这是一个测试。 processed_text normalize_punctuation(raw_text) print(processed_text) # 输出: 你好,世界!这是一个测试.检查分词器Tokenizer确认ChatTTS使用的分词器是否正确处理了这些标点。有时需要更新分词器的词汇表或处理规则。微调模型高级如果问题严重且你有训练数据可以考虑用包含正确标点处理的数据对模型进行少量微调。5.2 音频卡顿或爆音问题问题生成的音频在播放时出现卡顿、噼啪声或爆音。诊断流程检查采样率确保生成的音频波形采样率如24000Hz与播放设备或后续处理环节期望的采样率一致。如果不一致需要进行重采样。import librosa # 假设原始音频是24000Hz需要转换为16000Hz original_sr 24000 target_sr 16000 audio_resampled librosa.resample(audio_waveform, orig_sroriginal_sr, target_srtarget_sr)检查幅值Clipping音频波形的幅值超过[-1.0, 1.0]范围对于float32格式会导致削波失真产生爆音。在保存或播放前进行限幅。audio_waveform np.clip(audio_waveform, -1.0, 1.0)检查模型输出直接检查模型输出的原始音频波形看看是否存在异常的数值如NaN或Inf。这可能是模型推理过程中的数值不稳定造成的。查看声学模型Vocoder输入检查输入给声学模型的梅尔频谱等特征是否存在异常。特征异常会导致声学模型生成劣质音频。硬件/驱动问题在极少数情况下可能是GPU驱动或CUDA版本存在兼容性问题。尝试更新驱动或使用不同版本的ONNX Runtime/CUDA。6. 延伸思考还能做些什么量化部署为了进一步优化模型大小和推理速度可以尝试量化Quantization。ONNX Runtime支持将FP32模型量化为INT8能在几乎不损失精度的情况下显著提升速度并降低内存占用。可以使用ONNX Runtime提供的量化工具或PyTorch的量化功能导出量化模型。开发RESTful API接口将你的TTS引擎封装成HTTP服务方便其他系统调用。可以使用FastAPI等现代框架快速搭建。from fastapi import FastAPI, HTTPException from pydantic import BaseModel import numpy as np import io from scipy.io import wavfile app FastAPI() class TTSRequest(BaseModel): text: str speaker_id: int 0 # 假设支持多说话人 speed: float 1.0 app.post(/synthesize) async def synthesize_speech(request: TTSRequest): try: # 1. 文本预处理 processed_text preprocess_text(request.text) # 2. 推理调用前面封装好的推理模块 audio_np tts_pipeline.infer(processed_text, request.speaker_id, request.speed) # 3. 将numpy数组转换为WAV字节流 buffer io.BytesIO() wavfile.write(buffer, 24000, audio_np) buffer.seek(0) return StreamingResponse(buffer, media_typeaudio/wav) except Exception as e: raise HTTPException(status_code500, detailstr(e))流式合成Streaming TTS对于非常长的文本可以研究如何将模型拆分成流式模式生成一部分就输出一部分音频减少用户端等待时间。这需要对模型结构有更深的理解和修改。写在最后从零开始部署和优化一个像ChatTTS这样的语音合成模型确实是一个系统工程涉及环境、模型转换、性能优化和问题调试多个方面。通过采用DockerONNX Runtime的方案我们成功地将部署过程标准化并获得了显著的性能提升。希望这篇笔记里提到的具体步骤、代码示例和踩坑经验能帮助你更顺利地完成自己的TTS项目部署。当然每个项目都有其独特性你可能需要根据ChatTTS的具体代码和模型结构调整一些细节。但万变不离其宗掌握环境隔离、模型转换、推理优化和问题排查这些核心思路就能应对大部分挑战。如果遇到新的问题多查查ONNX Runtime文档和社区通常都能找到解决方案。