batch inference use tensorrt

[BooHwang]

Hi, guys, I use mmclassification tools train a classification model, and use mmdeploy tools trans pt model to onnx and trt engine, but then, when I use below code to inference, it get right result, because it use right preprocess in model cfg:

from mmdeploy.apis import inference_model
import numpy as np
import cv2

CLASSES = {0: "beautiful", 1: "ugly"}
device = 'cuda:0'
model_cfg = ('/root/mmclassification/configs/efficientnet/efficientnet-b5_8xb32-01norm.py')
deploy_cfg = ('/root/mmclassification/mmdeploy/configs/mmcls/classification_tensorrt_dynamic-456x456.py')
# backend_files = ['/root/mmclassification/exp/best_accuracy_top-1.engine']
backend_files = ['/root/mmclassification/deploy_model/end2end.engine']
img = '/data/datasets/03700012.jpg'
img_np = cv2.imread(img)
model_output = inference_model(model_cfg, deploy_cfg,
                                backend_files, img_np, device)
pred_score = np.max(model_output[0], axis=1)
pred_label = np.argmax(model_output[0], axis=1)[0]
result = {'pred_label': pred_label, 'pred_score': float(pred_score), 'pred_class': CLASSES[pred_label]}
print(result)

However, the inference_model api just adapt to one image inference, I want to inference batch images, so, I use below code:

from mmdeploy.backend.tensorrt import TRTWrapper
import numpy as np
import cv2
import torch
from torchvision import transforms

img_list = list()
data_transform = transforms.Compose([
            transforms.ToPILImage(),
            transforms.Resize(456),
            transforms.CenterCrop(456),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
        
for i in range(32):
    img_np = cv2.imread(img)
    img_torch = data_transform(img_np)
    img_np = img_torch.numpy()[:, :, ::-1]
    img_list.append(img_np)
img_np_nchw = np.stack(img_list, axis=0).astype(np.float32)
img_tensor = torch.from_numpy(img_np_nchw)
inputs = dict(input=img_tensor.cuda())

model = TRTWrapper(backend_files[0], ['output'])
model_output = model(inputs)
print(model_output)
output = model_output["output"].cpu().numpy()
for i in range(output.shape[0]):
    pred_score = np.max(output[i, :])
    pred_label = np.argmax(output[i, :])
    result = {'pred_label': pred_label, 'pred_score': float(pred_score), 'pred_class': CLASSES[pred_label]}

But the result get wrong, I'm doubt it wrong in preprocess, anyone can help?

[lvhan028]

@AllentDan could you help answer this question?

[AllentDan]

Hi, @BooHwang . What did you mean by wrong result? I tested the script, the results made sense to me.