CLIP图文多模态模型onnxruntime和tensorrt推理

首先下载github项目：https://github.com/Lednik7/CLIP-ONNX
修改clip_onnx/utils.py第61行opset_version=12为opset_version=15 , 运行测试脚本：

import clip
from PIL import Image
import numpy as np

# ONNX不支持CUDA
model, preprocess = clip.load("ViT-B/32", device="cpu", jit=False)

# 批量优先
image = preprocess(Image.open("CLIP.png")).unsqueeze(0).cpu()  # [1, 3, 224, 224]
image_onnx = image.detach().cpu().numpy().astype(np.float32)

# 批量优先
text = clip.tokenize(["a diagram", "a dog", "a cat"]).cpu()  # [3, 77]
text_onnx = text.detach().cpu().numpy().astype(np.int32)


from clip_onnx import clip_onnx

onnx_model = clip_onnx(model, visual_path="clip_visual.onnx", textual_path="clip_textual.onnx")
onnx_model.convert2onnx(image, text, verbose=True)
onnx_model.start_sessions(providers=["CPUExecutionProvider"])  # CPU模式

image_features = onnx_model.encode_image(image_onnx)
text_features = onnx_model.encode_text(text_onnx)

logits_per_image, logits_per_text = onnx_model(image_onnx, text_onnx)
probs = logits_per_image.softmax(dim=-1).detach().cpu().numpy()
print("标签概率：", probs)

onnx_model = clip_onnx(None)
onnx_model.load_onnx(visual_path="clip_visual.onnx", textual_path="clip_textual.onnx", logit_scale=100.0000)  # model.logit_scale.exp()
onnx_model.start_sessions(providers=["CPUExecutionProvider"])
image_features = onnx_model.encode_image(image_onnx)
text_features = onnx_model.encode_text(text_onnx)

logits_per_image, logits_per_text = onnx_model(image_onnx, text_onnx)
probs = logits_per_image.softmax(dim=-1).detach().cpu().numpy()
print("标签概率：", probs)

onnxruntime推理

import clip
from PIL import Image
import numpy as np
import onnxruntime


model, preprocess = clip.load("ViT-B/32", device="cpu", jit=False)
image = preprocess(Image.open("CLIP.png")).unsqueeze(0).cpu()  # [1, 3, 224, 224]
image_input = image.detach().cpu().numpy().astype(np.float32)
text = clip.tokenize(["a diagram", "a dog", "a cat"]).cpu()  # [3, 77]
text_input = text.detach().cpu().numpy().astype(np.int32)


class clip_onnx:
    def __init__(self, model=None, visual_path: str = "clip_visual.onnx", textual_path: str = "clip_textual.onnx", logit_scale=None):
        self.model = model
        self.visual_path = visual_path
        self.textual_path = textual_path
        self.logit_scale = logit_scale

    def start_sessions(self, providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']):
        self.visual_session = onnxruntime.InferenceSession(self.visual_path, providers=providers)
        self.textual_session = onnxruntime.InferenceSession(self.textual_path, providers=providers)

    def __call__(self, image, text, device: str = "cpu"):
        onnx_input_image = {self.visual_session.get_inputs()[0].name: image}
        image_features, = self.visual_session.run(None, onnx_input_image)
        onnx_input_text = {self.textual_session.get_inputs()[0].name: text}
        text_features, = self.textual_session.run(None, onnx_input_text)
        image_features = image_features /np.linalg.norm(image_features, axis=-1, keepdims=True)
        text_features = text_features / np.linalg.norm(text_features, axis=-1, keepdims=True)
        logits_per_image = self.logit_scale * image_features @ text_features.T
        logits_per_text = logits_per_image.T
        return logits_per_image, logits_per_text


onnx_model = clip_onnx(visual_path="clip_visual.onnx", textual_path="clip_textual.onnx", logit_scale=100.0000)  
onnx_model.start_sessions(providers=["CPUExecutionProvider"])
logits_per_image, logits_per_text = onnx_model(image_input, text_input)

def softmax(x, axis=-1):
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e_x / np.sum(e_x, axis=axis, keepdims=True)

probs = softmax(logits_per_image)
print("标签概率：", probs)

tensorrt推理

要使得tensorrt正常推理还需注释掉clip_onnx/utils.py第63行的dynamic_axes，再生成onnx模型。
通过trtexec 转换onnx模型得到engine文件：

/docker_share/TensorRT-8.6.1.6/bin/trtexec --onnx=./clip_textual.onnx --saveEngine=./clip_textual.engine 
/docker_share/TensorRT-8.6.1.6/bin/trtexec --onnx=./clip_visual.onnx --saveEngine=./clip_visual.engine

import clip
from PIL import Image
import numpy as np
import tensorrt as trt
import pycuda.autoinit 
import pycuda.driver as cuda 


model, preprocess = clip.load("ViT-B/32", device="cpu", jit=False)
image = preprocess(Image.open("CLIP.png")).unsqueeze(0).cpu()  # [1, 3, 224, 224]
image_input = image.detach().cpu().numpy().astype(np.float32)
text = clip.tokenize(["a diagram", "a dog", "a cat"]).cpu()  # [3, 77]
text_input = text.detach().cpu().numpy().astype(np.int32)


class clip_tensorrt:
    def __init__(self, model=None, visual_path: str = "clip_visual.engine", textual_path: str = "clip_textual.engine", logit_scale=None):
        self.model = model
        self.visual_path = visual_path
        self.textual_path = textual_path
        self.logit_scale = logit_scale

    def start_engines(self, providers=['CUDAExecutionProvider', 'CPUExecutionProvider']):
        logger = trt.Logger(trt.Logger.WARNING)
        with open(self.visual_path, "rb") as f, trt.Runtime(logger) as runtime:
            self.visual_engine = runtime.deserialize_cuda_engine(f.read())
        self.visual_context = self.visual_engine.create_execution_context()
        self.visual_inputs_host = cuda.pagelocked_empty(trt.volume(self.visual_context.get_binding_shape(0)), dtype=np.float32)
        self.visual_outputs_host = cuda.pagelocked_empty(trt.volume(self.visual_context.get_binding_shape(1)), dtype=np.float32)
        self.visual_inputs_device = cuda.mem_alloc(self.visual_inputs_host.nbytes)
        self.visual_outputs_device = cuda.mem_alloc(self.visual_outputs_host.nbytes)
        self.visual_stream = cuda.Stream()

        with open(self.textual_path, "rb") as f, trt.Runtime(logger) as runtime:
            self.textual_engine = runtime.deserialize_cuda_engine(f.read())
        self.textual_context = self.textual_engine.create_execution_context()
        self.textual_inputs_host = cuda.pagelocked_empty(trt.volume(self.textual_context.get_binding_shape(0)), dtype=np.int32)
        self.textual_outputs_host = cuda.pagelocked_empty(trt.volume(self.textual_context.get_binding_shape(1)), dtype=np.float32)
        self.textual_inputs_device = cuda.mem_alloc(self.textual_inputs_host.nbytes)
        self.textual_outputs_device = cuda.mem_alloc(self.textual_outputs_host.nbytes)
        self.textual_stream = cuda.Stream()
        
    def __call__(self, image, text):
        np.copyto(self.visual_inputs_host, image.ravel())
        with self.visual_engine.create_execution_context() as context:
            cuda.memcpy_htod_async(self.visual_inputs_device, self.visual_inputs_host, self.visual_stream)
            context.execute_async_v2(bindings=[int(self.visual_inputs_device), int(self.visual_outputs_device)], stream_handle=self.visual_stream.handle)
            cuda.memcpy_dtoh_async(self.visual_outputs_host, self.visual_outputs_device, self.visual_stream)
            self.visual_stream.synchronize()  

        np.copyto(self.textual_inputs_host, text.ravel())
        with self.textual_engine.create_execution_context() as context:
            cuda.memcpy_htod_async(self.textual_inputs_device, self.textual_inputs_host, self.textual_stream)
            context.execute_async_v2(bindings=[int(self.textual_inputs_device), int(self.textual_outputs_device)], stream_handle=self.textual_stream.handle)
            cuda.memcpy_dtoh_async(self.textual_outputs_host, self.textual_outputs_device, self.textual_stream)
            self.textual_stream.synchronize()  
        
        image_features = self.visual_outputs_host.reshape(1,512)
        text_features = self.textual_outputs_host.reshape(3,512)
        image_features = image_features /np.linalg.norm(image_features, axis=-1, keepdims=True)
        text_features = text_features / np.linalg.norm(text_features, axis=-1, keepdims=True)
        logits_per_image = self.logit_scale * image_features @ text_features.T
        logits_per_text = logits_per_image.T
        return logits_per_image, logits_per_text


onnx_model = clip_tensorrt(visual_path="clip_visual.engine", textual_path="clip_textual.engine", logit_scale=100.0000)  
onnx_model.start_engines()
logits_per_image, logits_per_text = onnx_model(image_input, text_input)

def softmax(x, axis=-1):
    e_x = np.exp(x - np.max(x, axis=axis, keepdims=True))
    return e_x / np.sum(e_x, axis=axis, keepdims=True)

probs = softmax(logits_per_image)
print("标签概率：", probs)

CLIP.png

正常运行得到结果

标签概率： [[0.9927926  0.00421789 0.00298956]]