Skip to content

Commit 269ae9c

Browse files
Li-HongdaLi-Hongda
committed
Surport yolonas
1 parent 84bbd23 commit 269ae9c

13 files changed

Lines changed: 268 additions & 11 deletions

File tree

‎README.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ This repo use TensorRT-8.x to deploy well-trained models, both image preprocessi
3737
+ 2023.05.19 🚀 Support cuda mask postprocess and support rtdetr.
3838
+ 2023.05.21 🚀 Support yolov6.
3939
+ 2023.05.26 🚀 Support dynamic batch inference.
40+
+ 2023.06.07 🚀 Support yolox and yolo-nas.
4041
</details>
4142

4243
## 3.Support Models
@@ -51,6 +52,7 @@ This repo use TensorRT-8.x to deploy well-trained models, both image preprocessi
5152
- [x] [YOLOv8-seg](https://github.com/ultralytics/ultralytics)<br>
5253
- [x] [YOLOX](https://github.com/Megvii-BaseDetection/YOLOX)<br>
5354
- [x] [RT-DETR](https://github.com/PaddlePaddle/PaddleDetection/tree/develop/configs/rtdetr)<br>
55+
- [x] [YOLO-NAS](https://github.com/Deci-AI/super-gradients)<br>
5456
</details>
5557

5658
All speed tests were performed on RTX 3090 with COCO Val set.The time calculated here is the sum of the time of image loading, preprocess, inference and postprocess, so it's going to be slower than what's reported in the paper.
@@ -65,6 +67,8 @@ All speed tests were performed on RTX 3090 with COCO Val set.The time calculated
6567
| YOLOv7 | 1 | FP32 | 640x640 | 107 |
6668
| YOLOv8-s | 1 | FP32 | 640x640 | 171 |
6769
| YOLOv8-seg-s | 1 | FP32 | 640x640 | 122 |
70+
| YOLOX-s | 1 | FP32 | 640x640 | 156 |
71+
| YOLO-NAS-s | 1 | FP32 | 640x640 | 165 |
6872
| RT-DETR | 1 | FP32 | 640x640 | 106 |
6973
</div>
7074

‎configs/yolonas.yaml‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
yolonas:
2+
onnx_file: "../weights/yolonas/yolonas-s.onnx"
3+
engine_file: "../weights/yolonas/yolonas-s.trt"
4+
type: "coco"
5+
mode: "fp32"
6+
dynamic: 0
7+
batchSize: 1
8+
imageWidth: 640
9+
imageHeight: 640
10+
conf_thr: 0.25
11+
nms_thr: 0.45
12+
strides: [8, 16, 32]
13+
imgScale: 255
14+
imgMean: [ 0, 0, 0 ]
15+
imgStd: [ 1, 1, 1 ]

‎include/build.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
#include "yolov7.h"
88
#include "yolov8.h"
99
#include "yolox.h"
10+
#include "yolonas.h"
1011
#include "rtdetr.h"
1112

1213
std::shared_ptr<Model> build_model(std::string model_arch, std::string cfg);

‎include/cuda_function.h‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,9 @@ void yolov8_postprocess_box(float* predict, int num_bboxes, int num_classes, int
3535
void rtdetr_postprocess_box(float* predict_box, float* predict_cls, int num_bboxes, int num_classes, int num_out,
3636
float conf_thr, int imageWidth, int imageHeight, AffineMatrix mat, cudaStream_t stream, float* dst);
3737

38+
void yolonas_postprocess_box(float* predict, int num_bboxes, int num_classes, int num_out, float conf_thr,
39+
float nms_thr, AffineMatrix mat, cudaStream_t stream, float* dst);
40+
3841
void postprocess_box_mask(float* predict, int num_bboxes, int num_classes, int num_out,
3942
float conf_thr, float nms_thr, AffineMatrix mat, cudaStream_t stream, float* dst);
4043

‎include/yolonas.h‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
#ifndef YOLONAS_H
2+
#define YOLONAS_H
3+
4+
#include "yolov8.h"
5+
6+
class YOLONAS : public YOLOv8 {
7+
public:
8+
explicit YOLONAS(const YAML::Node &config);
9+
std::vector<Detections> PostProcess(const std::vector<cv::Mat> &imgBatch, float* output);
10+
};
11+
12+
#endif

‎object_detection/CMakeLists.txt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,11 +66,12 @@ add_subdirectory(yolov6)
6666
add_subdirectory(yolov7)
6767
add_subdirectory(yolov8)
6868
add_subdirectory(yolox)
69+
add_subdirectory(yolonas)
6970
add_subdirectory(rtdetr)
7071
set(LIBRARY_OUTPUT_PATH ../libs)
7172
add_library(build SHARED ../src/build.cpp)
7273

7374
set(EXECUTABLE_OUTPUT_PATH ../../bin)
7475
add_executable(object_detection main.cpp)
75-
target_link_libraries(object_detection yaml-cpp build yolov5 yolov6 yolov7 yolov8 yolox rtdetr cudart)
76+
target_link_libraries(object_detection yaml-cpp build yolov5 yolov6 yolov7 yolov8 yolox yolonas rtdetr cudart)
7677

‎object_detection/yolonas/CMakeLists.txt‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
cmake_minimum_required(VERSION 3.10)
2+
3+
set(CMAKE_BUILD_TYPE "Debug")
4+
# set(CMAKE_BUILD_TYPE "Release")
5+
6+
project(yolonas)
7+
8+
set(CMAKE_CXX_STANDARD 14)
9+
10+
# YAML
11+
set(YAML_LIBRARY ../../yaml-cpp/build)
12+
set(YAML_INCLUDE ../../yaml-cpp/include)
13+
link_directories(${YAML_LIBRARY})
14+
list(APPEND ALL_INCLUDE ${YAML_INCLUDE})
15+
16+
17+
# Project
18+
set(PROJECT_INCLUDE ../../include)
19+
message(STATUS "Find project include at ${PROJECT_INCLUDE}")
20+
list(APPEND ALL_INCLUDE ${PROJECT_INCLUDE})
21+
set(PROJECT_LIBRARY ../../src/common.cpp
22+
../../src/basemodel.cpp
23+
../../src/detection.cpp
24+
../../src/yolo.cpp
25+
../../src/yolonas.cpp
26+
../../src/cuda_function.cu)
27+
28+
include_directories(${ALL_INCLUDE})
29+
30+
set(LIBRARY_OUTPUT_PATH ../../libs)
31+
add_library(yolonas SHARED ${PROJECT_LIBRARY} ${SAMPLES_LIBRARY})
32+
target_link_libraries(yolonas ${ALL_LIBS} yaml-cpp)

‎object_detection/yolonas/export.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
import argparse
2+
import onnx
3+
import onnxsim
4+
import torch
5+
import torch.nn as nn
6+
from super_gradients.training import models
7+
from super_gradients.common.object_names import Models
8+
9+
10+
11+
class YOLONAS(nn.Module):
12+
def __init__(self, model):
13+
super().__init__()
14+
self.model = model
15+
self.model.eval()
16+
17+
def forward(self, input):
18+
19+
output = self.model(input)
20+
return torch.cat(output, dim=-1)
21+
22+
23+
def parse_args():
24+
parser = argparse.ArgumentParser()
25+
parser.add_argument('--model', type=str, default='yolo_nas_m',
26+
choices=['yolo_nas_s','yolo_nas_m', 'yolo_nas_l'] ,
27+
help='model.pt')
28+
parser.add_argument('--save-model', type=str, default='yolonas-m.onnx',
29+
help='model.onnx')
30+
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640],
31+
help='image (h, w)')
32+
parser.add_argument('--batch-size', type=int, default=1, help='batch size')
33+
parser.add_argument('--half', action='store_true', help='FP16 export')
34+
parser.add_argument('--dynamic', action='store_true', help='dynamic axes')
35+
parser.add_argument('--simplify', action='store_false', help='simplify model')
36+
parser.add_argument('--opset', type=int, default=11, help='opset version')
37+
args = parser.parse_args()
38+
return args
39+
40+
41+
def main(model,
42+
save_model,
43+
img_size,
44+
batch_size,
45+
opset = 11,
46+
half = False,
47+
dynamic = False,
48+
simplify = True):
49+
model = models.get(model, pretrained_weights="coco")
50+
model.prep_model_for_conversion(input_size=[1, 3, 640, 640])
51+
52+
model = YOLONAS(model)
53+
model.eval()
54+
if dynamic:
55+
dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'},
56+
'output0': {0: 'batch', 1: 'anchors'}}
57+
58+
img_size *= 2 if len(img_size) == 1 else 1
59+
dummy_input = torch.zeros(batch_size, 3, *img_size)
60+
61+
torch.onnx.export(model,
62+
dummy_input,
63+
save_model,
64+
input_names=['images'],
65+
output_names=['output0'],
66+
opset_version=opset,
67+
do_constant_folding=True,
68+
dynamic_axes=dynamic or None)
69+
model_onnx = onnx.load(save_model)
70+
onnx.checker.check_model(model_onnx)
71+
if simplify:
72+
model_onnx, check = onnxsim.simplify(model_onnx)
73+
assert check, 'simplify failed'
74+
onnx.save(model_onnx, save_model)
75+
76+
if __name__ == '__main__':
77+
args = parse_args()
78+
main(**vars(args))

‎object_detection/yolonas/main.cpp‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
#include "../../include/yolonas.h"
2+
3+
4+
int main(int argc, char** argv) {
5+
std::string inputpath = argv[1];
6+
std::string cfg_dir = "../configs";
7+
std::string cfg_suffix = ".yaml";
8+
std::string savedir = "../results";
9+
auto savepath = savedir + "/" + "yolonas" + "/";
10+
auto cfg = cfg_dir + "/" + "yolonas" + cfg_suffix;
11+
YAML::Node root = YAML::LoadFile(cfg);
12+
YOLONAS YOLONAS(root["yolonas"]);
13+
YOLONAS.LoadEngine();
14+
YOLONAS.Inference(inputpath, savepath);
15+
return 0;
16+
}

‎src/basemodel.cpp‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ Model::Model(const YAML::Node &config) {
1818

1919
Model::~Model() {
2020
cudaStreamDestroy(stream);
21-
for (int i = 0; i < engine->getNbBindings(); i++){
21+
for (int i = 0; i < engine->getNbBindings(); i++) {
2222
CUDA_CHECK(cudaFree(gpu_buffers[i]));
2323
}
2424
};
@@ -91,7 +91,7 @@ bool Model::ReadTrtFile() {
9191
// sample::gLogInfo << "deserialize done" << std::endl;
9292
}
9393

94-
void Model::LoadEngine(){
94+
void Model::LoadEngine() {
9595
// create and load engine
9696
std::fstream existEngine;
9797
existEngine.open(engine_file, std::ios::in);

0 commit comments

Comments
 (0)