LaMa 大面积图像修复生产部署完整指南ONNX 导出到 TensorRT FP16 加速的三道门槛【免费下载链接】lama LaMa Image Inpainting, Resolution-robust Large Mask Inpainting with Fourier Convolutions, WACV 2022项目地址: https://gitcode.com/GitHub_Trending/la/lamabig-lama 是本仓库的大面积掩码图像修复模型原生 PyTorch 推理单张以秒计而训练栈的 autograd 与 hook 在线上全是纯开销。这篇文章围绕 big-lama走通「提速、验稳、上线」三道门槛ONNX 导出、TensorRT 引擎构建、服务化封装。先看选型表三条推理路线的适用边界结论先行GPU 型号固定且直接扛流量直接上 TensorRTGPU 型号未定或要快速量出加速上限ONNX Runtime 做过渡PyTorch 原生只留作开发基准与数值对照。路线优化发生在哪层部署约束相对 PyTorch 原生的保守预期PyTorch 原生动态图解释执行无融合无基准 1×ONNX Runtime静态图算子融合多后端需 onnxruntime 匹配的执行提供器约 1~2×CUDA EPTensorRT按具体 GPU 重编译层融合 内核调优 FP16引擎与 GPU 型号绑定FP32 约 1.5~2.5×FP16 约 2~4×一句话解读三条路线输出应当数值一致差异只在图优化深度与硬件适配程度这也是后文「先对数再谈快」的由来。先跑通 PyTorch 基准用官方 predict 流程锁定参照结果动手导出前先把官方推理路径跑一遍手里有随时可比对的参照输出。git clone https://gitcode.com/GitHub_Trending/la/lama cd lama export TORCH_HOME$(pwd) export PYTHONPATH$(pwd) python3 bin/predict.py model.path$(pwd)/big-lama \ indir$(pwd)/LaMa_test_images outdir$(pwd)/output \ model.checkpointlast.ckptmodel.checkpoint默认best.ckpt指定加载哪个 checkpoint。解压big-lama.zip得到big-lama/last.ckpt所以必须显式传model.checkpointlast.ckpt传错文件名会整图修复失败或黑屏。pad_out_to_modulo值为 8见 configs/prediction/default.yaml官方流程把输入补到 8 的倍数再进模型。不补的话输出与输入尺寸对不齐服务侧还得加裁剪逻辑。refine: False迭代细化器默认关闭。它在同一配置里是n_iters: 15、max_scales: 3延迟远高于主模型线上要开就单独测一版别混进主链路。GPU 路径可直接走 docker/2_predict_with_gpu.sh。基准跑通后记录一张真实图像的修复输出张量与单张耗时作为后续所有加速结果的对照基线。怎么构建 big-lama 的 ONNX 导出脚本输入通道与 opset先确认模型类。configs/training/big-lama.yaml 里 generator 段是kind: ffc_resnet由 saicinpainting/training/modules/init.py 的make_generator分发到FFCResNetGenerator定义在 saicinpainting/training/modules/ffc.py。别用GlobalGenerator那是lama-regular的结构类选错权重直接加载不上。model FFCResNetGenerator(**OmegaConf.to_container(cfg[generator])) model.load_state_dict(torch.load(big-lama/last.ckpt, map_locationcpu)[state_dict]) model.eval() torch.onnx.export(model, torch.randn(1, 4, 512, 512), big-lama.onnx, opset_version14, do_constant_foldingTrue, dynamic_axes{input: {2: h, 3: w}, output: {2: h, 3: w}})input_nc: 4同配置文件里concat_mask: truegenerator 输入是 3 通道图像 1 通道掩码拼接。传 3 通道导出不会报错但 ONNX 推理时通道数直接对不上。opset_version: 14ffc.py 的FourierUnit前向依赖torch.fft.rfftn低 opset 落不了地。导出失败先看torch.onnx报的具体算子名环境锁在 PyTorch 1.8本仓库 conda_env.yml 的版本时优先把导出侧 PyTorch 升到 1.12 再试。dynamic_axes把 H/W 声明为动态。不做的话每个分辨率导一个模型、每档尺寸建一次引擎运维上收不住。官方 bin/to_jit.py 是另一条路用JITWrapper把 image、mask 分开喂给训练模型再 trace拼接在内部完成。你直接导出FFCResNetGenerator的话拼接得放在服务侧预处理好。怎么建 big-lama 的 FP16 TensorRT 引擎workspace 与动态 profile 的 3 个参数import tensorrt as trt logger trt.Logger(trt.Logger.WARNING) builder trt.Builder(logger) net builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) trt.OnnxParser(net, logger).parse(open(big-lama.onnx, rb).read()) cfg builder.create_builder_config() cfg.max_workspace_size 1 30 prof builder.create_optimization_profile() prof.set_shape(input, (1,4,128,128), (1,4,512,512), (1,4,1024,1024)) cfg.add_optimization_profile(prof) cfg.set_flag(trt.BuilderFlag.FP16) open(big-lama.engine, wb).write(builder.build_serialized_network(net, cfg))max_workspace_size: 1GB1 30FFT 与 18 个残差块的中间结果吃显存。给小了构建失败或放弃融合选择显存不够先翻倍再排查别的。FP16 标志先用 FP32 建对照版验证一致再开 FP16。只有一版引擎时FP16 出 NaN 你分不清是精度问题还是图错误确认是精度问题后可对个别层用 layer precision override 单独锁回 FP32。min/opt/max profileTensorRT 不认识「动态」只认区间。空间维给 128/512/1024 三档opt设为业务最高频尺寸内核选择会明显更优不传要么构建报错要么只能跑死尺寸。FFC 基于 FFT空间维尽量保持 2 的幂对齐。怎么验证引擎与 PyTorch 的数值一致性容差、预热与抽查顺序不能反先确认输出一致再谈速度。out_pt model(x).detach() out_trt run_trt(engine, x) assert torch.allclose(out_pt.float(), out_trt, atol1e-4) # FP32 引擎 for _ in range(10): # 预热 run_trt(engine, x) ts [] for _ in range(50): t0 time.perf_counter() run_trt(engine, x) ts.append(time.perf_counter() - t0) print(np.mean(ts), s)atolFP32 引擎 1e-4 / FP16 引擎 1e-2一致性的量级门槛。FP32 引擎到不了 1e-4先怀疑图错误opset、动态轴别急着归咎精度FP16 落在 1e-2 内还要看肉眼抽查。预热 10 次、取 50 次以上均值首次调用含内核加载与显存分配单次结果没有统计意义。抽查样本 2~3 张测试输入用仓库评估流程的真实图像构造 4 通道张量而不是随机噪声——只有真实图像能暴露数值边界问题再目检修复区没有新伪影。怎么把推理封装成服务预对齐、batch 与性能区间 服务侧预处理与官方流程对齐补 8 的倍数、拼接 4 通道。ph, pw (8 - h % 8) % 8, (8 - w % 8) % 8 x torch.cat([F.pad(img, (0, pw, 0, ph)), F.pad(mask, (0, pw, 0, ph))], dim0)掩码命名与仓库测试集约定一致[image]_maskXXX.png同目录回归测试可直接拿 saicinpainting/evaluation/masks/countless/images/segmentation.png 这类真实样例当固定输入。batch从 1 起步batch1 先固化全链路再验 batch2 的一致性与延迟通过才扩。直接上大 batchFFT 类算子的显存占用非线性上涨爆了也不好归因。并发单引擎实例不并发每个 worker 一个引擎实例或用请求队列串行化执行。共享引擎会偶发挂起。⚡ 性能预期512×512 单张、A10 级 GPU保守区间PyTorch 原生秒级ONNX RuntimeCUDA EP约其 0.5~0.8 倍TensorRT FP32 约 0.4~0.6 倍FP16 约 0.2~0.5 倍。分辨率提到 1024 后绝对耗时上升但 TensorRT 的相对优势通常更明显大 kernel 更吃融合优化。这套「导出 → 建引擎 → 一致性验证」流水线对 configs/training/ 下lama-regular、big-lama-celeba等其他配置同样适用换个 kind 重新导出即可。下一步可构建 INT8 校准集榨取最后的吞吐或用 configs/training/trainer/ 里的 benchmark 配置做训练侧交叉对照。【免费下载链接】lama LaMa Image Inpainting, Resolution-robust Large Mask Inpainting with Fourier Convolutions, WACV 2022项目地址: https://gitcode.com/GitHub_Trending/la/lama创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考