Skip to content

Commit

Permalink
Update README
Browse files Browse the repository at this point in the history
  • Loading branch information
Linaom1214 committed Jul 27, 2022
1 parent 5433b4d commit 08bb095
Show file tree
Hide file tree
Showing 2 changed files with 336 additions and 168 deletions.
335 changes: 167 additions & 168 deletions README.md
Original file line number Diff line number Diff line change
@@ -1,169 +1,168 @@
# YOLO Series TensorRT Python/C++

## Support
YOLOv7、YOLOv6、 YOLOX、 YOLOV5、

## Update
- 2022.7.8 support YOLOV7
- 2022.7.3 support TRT int8 post-training quantization

## Prepare TRT Env
`Python`
```
pip install --upgrade setuptools pip --user
pip install nvidia-pyindex
pip install --upgrade nvidia-tensorrt
pip install pycuda
```
`C++`

[By Docker](https://github.com/NVIDIA/TensorRT/blob/main/docker/ubuntu-20.04.Dockerfile)

## Quick Start

Here is a Python Demo mybe help you quickly understand this repo [Link](https://aistudio.baidu.com/aistudio/projectdetail/4263301?contributionType=1&shared=1)

## YOLOv7 [C++, Python Support]

![](yolov7/3_yolov7.jpg)

```shell
https://github.com/WongKinYiu/yolov7.git
```
修改代码:将 yolo.py 对应行修改如下:
https://github.com/WongKinYiu/yolov7/blob/5f1b78ad614b45c5a98e7afdd295e20033d5ad3c/models/yolo.py#L57

```python
return x if self.training else (torch.cat(z, 1), ) if not self.export else (torch.cat(z, 1), x)
```

### 导出onnx
```shell
python models/export.py --weights ../yolov7.pt --grid
```

### 转化为TensorRT Engine

```
python export.py -o onnx-name -e trt-name -p fp32/16/int8
```
### 测试

```
cd yolov7
python trt.py
```

### C++

C++ [Demo](yolov7/cpp/README.md)





## YOLOv6 [C++, Python Support]

| model | input | | FPS | Device | Language |
| -------- | -------- | -------- | ------- | ------- | ------|
| yolov6s | 640*640 | FP16 | 360FPS | A100 | Python |
| yolov6s | 640*640 | FP32 | 350FPS | A100| Python |
| yolov6s | 640*640 | FP32 | 330FPS | 1080Ti | C++ |
| yolov6s | 640*640 | FP32 | 300FPS | 1080Ti | Python |

[YOLOv6 bilibili](https://www.bilibili.com/video/BV1x3411w7T6?share_source=copy_web)

![](yolov6/3_yolov6.jpg)
```shell
git clone https://github.com/meituan/YOLOv6.git
```
### 导出onnx
```shell
python deploy/ONNX/export_onnx.py --weights yolov6s.pt --img 640 --batch 1
```

### 转化为TensorRT Engine

```
python export.py -o onnx-name -e trt-name -p fp32/16/int8
```
### 测试

```
cd yolov6
python trt.py
```

### C++

C++ [Demo](yolov6/cpp/README.md)

## YOLOX [Python Support]
![](yolox/3_yolox.jpg)
### 导出ONNX

```
git clone https://github.com/Megvii-BaseDetection/YOLOX.git
```
```python

修改 export_onnx.py 为 model.head.decode_in_inference = True

修改 yolox/models/yolox_head.py文件

# [batch, n_anchors_all, 85]
# outputs = torch.cat(
# [x.flatten(start_dim=2) for x in outputs], dim=2

# ).permute(0, 2, 1)
outputs = torch.cat(
[x.view(-1,int(x.size(1)),int(x.size(2)*x.size(3))) for x in outputs], dim=2

).permute(0, 2, 1)

# outputs[..., :2] = (outputs[..., :2] + grids) * strides
# outputs[..., 2:4] = torch.exp(outputs[..., 2:4]) * strides
# return outputs
xy = (outputs[..., :2] + grids) * strides
wh = torch.exp(outputs[..., 2:4]) * strides
return torch.cat((xy, wh, outputs[..., 4:]), dim=-1)

```
```python
python3 tools/export_onnx.py --output-name yolox_s.onnx -n yolox-s -c yolox_s.pth
```
### 转化为TensorRT Engine
```
python export.py -o onnx-name -e trt-name -p fp32/16/int8
```
### 测试

```
cd yolovx
python trt.py
```

## YOLOV5 [Python Support]
![](yolov5/3_yolov5.jpg)

### 导出ONNX

```
git clone https://github.com/ultralytics/yolov5.git
```

```python
python path/to/export.py --weights yolov5s.pt --include onnx
```

### 转化为TensorRT Engine

```
python export.py -o onnx-name -e trt-name -p fp32/16/int8
```
### 测试

```
cd yolov5
python trt.py
# YOLO Series TensorRT Python/C++
## [简体中文](README_CN.md)

## Support
YOLOv7、YOLOv6、 YOLOX、 YOLOV5、

## Update
- 2022.7.8 support YOLOV7
- 2022.7.3 support TRT int8 post-training quantization

## Prepare TRT Env
`Python`
```
pip install --upgrade setuptools pip --user
pip install nvidia-pyindex
pip install --upgrade nvidia-tensorrt
pip install pycuda
```
`C++`

[By Docker](https://github.com/NVIDIA/TensorRT/blob/main/docker/ubuntu-20.04.Dockerfile)

## Quick Start

Here is a Python Demo mybe help you quickly understand this repo [Link](https://aistudio.baidu.com/aistudio/projectdetail/4263301?contributionType=1&shared=1)

## YOLOv7 [C++, Python Support]

![](yolov7/3_yolov7.jpg)

```shell
https://github.com/WongKinYiu/yolov7.git
```

Modify code: Modify the yolo.py as follows:
https://github.com/WongKinYiu/yolov7/blob/5f1b78ad614b45c5a98e7afdd295e20033d5ad3c/models/yolo.py#L57

```python
return x if self.training else (torch.cat(z, 1), ) if not self.export else (torch.cat(z, 1), x)
```

### Export onnx
```shell
python models/export.py --weights ../yolov7.pt --grid
```

### Convert to TensorRT Engine

```
python export.py -o onnx-name -e trt-name -p fp32/16/int8
```
### Test

```
cd yolov7
python trt.py
```

### C++

C++ [Demo](yolov7/cpp/README.md)


## YOLOv6 [C++, Python Support]

| model | input | | FPS | Device | Language |
| -------- | -------- | -------- | ------- | ------- | ------|
| yolov6s | 640*640 | FP16 | 360FPS | A100 | Python |
| yolov6s | 640*640 | FP32 | 350FPS | A100| Python |
| yolov6s | 640*640 | FP32 | 330FPS | 1080Ti | C++ |
| yolov6s | 640*640 | FP32 | 300FPS | 1080Ti | Python |

[YOLOv6 bilibili](https://www.bilibili.com/video/BV1x3411w7T6?share_source=copy_web)

![](yolov6/3_yolov6.jpg)
```shell
git clone https://github.com/meituan/YOLOv6.git
```
### Export onnx
```shell
python deploy/ONNX/export_onnx.py --weights yolov6s.pt --img 640 --batch 1
```

### Convert to TensorRT Engine

```
python export.py -o onnx-name -e trt-name -p fp32/16/int8
```
### Test

```
cd yolov6
python trt.py
```

### C++

C++ [Demo](yolov6/cpp/README.md)

## YOLOX [Python Support]
![](yolox/3_yolox.jpg)
### Export ONNX

```
git clone https://github.com/Megvii-BaseDetection/YOLOX.git
```
```python

Modify export_onnx.py as model.head.decode_in_inference = True

Modify yolox/models/yolox_head.py

# [batch, n_anchors_all, 85]
# outputs = torch.cat(
# [x.flatten(start_dim=2) for x in outputs], dim=2

# ).permute(0, 2, 1)
outputs = torch.cat(
[x.view(-1,int(x.size(1)),int(x.size(2)*x.size(3))) for x in outputs], dim=2

).permute(0, 2, 1)

# outputs[..., :2] = (outputs[..., :2] + grids) * strides
# outputs[..., 2:4] = torch.exp(outputs[..., 2:4]) * strides
# return outputs
xy = (outputs[..., :2] + grids) * strides
wh = torch.exp(outputs[..., 2:4]) * strides
return torch.cat((xy, wh, outputs[..., 4:]), dim=-1)

```
```python
python3 tools/export_onnx.py --output-name yolox_s.onnx -n yolox-s -c yolox_s.pth
```
### Convert to TensorRT Engine
```
python export.py -o onnx-name -e trt-name -p fp32/16/int8
```
### Test

```
cd yolovx
python trt.py
```

## YOLOV5 [Python Support]
![](yolov5/3_yolov5.jpg)

### Export ONNX

```
git clone https://github.com/ultralytics/yolov5.git
```

```python
python path/to/export.py --weights yolov5s.pt --include onnx
```

### Convert to TensorRT Engine

```
python export.py -o onnx-name -e trt-name -p fp32/16/int8
```
### Test

```
cd yolov5
python trt.py
```
Loading

0 comments on commit 08bb095

Please sign in to comment.