# yolox_tensorrt

**Repository Path**: dominic23331/yolox_tensorrt

## Basic Information

- **Project Name**: yolox_tensorrt
- **Description**: tensorrt加速yolox
- **Primary Language**: Unknown
- **License**: MulanPSL-2.0
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 4
- **Forks**: 0
- **Created**: 2022-09-23
- **Last Updated**: 2024-03-09

## Categories & Tags

**Categories**: Uncategorized

**Tags**: TensorRt, onnx

## README

# yolox_tensorrt

#### 介绍

本仓库为yolox在tensorrt中的部署，支持在PC端的部署以及在英伟达jetson nano上部署

#### 所需环境

1、numpy

2、tensorrt == 8.2.4.2（8.0及以上应该都可以，切勿小于8.0版本）

3、pycuda

4、pyyaml （用于设置初始化文件）

5、opencv-python

6、pillow

7、prettytable （用于显示网络参数，非必须）

#### 权值文件下载

链接：https://pan.baidu.com/s/1HuDrb1xtGKF14qM5nOZekQ 
提取码：yolo 
（权值文件为bubliiiig大神在coco数据集中预训练的模型转化的onnx，侵删）

#### 使用说明

1、将onnx文件预推理为tensorrt的engine文件，并放入engine文件夹下

```
trtexec --onnx=yolox_s.onnx --saveEngine=yolox_s.engine
```

2、设置config文件，以下为默认设置

```
# 输入网络的张量形状
input_shape: [1, 3, 640, 640]

# 归一化参数
mean: [0.485, 0.456, 0.406]
std: [0.229, 0.224, 0.225]

# 主干网络每个stage引出的图像尺寸
stage: [[80, 80], [40, 40], [20, 20]]

# nms中置信度阈值与iou阈值
conf_thres: 0.3
iou_thres: 0.5

# 类别
classes: [ 'person', 'bicycle', 'car', 'motorbike', 'aeroplane', 'bus', 'train',
           'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter',
           'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
           'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase',
           'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
           'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife',
           'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog',
           'pizza', 'donut', 'cake', 'chair', 'sofa', 'pottedplant', 'bed', 'diningtable', 'toilet',
           'tvmonitor', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven',
           'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear',
           'hair drier', 'toothbrush' ]
```

3、在inference.py中设置推理方式

```python3
    yolo = Yolo("engine/yolox_m.engine", "config/yolox.yaml", trt.Logger.ERROR)
    #---------------------------------------------------------#
    #    mode用于指定运行的模式
    #       "predict"      表示单张图片的预测
    #       "video"        表示视频预测，可预测视频文件或摄像头
    #---------------------------------------------------------#
    mode = "video"

    #---------------------------------------------------------#
    #    "show_type"       指定显示模式,可选择cv2或PIL
    #    "print_info"      指定是否打印预测结果
    #---------------------------------------------------------#
    show_type = "PIL"
    print_info = True

    # ---------------------------------------------------------#
    #    "video_path"      指定视频路径,0表示使用摄像头
    #    "video_save_path" 指定视频保存位置，如果为空表示不保存视频
    #    "video_fps"       指定保存视频的帧率
    # ---------------------------------------------------------#
    video_path = 0
    video_save_path = ""
    video_fps = 25.
```

4、运行inferen.py

```
python inference.py
```

#### 建议

1、在运行结果有问题时，将logger的版本指定为warning，观察是否有版本不匹配的情况

2、在jetson nano当中运行时，建议使用yolox_nano模型，并将图像大小设置为416x416

3、由于numpy版本问题，有的版本的numpy会在nms时报错找不到keepdims关键字，这时做以下改动即可

```python3
class_pred = np.argmax(prediction[:, 5: 5 + num_classes], axis=1, keepdims=True)
# 改为
class_pred = np.argmax(prediction[:, 5: 5 + num_classes], axis=1)
class_pred = class_pred.reshape(-1, 1)
```

4、prettytable库不是必须的，你可以将其注释掉，以自己喜欢的方式打印网络结构

#### Reference

1. https://github.com/bubbliiiing/yolox-pytorch (感谢bubbliiiing导师)
2. https://github.com/NVIDIA/trt-samples-for-hackathon-cn/tree/master/cookbook