diff --git a/yolov5/README.md b/yolov5/README.md index 0755d4e..d5064cf 100644 --- a/yolov5/README.md +++ b/yolov5/README.md @@ -1,12 +1,41 @@ -# yolov5 +# YOLOv5 -The Pytorch implementation is [ultralytics/yolov5](https://github.com/ultralytics/yolov5). +TensorRTx inference code base for [ultralytics/yolov5](https://github.com/ultralytics/yolov5). + +## Contributors + + + + + + + + + + + + + + + + + + + + + + + + + + + ## Different versions of yolov5 -Currently, we support yolov5 v1.0, v2.0, v3.0, v3.1, v4.0, v5.0, v6.0, v6.2. +Currently, we support yolov5 v1.0, v2.0, v3.0, v3.1, v4.0, v5.0, v6.0, v6.2 -- For yolov5 v6.2, download .pt from [yolov5 release v6.2](https://github.com/ultralytics/yolov5/releases/tag/v6.2), `git clone -b v6.2 https://github.com/ultralytics/yolov5.git` and `git clone -b yolov5-v6.2 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in current page. +- For yolov5 v6.2, download .pt from [yolov5 release v6.2](https://github.com/ultralytics/yolov5/releases/tag/v6.2), `git clone -b v6.2 https://github.com/ultralytics/yolov5.git` and `git clone -b yolov5-v6.2 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v6.2](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v6.2/yolov5) - For yolov5 v6.0, download .pt from [yolov5 release v6.0](https://github.com/ultralytics/yolov5/releases/tag/v6.0), `git clone -b v6.0 https://github.com/ultralytics/yolov5.git` and `git clone -b yolov5-v6.0 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v6.0](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v6.0/yolov5). - For yolov5 v5.0, download .pt from [yolov5 release v5.0](https://github.com/ultralytics/yolov5/releases/tag/v5.0), `git clone -b v5.0 https://github.com/ultralytics/yolov5.git` and `git clone -b yolov5-v5.0 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v5.0](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v5.0/yolov5). - For yolov5 v4.0, download .pt from [yolov5 release v4.0](https://github.com/ultralytics/yolov5/releases/tag/v4.0), `git clone -b v4.0 https://github.com/ultralytics/yolov5.git` and `git clone -b yolov5-v4.0 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v4.0](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v4.0/yolov5). @@ -26,13 +55,15 @@ Currently, we support yolov5 v1.0, v2.0, v3.0, v3.1, v4.0, v5.0, v6.0, v6.2. - BBox confidence thresh in yolov5.cpp - Batch size in yolov5.cpp -## How to Run, yolov5s as example +## Build and Run + +### Detection 1. generate .wts from pytorch with .pt, or download .wts from model zoo ``` // clone code according to above #Different versions of yolov5 -// download https://github.com/ultralytics/yolov5/releases/download/v6.0/yolov5s.pt +// download https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5s.pt cp {tensorrtx}/yolov5/gen_wts.py {ultralytics}/yolov5 cd {ultralytics}/yolov5 python gen_wts.py -w yolov5s.pt -o yolov5s.wts @@ -72,7 +103,18 @@ python yolov5_trt.py python yolov5_trt_cuda_python.py ``` -5. optional, run yolov5 classification models with similar steps +### Classification + +``` +# Download ImageNet labels +wget https://github.com/joannzhang00/ImageNet-dataset-classes-labels/blob/main/imagenet_classes.txt + +# Build and serialize TensorRT engine +./yolov5-cls -s yolov5s-cls.wts yolov5s-cls.engine s + +# Run inference +./yolov5-cls -d yolov5s-cls.engine ../samples +``` # INT8 Quantization @@ -85,11 +127,7 @@ python yolov5_trt_cuda_python.py 4. serialize the model and test

- -

- -

- +

## More Information diff --git a/yolov5/gen_wts.py b/yolov5/gen_wts.py index 621469f..4e1d670 100644 --- a/yolov5/gen_wts.py +++ b/yolov5/gen_wts.py @@ -28,17 +28,18 @@ def parse_args(): pt_file, wts_file, m_type = parse_args() +print(f'Generating .wts for {m_type} model') # Initialize device = select_device('cpu') # Load model +print(f'Loading {pt_file}') model = torch.load(pt_file, map_location=device) # load to FP32 model = model['ema' if model.get('ema') else 'model'].float() if m_type == "detect": # update anchor_grid info - anchor_grid = model.model[-1].anchors * \ - model.model[-1].stride[..., None, None] + anchor_grid = model.model[-1].anchors * model.model[-1].stride[..., None, None] # model.model[-1].anchor_grid = anchor_grid delattr(model.model[-1], 'anchor_grid') # model.model[-1] is detect layer # The parameters are saved in the OrderDict through the "register_buffer" method, and then saved to the weight. @@ -47,6 +48,7 @@ if m_type == "detect": model.to(device).eval() +print(f'Writing into {wts_file}') with open(wts_file, 'w') as f: f.write('{}\n'.format(len(model.state_dict().keys()))) for k, v in model.state_dict().items(): diff --git a/yolov5/yolov5_cls.cpp b/yolov5/yolov5_cls.cpp index 6ea2504..15c8619 100644 --- a/yolov5/yolov5_cls.cpp +++ b/yolov5/yolov5_cls.cpp @@ -69,7 +69,10 @@ std::vector topk(const std::vector& vec, int k) { std::vector read_classes(std::string file_name) { std::vector classes; std::ifstream ifs(file_name, std::ios::in); - assert(ifs.is_open()); + if (!ifs.is_open()) { + std::cerr << file_name << " is not found, pls refer to README and download it." << std::endl; + assert(0); + } std::string s; while (std::getline(ifs, s)) { classes.push_back(s); @@ -250,7 +253,7 @@ int main(int argc, char** argv) { std::cerr << "read_files_in_dir failed." << std::endl; return -1; } - auto classes = read_classes("../imagenet_classes.txt"); + auto classes = read_classes("imagenet_classes.txt"); static float data[BATCH_SIZE * 3 * INPUT_H * INPUT_W]; static float prob[BATCH_SIZE * OUTPUT_SIZE];