From e53982860075b4cf9ad95a8b78c3f1921053b5c6 Mon Sep 17 00:00:00 2001 From: triple-Mu Date: Fri, 13 Jan 2023 14:28:45 +0800 Subject: [PATCH] Add seg README --- README.md | 79 ++++++++++++++++++------------------- build.py | 4 ++ docs/API-Build.md | 31 +++++++++++++++ docs/Segment.md | 99 +++++++++++++++++++++++++++++++++++++++++++++++ models/engine.py | 9 +++-- 5 files changed, 178 insertions(+), 44 deletions(-) create mode 100644 docs/API-Build.md create mode 100644 docs/Segment.md diff --git a/README.md b/README.md index 2c4c546..7227a75 100644 --- a/README.md +++ b/README.md @@ -40,15 +40,18 @@ wget https://github.com/ultralytics/ultralytics/releases/download/v8.0.0/yolov8x.pt ``` -# Build TensorRT engine by ONNX +# Build TensorRT Engine by ONNX + ## Export ONNX by `ultralytics` API +### Export Your Own ONNX model + You can export your onnx model by `ultralytics` API and add postprocess into model at the same time. ``` shell -python export.py \ +python3 export.py \ --weights yolov8s.pt \ --iou-thres 0.65 \ --conf-thres 0.25 \ @@ -72,9 +75,9 @@ python export.py \ You will get an onnx model whose prefix is the same as input weights. -## Preprocessed ONNX model +### Just Taste First -If you just want to taste first, you can dowload the onnx model which are exported by `YOLOv8` package and modified by me. +If you just want to taste first, you can download the onnx model which are exported by `YOLOv8` package and modified by me. [**YOLOv8-n**](https://triplemu.oss-cn-beijing.aliyuncs.com/YOLOv8/ONNX/yolov8n_nms.onnx?OSSAccessKeyId=LTAI5tN1dgmZD4PF8AJUXp3J&Expires=1772936700&Signature=r6HgJTTcCSAxQxD9bKO9qBTtigQ%3D) @@ -86,14 +89,14 @@ If you just want to taste first, you can dowload the onnx model which are export [**YOLOv8-x**](https://triplemu.oss-cn-beijing.aliyuncs.com/YOLOv8/ONNX/yolov8x_nms.onnx?OSSAccessKeyId=LTAI5tN1dgmZD4PF8AJUXp3J&Expires=1673936778&Signature=3o%2F7QKhiZg1dW3I6sDrY4ug6MQU%3D) -## 1. By TensorRT ONNX Python api +## Export Engine by TensorRT Python api You can export TensorRT engine from ONNX by [`build.py` ](build.py). Usage: ``` shell -python build.py \ +python3 build.py \ --weights yolov8s.onnx \ --iou-thres 0.65 \ --conf-thres 0.25 \ @@ -113,14 +116,17 @@ python build.py \ You can modify `iou-thres` `conf-thres` `topk` by yourself. -## 2. By trtexec tools +## 2. Export Engine by Trtexec Tools You can export TensorRT engine by [`trtexec`](https://github.com/NVIDIA/TensorRT/tree/main/samples/trtexec) tools. Usage: ``` shell -/usr/src/tensorrt/bin/trtexec --onnx=yolov8s.onnx --saveEngine=yolov8s.engine --fp16 +/usr/src/tensorrt/bin/trtexec \ +--onnx=yolov8s.onnx \ +--saveEngine=yolov8s.engine \ +--fp16 ``` **If you installed TensorRT by a debian package, then the installation path of `trtexec` @@ -128,41 +134,27 @@ is `/usr/src/tensorrt/bin/trtexec`** **If you installed TensorRT by a tar package, then the installation path of `trtexec` is under the `bin` folder in the path you decompressed** +# Build TensorRT Engine by TensorRT API +Please see more information in [`API-Build.md`](docs/API-Build.md) -# Build TensorRT engine by API +***Notice !!!*** We don't support YOLOv8-seg model now !!! -When you want to build engine by api. You should generate the pickle weights parameters first. +# Inference -``` shell -python gen_pkl.py -w yolov8s.pt -o yolov8s.pkl -``` - -You will get a `yolov8s.pkl` which contain the operators' parameters. And you can rebuild `yolov8s` model in TensorRT api. - -``` -python build.py \ ---weights yolov8s.pkl \ ---iou-thres 0.65 \ ---conf-thres 0.25 \ ---topk 100 \ ---fp16 \ ---input-shape 1 3 640 640 \ ---device cuda:0 -``` - -***Notice !!!*** Now we only support static input shape model build by TensorRT api. You'd best give the legal`input-shape`. - -# Infer images by the engine which you export or build - -## 1. Python infer +## 1. Infer with python script You can infer images with the engine by [`infer.py`](infer.py) . Usage: ``` shell -python3 infer.py --engine yolov8s.engine --imgs data --show --out-dir outputs --device cuda:0 +python3 infer.py \ +--engine yolov8s.engine \ +--imgs data \ +--show \ +--out-dir outputs \ +--device cuda:0 ``` #### Description of all arguments @@ -174,13 +166,13 @@ python3 infer.py --engine yolov8s.engine --imgs data --show --out-dir outputs -- - `--device` : The CUDA deivce you use. - `--profile` : Profile the TensorRT engine. -## 2. C++ infer +## 2. Infer with C++ -You can infer with c++ in [`csrc/end2end`](csrc/end2end) . +You can infer with c++ in [`csrc/detect`](csrc/detect) . -Build: +### Build: -Please set you own librarys in [`CMakeLists.txt`](csrc/end2end/CMakeLists.txt) and modify you own config in [`config.h`](csrc/end2end/include/config.h) such as `CLASS_NAMES` and `COLORS`. +Please set you own librarys in [`CMakeLists.txt`](csrc/detect/CMakeLists.txt) and modify you own config in [`config.h`](csrc/detect/include/config.h) such as `CLASS_NAMES` and `COLORS`. ``` shell export root=${PWD} @@ -203,6 +195,15 @@ Usage: ./yolov8 yolov8s.engine data/test.mp4 # the video path ``` +# TensorRT Segment Deploy + +Please see more information in [`Segment.md`](docs/Segment.md) + +# DeepStream Detection Deploy + +See more in [`README.md`](csrc/deepstream/README.md) + + # Profile you engine If you want to profile the TensorRT engine: @@ -212,7 +213,3 @@ Usage: ``` shell python3 infer.py --engine yolov8s.engine --profile ``` - -# DeepStream Deploy - -See more in [`README.md`](csrc/deepstream/README.md) diff --git a/build.py b/build.py index a0f687f..a097eea 100644 --- a/build.py +++ b/build.py @@ -36,6 +36,9 @@ def parse_args(): type=str, default='cuda:0', help='TensorRT builder device') + parser.add_argument('--seg', + action='store_true', + help='Build seg model by onnx') args = parser.parse_args() assert len(args.input_shape) == 4 return args @@ -43,6 +46,7 @@ def parse_args(): def main(args): builder = EngineBuilder(args.weights, args.device) + builder.seg = args.seg builder.build(fp16=args.fp16, input_shape=args.input_shape, iou_thres=args.iou_thres, diff --git a/docs/API-Build.md b/docs/API-Build.md new file mode 100644 index 0000000..4f2055d --- /dev/null +++ b/docs/API-Build.md @@ -0,0 +1,31 @@ +# Build TensorRT Engine By TensorRT Python API + +When you want to build engine by API. You should generate the pickle weights parameters first. + +``` shell +python3 gen_pkl.py -w yolov8s.pt -o yolov8s.pkl +``` + +You will get a `yolov8s.pkl` which contain the operators' parameters. + +And you can rebuild `yolov8s` model in TensorRT api. + +``` +python3 build.py \ +--weights yolov8s.pkl \ +--iou-thres 0.65 \ +--conf-thres 0.25 \ +--topk 100 \ +--fp16 \ +--input-shape 1 3 640 640 \ +--device cuda:0 +``` + +***Notice !!!*** + +Now we only support static input shape model build by TensorRT api. +You'd best give the legal `input-shape`. + +***Notice !!!*** + +Now we don't support YOLOv8-seg model building by API. It will be supported later. diff --git a/docs/Segment.md b/docs/Segment.md new file mode 100644 index 0000000..c026ca2 --- /dev/null +++ b/docs/Segment.md @@ -0,0 +1,99 @@ +# YOLOv8-seg Model with TensorRT + +Instance segmentation models are currently experimental. + +Our conversion route is : + YOLOv8 PyTorch model -> ONNX -> TensorRT Engine + +***Notice !!!*** We don't support TensorRT API building !!! + +# Export Your Own ONNX model + +You can export your onnx model by `ultralytics` API. + +``` shell +python3 export_seg.py \ +--weights yolov8s-seg.pt \ +--opset 11 \ +--sim \ +--input-shape 1 3 640 640 \ +--device cuda:0 +``` + +#### Description of all arguments + +- `--weights` : The PyTorch model you trained. +- `--opset` : ONNX opset version, default is 11. +- `--sim` : Whether to simplify your onnx model. +- `--input-shape` : Input shape for you model, should be 4 dimensions. +- `--device` : The CUDA deivce you export engine . + +You will get an onnx model whose prefix is the same as input weights. + +This onnx model doesn't contain postprocessing. + +# Export Engine by TensorRT Python api + +You can export TensorRT engine from ONNX by [`build.py` ](../build.py). + +Usage: + +``` shell +python3 build.py \ +--weights yolov8s-seg.onnx \ +--fp16 \ +--device cuda:0 \ +--seg +``` + +#### Description of all arguments + +- `--weights` : The ONNX model you download. +- `--fp16` : Whether to export half-precision engine. +- `--device` : The CUDA deivce you export engine. +- `--seg` : Whether to export seg engine. + +# Export Engine by Trtexec Tools + +You can export TensorRT engine by [`trtexec`](https://github.com/NVIDIA/TensorRT/tree/main/samples/trtexec) tools. + +Usage: + +``` shell +/usr/src/tensorrt/bin/trtexec \ +--onnx=yolov8s-seg.onnx \ +--saveEngine=yolov8s-seg.engine \ +--fp16 +``` + +# Inference + +## Infer with python script + +You can infer images with the engine by [`infer.py`](../infer.py) . + +Usage: + +``` shell +python3 infer.py \ +--engine yolov8s-seg.engine \ +--imgs data \ +--show \ +--out-dir outputs \ +--device cuda:0 \ +--seg +``` + +#### Description of all arguments + +- `--engine` : The Engine you export. +- `--imgs` : The images path you want to detect. +- `--show` : Whether to show detection results. +- `--out-dir` : Where to save detection results images. It will not work when use `--show` flag. +- `--device` : The CUDA deivce you use. +- `--profile` : Profile the TensorRT engine. +- `--seg` : Infer with seg model. + +## Infer with C++ + +***Notice !!!*** COMMING SOON !!! diff --git a/models/engine.py b/models/engine.py index 8b55ff5..a558b0f 100644 --- a/models/engine.py +++ b/models/engine.py @@ -9,6 +9,7 @@ import torch class EngineBuilder: + seg = False def __init__( self, @@ -77,9 +78,10 @@ class EngineBuilder: topk: int = 100): parser = trt.OnnxParser(self.network, self.logger) onnx_model = onnx.load(str(self.checkpoint)) - onnx_model.graph.node[-1].attribute[2].i = topk - onnx_model.graph.node[-1].attribute[3].f = conf_thres - onnx_model.graph.node[-1].attribute[4].f = iou_thres + if not self.seg: + onnx_model.graph.node[-1].attribute[2].i = topk + onnx_model.graph.node[-1].attribute[3].f = conf_thres + onnx_model.graph.node[-1].attribute[4].f = iou_thres if not parser.parse(onnx_model.SerializeToString()): raise RuntimeError( @@ -110,6 +112,7 @@ class EngineBuilder: conf_thres: float = 0.25, topk: int = 100, ): + assert not self.seg from .api import SPPF, C2f, Conv, Detect, get_depth, get_width with open(self.checkpoint, 'rb') as f: