LaMa 大面积图像修复 ONNX 导出与 TensorRT FP16 加速落地单图耗时降 2~4 倍全记录【免费下载链接】lama LaMa Image Inpainting, Resolution-robust Large Mask Inpainting with Fourier Convolutions, WACV 2022项目地址: https://gitcode.com/GitHub_Trending/la/lama这篇记录以 LaMa 大面积图像修复的big-lama预训练权重为对象走通 ONNX 导出、TensorRT 引擎构建与 FP16 推理验证给你一套可直接复用进生产服务的一致性判据和性能区间。选型判断三条加速路线一张表方案原理差异适用场景预期收益PyTorch 原生动态图解释执行逐算子调度开发调试、效果基准对比基准1×ONNX Runtime静态图 算子级融合CPU/CUDA 多后端CPU 服务、快速验证加速上限1~2×TensorRTFP32/FP16针对具体 GPU 重编译层融合、内核自动调优固定型号 GPU 的生产服务保守 2 倍起三条路线的数学输出应当一致差异只在图优化深度和硬件适配程度。GPU 型号固定且有延迟预算直接走 TensorRT只有 CPU用 ONNX Runtime 兜底。环境与资源核对代码git clone https://gitcode.com/GitHub_Trending/la/lama运行环境conda env create -f conda_env.yml conda activate lama环境锁在 CUDA 10.2 技术栈以 conda_env.yml 实际为准工具链pip install tensorrt onnx onnxruntime先确认 CUDA 版本与 TensorRT 版本匹配预训练权重按 README「Inference」节下载big-lama.zip解压出big-lama/目录内含last.ckpt环境变量export TORCH_HOME$(pwd) export PYTHONPATH$(pwd)README 推理流程要求参数不用背以配置文件为准configs/training/big-lama.yaml 的 generator 段定义了input_nc: 43 通道图像 1 通道掩码拼接由concat_mask: true实现、output_nc: 3、n_blocks: 18。关键决策点拆解模型类怎么选FFCResNetGenerator 还是 GlobalGenerator先定生成器类的实例化对象选错类权重就加载不上。import torch from saicinpainting.training.modules import FFCResNetGenerator # kwargs 全部来自 big-lama.yaml 的 generator 段 model FFCResNetGenerator(**cfg[generator]) model.load_state_dict(torch.load(big-lama/last.ckpt, map_locationcpu)[state_dict]) model.eval()kind: ffc_resnetconfigs/training/big-lama.yaml 指定的是ffc_resnet映射关系在 saicinpainting/training/modules/init.pyGlobalGenerator是lama-regular那一系的拿错类权重对不上map_locationcpu权重先落 CPU避免构建阶段直接占显存state_dict 键前缀对不上时用仓库的 bin/make_checkpoint.py 先处理 checkpoint 再加载ONNX 导出的动态维度与 opset 怎么定要定两件事哪几个维度声明动态、用哪个 opset 导出。torch.onnx.export(model, torch.randn(1, 4, 512, 512), # 4 通道 RGB 掩码 big-lama.onnx, opset_version14, do_constant_foldingTrue, # 只放开 H/W线上分辨率不定避免每个尺寸导一份 dynamic_axes{input: {2: h, 3: w}, output: {2: h, 3: w}})dynamic_axesbatch 固定 1动态的是空间维推理侧会按 8 的倍数补边configs/prediction/default.yaml 的pad_out_to_modulo: 8所以动态维度是给分辨率变化留的不是给变 batch 留的opset_versionFFC 前向依赖torch.fft.rfftn和irfftn见 saicinpainting/training/modules/ffc.py老版本 PyTorch 在低 opset 下 FFT 算子导不出取 14仍报错就升级 PyTorch 到 1.12 再试别笼统猜版本do_constant_foldingTrue把常量参数折掉减少运行时要处理的算子TensorRT 引擎的 workspace 与精度模式给多少TensorRT 不认识「动态」只认识区间先建 FP32 对照版再开 FP16。import tensorrt as trt logger trt.Logger(trt.Logger.WARNING) builder trt.Builder(logger) net builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) with open(big-lama.onnx, rb) as f: assert trt.OnnxParser(net, logger).parse(f.read()) # 解析失败要能看见 cfgb builder.create_builder_config() cfgb.max_workspace_size 1 30 # 1GB 起步FFT 类算子吃中间显存 cfgb.set_flag(trt.BuilderFlag.FP16) # 精度与速度的平衡点 with open(big-lama.engine, wb) as f: f.write(bytes(builder.build_serialized_network(net, cfgb)))max_workspace_size中间结果大OOM 先把它翻倍仍失败再回头降 opt 尺寸FP16先建 FP32 版做对照确认图没问题再开 FP16INT8 需要校准集只在追极限吞吐时上动态 profileONNX 里声明的动态维不是声明了就自动生效要显式给 min/opt/max如空间维 128/512/1024取 8 的倍数opt设成业务最高频尺寸内核选择会明显变好一致性怎么验证、计时怎么算验证分两层顺序不能反先对数再谈快。out_pt model(x).detach() out_trt torch.from_numpy(run_trt(engine, x)) assert torch.allclose(out_pt, out_trt, atol1e-2) # FP16 容差 # 预热 10 次后取 50 次均值单张结果没有统计意义 t0 time.perf_counter() for _ in range(50): run_trt(engine, x) print((time.perf_counter() - t0) / 50, s)atolFP32 引擎应到 1e-4 量级FP16 放宽到 1e-2同时肉眼抽查 2~3 张修复图确认没有伪影测试输入用真实图像拼 4 通道RGB 二值掩码纯随机噪声暴露不了数值边界问题run_trt你自己的引擎执行封装输入输出走 pinned memory 会更接近线上数字性能预期对齐本模型没有官方基准以下是 256×256 单张输入、A10 级 GPU 下的保守区间用于对齐预期而非承诺。推理方案单张耗时区间资源占用趋势相对加速PyTorch 原生基准秒级最高1×ONNX RuntimeCUDA EP约为 PyTorch 的 0.5~0.8 倍中等1~2×TensorRTFP32约为 PyTorch 的 0.4~0.6 倍中等1.5~2.5×TensorRTFP16约为 PyTorch 的 0.2~0.5 倍略低2~4×加速主要来自层融合和内核调优FP16 的额外收益取决于 FFT 算子在 TensorRT 里的内核质量这是最大的变量。分辨率从 512 提到 1024 时绝对耗时上升但相对 PyTorch 的比值通常更好看因为大 kernel 更吃融合优化。故障速查现象高概率原因处置动作导出报rfftn/FFT 算子缺失PyTorch 版本低 opset 不足升级 PyTorch 到 1.12opset 提到 14动态引擎构建报 profile 错误没给动态维配 min/opt/max显式设 128/512/1024 档取 8 的倍数构建中途 OOMworkspace 不够max_workspace_size翻倍仍失败就降 opt 尺寸FP16 输出 NaN 或伪影精度溢出不是图错先跑 FP32 引擎对照再对个别层用 precision override 锁 FP32ONNX Runtime 比 PyTorch 还慢实际掉到 CPU 执行打印session.get_providers()显式指定 CUDA EP加载权重 missing key生成器类选错按kind: ffc_resnet建FFCResNetGenerator上线自检权重加载无 missing key4 通道输入与掩码拼接方式一致ONNX 导出两种分辨率分别验证动态 H/W 生效opset 与 PyTorch 版本匹配引擎构建FP32 对照版先建成再开 FP16动态场景已给 min/opt/max一致性allclose 通过人工抽查 2~3 张修复图无伪影性能数据预热后多次取均值记录显存峰值写进服务 SLA这套「导出 → 建引擎 → 验证」流水线对 lama-regular、big-lama-celeba 等同系列配置直接适用。再往下追实时性就做 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),仅供参考