Skip to content

Commit 842adee

Browse files
fazligorkembalfazligorkembal
committed
Yolo26 Initialized
1 parent 3276610 commit 842adee

22 files changed

Lines changed: 3423 additions & 0 deletions

‎yolo26/.clang-format‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
# Google C/C++ Code Style settings (with 4-space)
2+
# Refered to https://github.com/kehanXue/google-style-clang-format/blob/master/.clang-format
3+
4+
Language: Cpp
5+
BasedOnStyle: Google
6+
AccessModifierOffset: -1
7+
AlignAfterOpenBracket: Align
8+
AlignConsecutiveAssignments: None
9+
AlignOperands: Align
10+
AllowAllArgumentsOnNextLine: true
11+
AllowAllConstructorInitializersOnNextLine: true
12+
AllowAllParametersOfDeclarationOnNextLine: false
13+
AllowShortBlocksOnASingleLine: Empty
14+
AllowShortCaseLabelsOnASingleLine: false
15+
AllowShortFunctionsOnASingleLine: Inline
16+
AllowShortIfStatementsOnASingleLine: Never # To avoid conflict, set this "Never" and each "if statement" should include brace when coding
17+
AllowShortLambdasOnASingleLine: Inline
18+
AllowShortLoopsOnASingleLine: false
19+
AlwaysBreakAfterReturnType: None
20+
AlwaysBreakTemplateDeclarations: Yes
21+
BinPackArguments: true
22+
BreakBeforeBraces: Custom
23+
BraceWrapping:
24+
AfterCaseLabel: false
25+
AfterClass: false
26+
AfterStruct: false
27+
AfterControlStatement: Never
28+
AfterEnum: false
29+
AfterFunction: false
30+
AfterNamespace: false
31+
AfterUnion: false
32+
AfterExternBlock: false
33+
BeforeCatch: false
34+
BeforeElse: false
35+
BeforeLambdaBody: false
36+
IndentBraces: false
37+
SplitEmptyFunction: false
38+
SplitEmptyRecord: false
39+
SplitEmptyNamespace: false
40+
BreakBeforeBinaryOperators: None
41+
BreakBeforeTernaryOperators: true
42+
BreakConstructorInitializers: BeforeColon
43+
BreakInheritanceList: BeforeColon
44+
ColumnLimit: 120
45+
CompactNamespaces: false
46+
ContinuationIndentWidth: 8
47+
Cpp11BracedListStyle: true
48+
DerivePointerAlignment: false # Make sure the * or & align on the left
49+
EmptyLineBeforeAccessModifier: LogicalBlock
50+
FixNamespaceComments: true
51+
IncludeBlocks: Preserve
52+
IndentCaseLabels: true
53+
IndentPPDirectives: None
54+
IndentWidth: 4
55+
KeepEmptyLinesAtTheStartOfBlocks: true
56+
MaxEmptyLinesToKeep: 1
57+
NamespaceIndentation: None
58+
ObjCSpaceAfterProperty: false
59+
ObjCSpaceBeforeProtocolList: true
60+
PointerAlignment: Left
61+
ReflowComments: false
62+
# SeparateDefinitionBlocks: Always # Only support since clang-format 14
63+
SpaceAfterCStyleCast: false
64+
SpaceAfterLogicalNot: false
65+
SpaceAfterTemplateKeyword: true
66+
SpaceBeforeAssignmentOperators: true
67+
SpaceBeforeCpp11BracedList: false
68+
SpaceBeforeCtorInitializerColon: true
69+
SpaceBeforeInheritanceColon: true
70+
SpaceBeforeParens: ControlStatements
71+
SpaceBeforeRangeBasedForLoopColon: true
72+
SpaceBeforeSquareBrackets: false
73+
SpaceInEmptyParentheses: false
74+
SpacesBeforeTrailingComments: 2
75+
SpacesInAngles: false
76+
SpacesInCStyleCastParentheses: false
77+
SpacesInContainerLiterals: false
78+
SpacesInParentheses: false
79+
SpacesInSquareBrackets: false
80+
Standard: c++11
81+
TabWidth: 8
82+
UseTab: Never

‎yolo26/CMakeLists.txt‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
cmake_minimum_required(VERSION 3.10)
2+
3+
project(yolo26)
4+
5+
add_definitions(-std=c++11)
6+
add_definitions(-DAPI_EXPORTS)
7+
set(CMAKE_CXX_STANDARD 11)
8+
set(CMAKE_BUILD_TYPE Debug)
9+
10+
set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc)
11+
enable_language(CUDA)
12+
13+
include_directories(${PROJECT_SOURCE_DIR}/include)
14+
include_directories(${PROJECT_SOURCE_DIR}/plugin)
15+
16+
# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
17+
if(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64")
18+
message("embed_platform on")
19+
include_directories(/usr/local/cuda/targets/aarch64-linux/include)
20+
link_directories(/usr/local/cuda/targets/aarch64-linux/lib)
21+
else()
22+
message("embed_platform off")
23+
24+
# cuda
25+
include_directories(/usr/local/cuda/include)
26+
link_directories(/usr/local/cuda/lib64)
27+
28+
# tensorrt
29+
include_directories(/workspace/shared/TensorRT-8.6.3/include)
30+
link_directories(/workspace/shared/TensorRT-8.6.3/lib)
31+
endif()
32+
33+
add_library(yololayerplugins SHARED ${PROJECT_SOURCE_DIR}/plugin/yololayer.cu)
34+
target_link_libraries(yololayerplugins nvinfer cudart)
35+
36+
find_package(OpenCV)
37+
include_directories(${OpenCV_INCLUDE_DIRS})
38+
39+
file(GLOB_RECURSE SRCS ${PROJECT_SOURCE_DIR}/src/*.cpp ${PROJECT_SOURCE_DIR}/src/*.cu)
40+
41+
add_executable(yolo26_det ${PROJECT_SOURCE_DIR}/yolo26_det.cpp ${SRCS})
42+
target_link_libraries(yolo26_det nvinfer)
43+
target_link_libraries(yolo26_det cudart)
44+
target_link_libraries(yolo26_det yololayerplugins)
45+
target_link_libraries(yolo26_det ${OpenCV_LIBS})

‎yolo26/README.md‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
## Introduction
2+
3+
Yolo26 model supports TensorRT-8.
4+
5+
Training code [link](https://github.com/ultralytics/ultralytics/archive/refs/tags/v8.4.0.zip)
6+
7+
## Environment
8+
9+
* cuda 12.4
10+
* cudnn 9.1.0.70
11+
* tensorrt 8.6.3
12+
* opencv 4.8.0
13+
* ultralytics 8.4.0
14+
15+
## Support
16+
17+
* [✅] Yolo26n-det support FP32/FP16 and C++ API
18+
19+
## COMING FEATURES
20+
* [⏳] Yolo26s-det, Yolo26m-det, Yolo26l-det, Yolo26sx-det
21+
* [⏳] Windows OS Support
22+
* [⏳] Ssupport Batched Inputs
23+
* [⏳] Support Quantization
24+
* [⏳] Yolo26-cls models
25+
* [⏳] Yolo26-pose models
26+
* [⏳] Yolo26-seg models
27+
* [⏳] Yolo26-obb models
28+
29+
## Config
30+
31+
* Choose the YOLO26 sub-model n/s/m/l/x from command line arguments. (For now only tested n model!)
32+
* Other configs please check [src/config.h](src/config.h)
33+
34+
## Build and Run
35+
36+
1. generate .wts from pytorch with .pt, or download .wts from model zoo
37+
38+
```shell
39+
# Download ultralytics
40+
wget https://github.com/ultralytics/ultralytics/archive/refs/tags/v8.4.0.zip -O ultralytics-8.4.0.zip
41+
# Unzip ultralytics
42+
unzip ultralytics-8.4.0.zip
43+
cd ultralytics-8.4.0
44+
# Download models
45+
wget https://github.com/ultralytics/assets/releases/download/v8.4.0/yolo26n.pt -O yolo26n.pt # to download other models, replace 'yolo26n.pt' with 'yolo26s.pt', 'yolo26m.pt', 'yolo26l.pt' or 'yolo26x.pt'
46+
# Generate .wts
47+
cp [PATH-TO-MAIN-FOLDER]/gen_wts.py .
48+
python gen_wts.py -w yolo26n.pt -o yolo26n.wts -t detect
49+
# A file 'yolo26n.wts' will be generated.
50+
```
51+
52+
2. build tensorrt_yolo26 and run
53+
```shell
54+
cd [PATH-TO-MAIN-FOLDER]
55+
mkdir build
56+
cd build
57+
cmake ..
58+
make
59+
```
60+
61+
### Detection
62+
```shell
63+
cp [PATH-TO-ultralytics]/yolo26n.wts .
64+
# Build and serialize TensorRT engine
65+
./yolo26_det -s yolo26n.wts yolo26n.engine [n/s/m/l/x]
66+
# Run inference
67+
./yolo26_det -d yolo26n.engine ../images [c/g]
68+
# results saved in build directory
69+
```
70+
71+
## More Information
72+
See the readme in [home page.](https://github.com/wang-xinyu/tensorrtx)

‎yolo26/gen_wts.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
import sys # noqa: F401
2+
import argparse
3+
import os
4+
import struct
5+
import torch
6+
7+
8+
def parse_args():
9+
parser = argparse.ArgumentParser(description='Convert .pt file to .wts')
10+
parser.add_argument('-w', '--weights', required=True,
11+
help='Input weights (.pt) file path (required)')
12+
parser.add_argument(
13+
'-o', '--output', help='Output (.wts) file path (optional)')
14+
parser.add_argument(
15+
'-t', '--type', type=str, default='detect', choices=['detect', 'cls', 'seg', 'pose', 'obb'],
16+
help='determines the model is detection/classification')
17+
args = parser.parse_args()
18+
if not os.path.isfile(args.weights):
19+
raise SystemExit('Invalid input file')
20+
if not args.output:
21+
args.output = os.path.splitext(args.weights)[0] + '.wts'
22+
elif os.path.isdir(args.output):
23+
args.output = os.path.join(
24+
args.output,
25+
os.path.splitext(os.path.basename(args.weights))[0] + '.wts')
26+
return args.weights, args.output, args.type
27+
28+
29+
pt_file, wts_file, m_type = parse_args()
30+
31+
print(f'Generating .wts for {m_type} model')
32+
33+
# Load model
34+
print(f'Loading {pt_file}')
35+
36+
# Initialize
37+
device = 'cpu'
38+
39+
# Load model
40+
model = torch.load(pt_file, map_location=device, weights_only=False)['model'].float() # load to FP32
41+
42+
if m_type in ['detect', 'seg', 'pose', 'obb']:
43+
anchor_grid = model.model[-1].anchors * model.model[-1].stride[..., None, None]
44+
45+
delattr(model.model[-1], 'anchors')
46+
47+
model.to(device).eval()
48+
49+
with open(wts_file, 'w') as f:
50+
f.write('{}\n'.format(len(model.state_dict().keys())))
51+
for k, v in model.state_dict().items():
52+
vr = v.reshape(-1).cpu().numpy()
53+
f.write('{} {} '.format(k, len(vr)))
54+
for vv in vr:
55+
f.write(' ')
56+
f.write(struct.pack('>f', float(vv)).hex())
57+
f.write('\n')

‎yolo26/include/block.h‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
#pragma once
2+
3+
#include <map>
4+
#include <string>
5+
#include <vector>
6+
#include "NvInfer.h"
7+
8+
using namespace std;
9+
std::map<std::string, nvinfer1::Weights> loadWeights(const std::string file);
10+
11+
nvinfer1::IScaleLayer* addBatchNorm2d(nvinfer1::INetworkDefinition* network,
12+
std::map<std::string, nvinfer1::Weights> weightMap, nvinfer1::ITensor& input,
13+
std::string lname, float eps);
14+
15+
nvinfer1::IElementWiseLayer* convBnSiLU(nvinfer1::INetworkDefinition* network,
16+
std::map<std::string, nvinfer1::Weights> weightMap, nvinfer1::ITensor& input,
17+
int ch, std::vector<int> k, int s, std::string lname, int g = 1);
18+
19+
nvinfer1::IElementWiseLayer* C3K2(nvinfer1::INetworkDefinition* network,
20+
std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input, int c1,
21+
int c2, int n, bool c3k, bool shortcut, bool atnn, float e, std::string lname);
22+
23+
nvinfer1::IElementWiseLayer* SPPF(nvinfer1::INetworkDefinition* network,
24+
std::map<std::string, nvinfer1::Weights> weightMap, nvinfer1::ITensor& input, int c1,
25+
int c2, int k, bool shortcut, std::string lname);
26+
27+
nvinfer1::IElementWiseLayer* C2PSA(nvinfer1::INetworkDefinition* network,
28+
std::map<std::string, nvinfer1::Weights>& weightMap, nvinfer1::ITensor& input,
29+
int c1, int c2, int n, float e, std::string lname);
30+
31+
nvinfer1::ILayer* DWConv(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights> weightMap,
32+
nvinfer1::ITensor& input, int ch, std::vector<int> k, int s, std::string lname);
33+
34+
nvinfer1::ILayer* conv(nvinfer1::INetworkDefinition* network, std::map<std::string, nvinfer1::Weights> weightMap,
35+
nvinfer1::ITensor& input, int ch, std::vector<int> k, int s, std::string lname, int g = 1,
36+
bool act = true);
37+
38+
nvinfer1::IPluginV2Layer* addYoloLayer(nvinfer1::INetworkDefinition* network, nvinfer1::ITensor& input,
39+
const std::vector<int>& strides, const std::vector<int>& fm_sizes,
40+
int stridesLength, bool is_segmentation, bool is_pose, bool is_obb,
41+
int anchorCount);

‎yolo26/include/config.h‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
#define USE_FP16
2+
// #define USE_FP32
3+
// #define USE_INT8
4+
5+
const static char* kInputTensorName = "images";
6+
const static char* kOutputTensorName = "output";
7+
const static char* kProtoTensorName = "proto";
8+
const static int kNumClass = 80;
9+
const static int kPoseNumClass = 1;
10+
const static int kNumberOfPoints = 17; // number of keypoints total
11+
// obb model's number of classes
12+
constexpr static int kObbNumClass = 15;
13+
const static int kObbNe = 1; // number of extra parameters
14+
const static int kBatchSize = 1;
15+
const static int kGpuId = 0;
16+
const static int kInputH = 640;
17+
const static int kInputW = 640;
18+
const static int kObbInputH = 1024;
19+
const static int kObbInputW = 1024;
20+
const static float kNmsThresh = 0.45f;
21+
const static float kConfThresh = 0.5f;
22+
const static float kConfThreshKeypoints = 0.5f; // keypoints confidence
23+
const static int kMaxInputImageSize = 3000 * 3000;
24+
const static int kMaxNumOutputBbox = 300;
25+
// Quantization input image folder path
26+
const static char* kInputQuantizationFolder = "./coco_calib";
27+
28+
// Classfication model's number of classes
29+
constexpr static int kClsNumClass = 1000;
30+
// Classfication model's input shape
31+
constexpr static int kClsInputH = 224;
32+
constexpr static int kClsInputW = 224;

‎yolo26/include/cuda_utils.h‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
#ifndef TRTX_CUDA_UTILS_H_
2+
#define TRTX_CUDA_UTILS_H_
3+
4+
#include <cuda_runtime_api.h>
5+
6+
#ifndef CUDA_CHECK
7+
#define CUDA_CHECK(callstr) \
8+
{ \
9+
cudaError_t error_code = callstr; \
10+
if (error_code != cudaSuccess) { \
11+
std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \
12+
assert(0); \
13+
} \
14+
}
15+
#endif // CUDA_CHECK
16+
17+
#endif // TRTX_CUDA_UTILS_H_

0 commit comments

Comments
 (0)