1. 赛题背景与方案整体思路拆解1.1 这个比赛到底在比什么先说说这个比赛的定位。生成式AI模型优化赛核心考察的不是谁模型训得好而是谁能在给定硬件条件下把已有模型的推理性能压榨到极致。说白了模型精度是主办方给的你要做的是让它在目标设备上跑得更快、更省显存、延迟更低同时尽量不掉点。这类比赛一般会给出一个基线模型通常是PyTorch格式指定目标硬件平台比如T4、A10、消费级显卡等然后让你在限定时间内提交优化后的推理方案。评分维度通常包括推理延迟、吞吐量、显存占用、输出质量保持度这几个方面。我参加的这一届赛题是基于ControlNet的生成式图像模型做推理加速。基线方案在T4上跑一张1080p的图大概要好几秒主办方希望看到的是数量级的提升。最终我们拿了第三名成绩是单张1080p图像生成延迟压到了原来的三分之一左右显存占用降低了约40%。1.2 为什么选这条技术路线拿到赛题之后我先把可能的优化方向列了一遍模型剪枝与量化直接减小模型体积和计算量推理引擎替换用TensorRT或ONNX Runtime替代原生PyTorch推理算子融合与图优化减少kernel launch开销和内存搬运调度策略优化多路并发时的batch调度和显存管理ControlNet条件分支的独立优化ControlNet本身有大量条件注入操作这部分有独立优化空间最终我选的主线是PyTorch → ONNX → TensorRT这条链路配合INT8量化和ControlNet分支的独立处理。原因很直接TensorRT在NVIDIA GPU上的推理性能是经过验证的FP16和INT8的加速比在T4这种卡上非常明显。ONNX作为中间格式既能做图优化又方便做量化校准。为什么不直接用TorchScript或者torch.compile实测下来在ControlNet这种包含大量动态shape和条件分支的模型上TorchScript的兼容性问题很多torch.compile在T4上的收益也不如TensorRT稳定。ONNX虽然转换过程有坑但一旦转成功后续的优化空间最大。1.3 整体方案架构整个方案分四层第一层是模型导出层。把PyTorch的ControlNetUNetVAE拆开分别导出成ONNX。为什么要拆因为ControlNet的条件注入是在UNet的每个down block和mid block上做的如果整体导出ONNX图会非常复杂TensorRT解析时容易出问题。拆开之后ControlNet单独跑一次得到条件特征UNet在推理时直接接收这些特征图结构清晰很多。第二层是量化校准层。用ONNX Runtime的量化工具做INT8校准校准集从训练集里随机采样几百张图。这里的关键是校准集的分布要覆盖实际推理时的输入分布否则量化后的精度损失会很大。第三层是TensorRT引擎构建层。把量化后的ONNX喂给TensorRT的builder设置好优化profile主要是batch size和分辨率范围生成plan文件。这一步的耗时很长但生成一次之后可以反复用。第四层是推理调度层。用TensorRT的Python API加载engine管理显存分配和流式执行。多路并发的时候用CUDA stream做异步调度避免串行等待。注意TensorRT engine是和硬件绑定的换卡必须重新build。比赛提交时如果目标硬件不确定建议同时准备ONNX和engine两套方案。2. 核心细节解析与实操要点2.1 PyTorch转ONNX的坑与解法PyTorch转ONNX看起来就一行torch.onnx.export的事但ControlNet这种模型转起来问题很多。我踩过的坑主要有三个第一个坑是动态shape。ControlNet的输入包括latent、timestep、condition image、text embedding其中latent的H和W是动态的。导出时必须用dynamic_axes参数指定哪些维度是动态的否则ONNX图会把shape写死后续换分辨率就报错。torch.onnx.export( model, (latent, timestep, cond, text_emb), controlnet.onnx, input_names[latent, timestep, cond, text_emb], output_names[out], dynamic_axes{ latent: {2: h, 3: w}, cond: {2: h, 3: w}, out: {2: h, 3: w} }, opset_version14 )opset版本选14而不是最新的17是因为TensorRT对opset 14的支持最稳定。我试过opset 17有些算子TRT解析不了会fallback到ONNX Runtime性能直接掉一半。第二个坑是attention模块的导出。ControlNet里的attention用的是PyTorch的scaled_dot_product_attention这个算子在ONNX里没有直接对应。我的做法是把它替换成手动实现的attention用matmul softmax matmul三步走这样ONNX能正确导出TensorRT也能识别成融合算子。第三个坑是GroupNorm的精度。PyTorch的GroupNorm在导出ONNX时如果opset低于14会被拆成多个小算子TRT优化时反而变慢。opset 14之后GroupNorm有原生支持这个问题就解决了。2.2 INT8量化的校准策略INT8量化是这次方案里收益最大的一步也是最容易掉点的一步。ONNX Runtime的量化工具支持静态量化和动态量化我选的是静态量化因为静态量化在推理时不需要额外计算量化参数速度更快。校准集的构建有几个讲究数量官方建议100-500张我实测下来200张左右就够了再多收益递减分布要覆盖不同的condition类型canny、depth、pose等不能只用一种预处理校准集的预处理必须和推理时完全一致包括归一化参数、resize方式from onnxruntime.quantization import quantize_static, CalibrationDataReader class ControlNetCalibReader(CalibrationDataReader): def __init__(self, calib_data): self.data calib_data self.idx 0 def get_next(self): if self.idx len(self.data): return None batch self.data[self.idx] self.idx 1 return { latent: batch[latent].numpy(), timestep: batch[timestep].numpy(), cond: batch[cond].numpy(), text_emb: batch[text_emb].numpy() } quantize_static( model_inputcontrolnet.onnx, model_outputcontrolnet_int8.onnx, calibration_data_readerControlNetCalibReader(calib_data), quant_formatQuantFormat.QDQ, per_channelTrue, weight_typeQuantType.QInt8 )per_channelTrue这个参数很关键。默认的per_tensor量化对weight的精度损失比较大per_channel按通道量化能明显改善。quant_format选QDQ而不是QOperator是因为QDQ格式在TensorRT里的兼容性更好。2.3 TensorRT引擎构建的参数选择TensorRT build engine的时候有几个参数直接决定最终性能精度模式trt.BuilderFlag.FP16和trt.BuilderFlag.INT8可以同时开。开了INT8之后TRT会自动选择哪些层用INT8、哪些层用FP16。实测下来ControlNet的卷积层大部分能走INT8attention部分TRT会保守地选FP16。优化profile必须设置min/opt/max三个shape。min设成1x4x64x64opt设成1x4x96x96对应768x768输出max设成1x4x128x128对应1024x1024输出。opt shape决定了TRT的kernel选择策略设成最常用的分辨率收益最大。profile builder.create_optimization_profile() profile.set_shape(latent, min(1, 4, 64, 64), opt(1, 4, 96, 96), max(1, 4, 128, 128)) config.add_optimization_profile(profile)workspace大小设成4GB。太小了TRT会放弃一些需要额外显存的优化策略太大了浪费。T4总共16GB显存4GB workspace加上模型本身和中间激活刚好够用。builder优化等级builder_config.builder_optimization_level 5这是最高等级TRT会花更多时间搜索最优kernel但build时间会变长。我build一次大概花了20分钟可以接受。2.4 ControlNet条件分支的独立处理ControlNet的推理流程是condition image先过ControlNet得到一组条件特征然后这些特征在UNet的每个block里做加法。如果整体推理每次都要重新跑ControlNet但实际上如果condition image不变ControlNet的输出是可以缓存的。我的做法是把ControlNet单独导出成一个engineUNet单独一个engine。推理时先跑ControlNet把输出特征存到显存里然后UNet推理时直接读这些特征。这样在多步采样比如20步DDIM的场景下ControlNet只需要跑一次省掉了19次重复计算。这个优化带来的收益非常明显。原本20步采样里ControlNet占了大约30%的计算量缓存之后这部分直接省掉了整体延迟降低了约25%。提示缓存ControlNet输出会额外占用显存大概几百MB。如果显存紧张可以只缓存部分block的特征或者用FP16存储缓存。3. 实操过程与核心环节实现3.1 环境搭建与版本匹配环境这块我踩了不少坑先把最终可用的版本组合列出来组件版本说明Ubuntu20.0422.04也可以但20.04的CUDA兼容性更好CUDA11.8T4支持的最高稳定版本cuDNN8.9和CUDA 11.8配套PyTorch2.0.12.1在ONNX导出时有bugONNX1.141.15的opset变化较大ONNX Runtime1.16GPU版本用于量化校准TensorRT8.68.5也可以8.6的INT8支持更好Python3.103.11有些包还没适配PyTorch和CUDA的版本对应关系一定要查清楚。我一开始用了PyTorch 2.1 CUDA 11.8导出ONNX时attention模块报错换成2.0.1就好了。ONNX Runtime的GPU版本要和CUDA版本匹配否则量化校准时会fallback到CPU速度慢十倍。安装TensorRT的时候建议用tar包安装而不是pip。pip安装的TensorRT有时候会缺一些组件比如trtexec命令行工具。tar包解压后把lib路径加到LD_LIBRARY_PATH里就行。export LD_LIBRARY_PATH/usr/local/TensorRT-8.6.1.6/lib:$LD_LIBRARY_PATH export PATH/usr/local/TensorRT-8.6.1.6/bin:$PATH3.2 模型拆分与导出实操先把ControlNet和UNet从原始pipeline里拆出来。原始代码里它们是通过diffusers库调用的我需要把它们单独实例化from diffusers import ControlNetModel, UNet2DConditionModel controlnet ControlNetModel.from_pretrained( lllyasviel/control_v11p_sd15_canny, torch_dtypetorch.float16 ).cuda().eval() unet UNet2DConditionModel.from_pretrained( runwayml/stable-diffusion-v1-5, subfolderunet, torch_dtypetorch.float16 ).cuda().eval()导出ControlNet的时候输入要构造好。ControlNet的forward签名是(sample, timestep, encoder_hidden_states, controlnet_cond, return_dict)其中sample是latentcontrolnet_cond是condition image经过VAE编码后的latent。dummy_latent torch.randn(1, 4, 96, 96).half().cuda() dummy_timestep torch.tensor([1]).cuda() dummy_text torch.randn(1, 77, 768).half().cuda() dummy_cond torch.randn(1, 4, 96, 96).half().cuda() torch.onnx.export( controlnet, (dummy_latent, dummy_timestep, dummy_text, dummy_cond), controlnet.onnx, input_names[latent, timestep, text_emb, cond], output_names[down_samples, mid_sample], dynamic_axes{ latent: {2: h, 3: w}, cond: {2: h, 3: w}, down_samples: {2: h, 3: w}, mid_sample: {2: h, 3: w} }, opset_version14 )注意ControlNet的输出是多个down block的特征和一个mid block的特征这些都要作为UNet的额外输入。UNet的导出类似但要多接收这些条件特征作为输入。3.3 量化校准的完整流程量化校准分三步准备校准数据、跑校准、验证精度。准备校准数据的时候我从COCO数据集里随机抽了200张图每张图用Canny算子提取边缘作为condition然后用CLIP text encoder编码一个固定的prompt比如a photo作为text embedding。latent用VAE encoder从原图编码得到。calib_data [] for img in calib_images: cond canny_edge(img) cond_latent vae.encode(cond).latent_dist.sample() img_latent vae.encode(img).latent_dist.sample() text_emb clip.encode(a photo) calib_data.append({ latent: img_latent, timestep: torch.tensor([1]), cond: cond_latent, text_emb: text_emb })跑完量化之后一定要验证精度。我的验证方法是用同一组输入分别跑FP16和INT8的ONNX模型比较输出的余弦相似度。如果相似度低于0.98说明量化掉点严重需要调整校准集或者改用per_channel量化。实测下来ControlNet的INT8量化后余弦相似度在0.985左右UNet在0.99左右生成图像肉眼几乎看不出差异。3.4 TensorRT引擎构建与推理封装引擎构建用Python API比trtexec命令行更灵活因为可以在代码里控制profile和flag。import tensorrt as trt logger trt.Logger(trt.Logger.WARNING) builder trt.Builder(logger) network builder.create_network( 1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) ) parser trt.OnnxParser(network, logger) with open(controlnet_int8.onnx, rb) as f: parser.parse(f.read()) config builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 4 30) config.set_flag(trt.BuilderFlag.FP16) config.set_flag(trt.BuilderFlag.INT8) profile builder.create_optimization_profile() profile.set_shape(latent, (1,4,64,64), (1,4,96,96), (1,4,128,128)) profile.set_shape(cond, (1,4,64,64), (1,4,96,96), (1,4,128,128)) config.add_optimization_profile(profile) engine builder.build_engine(network, config) with open(controlnet.engine, wb) as f: f.write(engine.serialize())推理封装的时候用execute_async_v3配合CUDA stream做异步执行。多路并发的时候每个请求分配一个streamTRT会自动做kernel级别的并行。context engine.create_execution_context() context.set_input_shape(latent, (1,4,96,96)) context.set_input_shape(cond, (1,4,96,96)) # 绑定显存 context.set_tensor_address(latent, latent_ptr) context.set_tensor_address(cond, cond_ptr) # ... 其他输入输出 context.execute_async_v3(stream_handle)3.5 多路并发调度策略比赛评分里有吞吐量这一项所以多路并发是必须做的。T4有16GB显存单路1080p推理大概占4GB理论上可以跑3-4路。但实际上因为中间激活和workspace的开销跑3路比较稳。调度策略上我用了一个简单的生产者-消费者模型主线程负责接收请求和预处理然后放到队列里多个worker线程各自持有一个TRT context从队列里取任务执行。每个worker绑定一个CUDA stream避免stream之间的同步开销。注意TRT的context不是线程安全的每个线程必须用自己的context。但engine可以共享engine是只读的。实测下来3路并发时单路延迟从原来的1.2秒增加到1.5秒左右但总吞吐量从0.83张/秒提升到2张/秒性价比很高。4. 常见问题与排查技巧实录4.1 ONNX导出报错速查报错信息原因解法Unsupported operator: aten::scaled_dot_product_attentionPyTorch 2.1的新算子ONNX不支持降级到PyTorch 2.0.1或手动替换attention实现Dynamic shape not supported for output输出没有指定dynamic_axes给所有动态输出加上dynamic_axesGroupNorm decomposition failedopset版本太低升级到opset 14RuntimeError: expected scalar type Half but found Float模型和输入dtype不一致统一用.half()或.float()4.2 TensorRT build失败的排查思路TRT build失败最常见的原因是ONNX图里有TRT不支持的算子。排查方法是先用trtexec --onnxmodel.onnx --verbose跑一遍看哪个层报错。如果是不支持的算子有两个解法一是用ONNX Runtime的onnxsim做图简化把一些冗余算子消掉二是用TRT的plugin机制自己实现但这个成本很高比赛时间紧的话不建议。另一个常见问题是显存不足。build的时候TRT会尝试各种kernel组合显存峰值可能比推理时高很多。如果build失败报OOM把workspace调小或者把max shape调小。4.3 量化掉点严重的处理量化掉点严重通常有三个原因校准集分布不对。比如校准集全是canny condition但实际推理时有depth和pose那量化参数就会偏。解法是校准集要覆盖所有condition类型。敏感层没有跳过。有些层对量化特别敏感比如attention的softmax之前的层。ONNX Runtime的量化工具支持用nodes_to_exclude跳过指定层。我的做法是先全量量化然后逐层排查把掉点最严重的几层排除。per_tensor量化精度不够。改成per_channel量化weight的精度损失会小很多。如果还不行可以试试混合量化weight用INT8activation用FP16。4.4 推理时的显存管理技巧TRT推理时显存管理有几个实用技巧预分配显存池不要在每次推理时malloc/free而是启动时一次性分配好所有输入输出buffer推理时复用用cudaMallocAsyncCUDA 11.2支持的异步显存分配比传统的cudaMalloc快很多及时释放中间激活TRT的context在执行完一次推理后中间激活的显存不会自动释放需要手动调用context.reset()或者复用同一块显存我在实际使用中发现预分配显存池能减少大约15%的推理延迟因为省掉了显存分配的开销。4.5 精度与速度的平衡经验最后分享几个精度和速度平衡的经验INT8量化在ControlNet上收益最大大约有1.8倍加速精度损失在可接受范围内。UNet的INT8收益小一些大约1.4倍因为UNet的attention部分TRT保守地选了FP16。FP16是必开的相比FP32有接近2倍的加速精度损失几乎为零。如果时间紧只做FP16也能拿到不错的成绩。ControlNet输出缓存是性价比最高的优化实现简单收益明显。但要注意缓存的显存开销如果显存紧张可以只缓存mid block的特征。多路并发的收益在T4上比较明显但超过3路之后收益递减因为显存带宽成为瓶颈。建议根据实际显存占用情况调整并发数不要盲目追求高并发。踩过几次坑之后我的体会是优化方案不要一次全上要逐步叠加每加一个优化就测一次精度和速度确保每一步都是正收益。有时候两个优化单独做都有效但叠加在一起反而因为资源竞争导致性能下降。