如何实现Padim模型在NVIDIA GPU上的OpenVINO高帧率推理？

核心问题说明 #

OpenVINO的device="gpu"参数仅适配英特尔系列GPU，对NVIDIA GPU不兼容，这是你设置该参数后构建失败的原因。要在NVIDIA GPU上实现高帧率推理，需改用NVIDIA生态的推理框架，以下是两种可行方案：

方案1：ONNX Runtime + CUDA后端推理 #

你已拥有model.onnx文件，直接使用ONNX Runtime的CUDA后端即可快速实现GPU推理，步骤如下：

1. 安装依赖 #

pip install onnxruntime-gpu

2. 推理代码实现 #

结合anomalib的后处理逻辑（参考metadata.json中的配置），示例代码如下：

import onnxruntime as ort
import numpy as np
from PIL import Image
import json

# 加载模型元数据
with open(metadata_path, 'r') as f:
    metadata = json.load(f)

# 初始化ONNX Runtime会话（优先使用CUDA）
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
session = ort.InferenceSession('model.onnx', providers=providers)

# 获取模型输入输出节点名称
input_name = session.get_inputs()[0].name
output_names = [out.name for out in session.get_outputs()]

# 图像预处理（与训练阶段逻辑一致）
def preprocess(image_path):
    image = Image.open(image_path).convert('RGB')
    img_size = metadata['transform']['image_size']
    image = image.resize((img_size[0], img_size[1]))
    # 归一化并调整维度为(1, C, H, W)
    image = np.array(image).astype(np.float32) / 255.0
    image = np.transpose(image, (2, 0, 1))[np.newaxis, ...]
    return image

# 推理及后处理
def predict(image_path):
    input_tensor = preprocess(image_path)
    outputs = session.run(output_names, {input_name: input_tensor})
    
    # 基于元数据计算异常结果
    anomaly_score = outputs[0]
    threshold = metadata['metrics']['image_threshold']
    pred_label = anomaly_score > threshold
    
    return {
        'anomaly_score': anomaly_score.item(),
        'pred_label': pred_label.item(),
        'anomaly_map': outputs[1] if len(outputs) > 1 else None
    }

# 测试推理
predictions = predict("test_image.jpg")

方案2：TensorRT引擎推理（追求极致帧率） #

TensorRT是NVIDIA官方的优化推理引擎，能最大化GPU性能，适合1000-2000fps的高需求场景：

1. 安装依赖 #

需匹配CUDA版本安装TensorRT（建议从NVIDIA官网下载对应wheel包）：

pip install tensorrt onnx

2. 将ONNX模型转换为TensorRT引擎 #

import tensorrt as trt
import onnx

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('model.onnx', 'rb') as f:
    if not parser.parse(f.read()):
        for idx in range(parser.num_errors):
            print(parser.get_error(idx))

# 配置引擎参数（启用FP16加速）
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 分配1GB工作内存
if builder.platform_has_fast_fp16:
    config.set_flag(trt.BuilderFlag.FP16)

# 构建并保存引擎
engine = builder.build_engine(network, config)
with open('model.trt', 'wb') as f:
    f.write(engine.serialize())

3. 使用TensorRT引擎推理 #

import tensorrt as trt
import numpy as np
import cuda
from PIL import Image
import json

# 加载元数据
with open(metadata_path, 'r') as f:
    metadata = json.load(f)

# 加载TensorRT引擎
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
runtime = trt.Runtime(TRT_LOGGER)
with open('model.trt', 'rb') as f:
    engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()

# 分配CUDA内存
def allocate_buffers(engine):
    inputs, outputs, bindings = [], [], []
    stream = trt.Stream()
    for binding in engine:
        size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
        dtype = trt.nptype(engine.get_binding_dtype(binding))
        host_mem = np.zeros(size, dtype=dtype)
        device_mem = cuda.mem_alloc(host_mem.nbytes)
        bindings.append(int(device_mem))
        if engine.binding_is_input(binding):
            inputs.append({'host': host_mem, 'device': device_mem})
        else:
            outputs.append({'host': host_mem, 'device': device_mem})
    return inputs, outputs, bindings, stream

inputs, outputs, bindings, stream = allocate_buffers(engine)
input_name = engine.get_binding_name(0)
output_names = [engine.get_binding_name(i) for i in range(1, engine.num_bindings)]

# 推理函数
def predict(image_path):
    input_tensor = preprocess(image_path)  # 复用方案1的preprocess函数
    # 数据传输到GPU
    np.copyto(inputs[0]['host'], input_tensor.flatten())
    cuda.memcpy_htod_async(inputs[0]['device'], inputs[0]['host'], stream)
    # 异步推理
    context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
    # 结果传输回CPU
    for out in outputs:
        cuda.memcpy_dtoh_async(out['host'], out['device'], stream)
    stream.synchronize()
    
    # 后处理
    anomaly_score = outputs[0]['host'].reshape(1)
    threshold = metadata['metrics']['image_threshold']
    pred_label = anomaly_score > threshold
    
    return {
        'anomaly_score': anomaly_score.item(),
        'pred_label': pred_label.item(),
        'anomaly_map': outputs[1]['host'].reshape((1, metadata['image_size'][0], metadata['image_size'][1])) if len(outputs) > 1 else None
    }

高帧率优化建议 #

• 启用FP16/INT8量化：TensorRT和ONNX Runtime均支持，可大幅提升推理速度
• 批量推理：一次性输入多张图像，充分利用GPU并行计算能力
• 异步处理：将预处理、推理、后处理步骤并行执行，减少等待时间
• 版本匹配：确保GPU驱动、CUDA、TensorRT/ONNX Runtime版本相互兼容

内容的提问来源于stack exchange，提问作者Karel Debedts