当前位置：首页 > news >正文

LaMa图像修复模型：ONNX导出与TensorRT加速实践指南

news 2026/6/29 10:33:15

LaMa图像修复模型：ONNX导出与TensorRT加速实践指南

【免费下载链接】lama项目地址: https://gitcode.com/gh_mirrors/lam/lama

在当今数字图像处理领域，LaMa（Large Mask Inpainting）模型凭借其出色的高分辨率图像修复能力，已成为业界关注的焦点。该模型采用傅里叶卷积技术，能够在保持图像质量的同时，有效处理大面积缺失区域。本文将深入探讨如何通过ONNX导出和TensorRT加速技术，显著提升LaMa模型的推理性能。

LaMa模型架构解析

LaMa模型的核心架构基于改进的ResNet结构，包含下采样、瓶颈层和上采样三个主要部分。根据配置文件分析，big-lama模型采用以下关键参数：

输入通道数：4（3通道图像 + 1通道掩码）
输出通道数：3（修复后的RGB图像）
基础特征图数量：64
下采样层数：3
残差块数量：18

模型的关键创新在于引入了FFC（Fourier Filter Convolution）块，这些块能够在频域和空域同时处理特征，从而更好地捕捉图像的全局结构信息。

环境准备与模型获取

项目环境搭建

首先需要获取项目代码并配置运行环境：

git clone https://gitcode.com/gh_mirrors/lam/lama cd lama conda env create -f conda_env.yml conda activate lama

预训练模型下载

获取官方预训练的big-lama模型权重：

wget https://huggingface.co/smartywu/big-lama/resolve/main/big-lama.zip unzip big-lama.zip

下载完成后，模型文件将保存在项目目录的big-lama文件夹中。

ONNX模型导出实战

模型配置加载

在导出ONNX模型之前，需要正确加载模型配置和权重。big-lama模型的配置文件位于configs/training/big-lama.yaml，其中定义了生成器的具体参数。

导出代码实现

创建专门的导出脚本来完成ONNX格式转换：

import torch import yaml from saicinpainting.training.modules.pix2pixhd import GlobalGenerator # 加载模型配置 with open("configs/training/big-lama.yaml", 'r') as f: config = yaml.safe_load(f) # 构建生成器模型 generator_config = config['generator'] model = GlobalGenerator( input_nc=generator_config['input_nc'], output_nc=generator_config['output_nc'], ngf=generator_config['ngf'], n_downsampling=generator_config['n_downsampling'], n_blocks=generator_config['n_blocks'], padding_type=generator_config['padding_type'], ffc_positions=generator_config.get('ffc_positions', None), ffc_kwargs=generator_config.get('ffc_kwargs', {}) ) # 加载预训练权重 checkpoint = torch.load("big-lama/last.ckpt", map_location='cpu') model.load_state_dict(checkpoint['state_dict'], strict=False) model.eval()

动态输入尺寸支持

为了适应不同分辨率的输入图像，需要在导出时配置动态尺寸：

# 创建示例输入 dummy_input = torch.randn(1, 4, 512, 512) # 导出ONNX模型 torch.onnx.export( model, dummy_input, "big-lama.onnx", opset_version=12, do_constant_folding=True, input_names=['input'], output_names=['output'], dynamic_axes={ 'input': {2: 'height', 3: 'width'}, 'output': {2: 'height', 3: 'width'} } )

TensorRT加速技术详解

TensorRT环境配置

确保系统已安装正确版本的TensorRT：

pip install tensorrt

引擎构建过程

将ONNX模型转换为优化的TensorRT引擎：

import tensorrt as trt # 创建日志记录器 TRT_LOGGER = trt.Logger(trt.Logger.WARNING) # 构建引擎 builder = trt.Builder(TRT_LOGGER) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser = trt.OnnxParser(network, TRT_LOGGER) # 解析ONNX模型 with open("big-lama-sim.onnx", 'rb') as model_file: parser.parse(model_file.read()) # 配置构建参数 config = builder.create_builder_config() config.max_workspace_size = 1 << 30 # 1GB config.set_flag(trt.BuilderFlag.FP16) # 启用FP16精度 # 序列化并保存引擎 serialized_engine = builder.build_serialized_network(network, config) with open("big-lama.engine", "wb") as f: f.write(serialized_engine)