New YOLOv5 conversion and support (>= v2.0)
This commit is contained in:
Marcos Luciano
@@ -2,8 +2,8 @@
|
||||
gpu-id=0
|
||||
net-scale-factor=0.0039215697906911373
|
||||
model-color-format=0
|
||||
custom-network-config=yolov5n.cfg
|
||||
model-file=yolov5n.wts
|
||||
custom-network-config=yolov5s.cfg
|
||||
model-file=yolov5s.wts
|
||||
model-engine-file=model_b1_gpu0_fp32.engine
|
||||
#int8-calib-file=calib.table
|
||||
labelfile-path=labels.txt
|
||||
|
||||
110
docs/YOLOR.md
Normal file
110
docs/YOLOR.md
Normal file
@@ -0,0 +1,110 @@
|
||||
# YOLOR usage
|
||||
|
||||
**NOTE**: You need to use the main branch of the YOLOR repo to convert the model.
|
||||
|
||||
**NOTE**: The cfg is required.
|
||||
|
||||
* [Convert model](#convert-model)
|
||||
* [Compile the lib](#compile-the-lib)
|
||||
* [Edit the config_infer_primary_yolor file](#edit-the-config_infer_primary_yolor-file)
|
||||
* [Edit the deepstream_app_config file](#edit-the-deepstream_app_config-file)
|
||||
* [Testing the model](#testing-the-model)
|
||||
|
||||
##
|
||||
|
||||
### Convert model
|
||||
|
||||
#### 1. Download the YOLOR repo and install the requirements
|
||||
|
||||
```
|
||||
git clone https://github.com/WongKinYiu/yolor.git
|
||||
cd yolor
|
||||
pip3 install -r requirements.txt
|
||||
```
|
||||
|
||||
**NOTE**: It is recommended to use a Python virtualenv.
|
||||
|
||||
#### 2. Copy conversor
|
||||
|
||||
Copy the `gen_wts_yolor.py` file from `DeepStream-Yolo/utils` directory to the `yolor` folder.
|
||||
|
||||
#### 3. Download the model
|
||||
|
||||
Download the `pt` file from [YOLOR](https://github.com/WongKinYiu/yolor) repo.
|
||||
|
||||
**NOTE**: You can use your custom model, but it is important to keep the YOLO model reference (`yolor_`) in you `cfg` and `weights`/`wts` filenames to generate the engine correctly.
|
||||
|
||||
#### 4. Convert model
|
||||
|
||||
Generate the `cfg` and `wts` files (example for YOLOR-CSP)
|
||||
|
||||
```
|
||||
python3 gen_wts_yolor.py -w yolor_csp.pt -c cfg/yolor_csp.cfg
|
||||
```
|
||||
|
||||
#### 5. Copy generated files
|
||||
|
||||
Copy the generated `cfg` and `wts` files to the `DeepStream-Yolo` folder
|
||||
|
||||
##
|
||||
|
||||
### Compile the lib
|
||||
|
||||
Open the `DeepStream-Yolo` folder and compile the lib
|
||||
|
||||
* DeepStream 6.1 on x86 platform
|
||||
|
||||
```
|
||||
CUDA_VER=11.6 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
* DeepStream 6.0.1 / 6.0 on x86 platform
|
||||
|
||||
```
|
||||
CUDA_VER=11.4 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
* DeepStream 6.1 on Jetson platform
|
||||
|
||||
```
|
||||
CUDA_VER=11.4 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
* DeepStream 6.0.1 / 6.0 on Jetson platform
|
||||
|
||||
```
|
||||
CUDA_VER=10.2 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
##
|
||||
|
||||
### Edit the config_infer_primary_yolor file
|
||||
|
||||
Edit the `config_infer_primary_yolor.txt` file according to your model (example for YOLOR-CSP)
|
||||
|
||||
```
|
||||
[property]
|
||||
...
|
||||
custom-network-config=yolor_csp.cfg
|
||||
model-file=yolor_csp.wts
|
||||
...
|
||||
```
|
||||
|
||||
##
|
||||
|
||||
### Edit the deepstream_app_config.txt file
|
||||
|
||||
```
|
||||
...
|
||||
[primary-gie]
|
||||
...
|
||||
config-file=config_infer_primary_yolor.txt
|
||||
```
|
||||
|
||||
##
|
||||
|
||||
### Testing the model
|
||||
|
||||
```
|
||||
deepstream-app -c deepstream_app_config.txt
|
||||
```
|
||||
135
docs/YOLOv5.md
Normal file
135
docs/YOLOv5.md
Normal file
@@ -0,0 +1,135 @@
|
||||
# YOLOv5 usage
|
||||
|
||||
**NOTE**: You can use the main branch of the YOLOv5 repo to convert all model versions.
|
||||
|
||||
**NOTE**: The yaml is not required.
|
||||
|
||||
* [Convert model](#convert-model)
|
||||
* [Compile the lib](#compile-the-lib)
|
||||
* [Edit the config_infer_primary_yoloV5 file](#edit-the-config_infer_primary_yolov5-file)
|
||||
* [Edit the deepstream_app_config file](#edit-the-deepstream_app_config-file)
|
||||
* [Testing the model](#testing-the-model)
|
||||
|
||||
##
|
||||
|
||||
### Convert model
|
||||
|
||||
#### 1. Download the YOLOv5 repo and install the requirements
|
||||
|
||||
```
|
||||
git clone https://github.com/ultralytics/yolov5.git
|
||||
cd yolov5
|
||||
pip3 install -r requirements.txt
|
||||
```
|
||||
|
||||
**NOTE**: It is recommended to use a Python virtualenv.
|
||||
|
||||
#### 2. Copy conversor
|
||||
|
||||
Copy the `gen_wts_yoloV5.py` file from `DeepStream-Yolo/utils` directory to the `yolov5` folder.
|
||||
|
||||
#### 3. Download the model
|
||||
|
||||
Download the `pt` file from [YOLOv5](https://github.com/ultralytics/yolov5/releases/) releases (example for YOLOv5s 6.1)
|
||||
|
||||
```
|
||||
wget https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5s.pt
|
||||
```
|
||||
|
||||
**NOTE**: You can use your custom model, but it is important to keep the YOLO model reference (`yolov5_`) in you `cfg` and `weights`/`wts` filenames to generate the engine correctly.
|
||||
|
||||
#### 4. Convert model
|
||||
|
||||
Generate the `cfg` and `wts` files (example for YOLOv5s)
|
||||
|
||||
```
|
||||
python3 gen_wts_yoloV5.py -w yolov5s.pt
|
||||
```
|
||||
|
||||
**NOTE**: To change the inference size (defaut: 640)
|
||||
|
||||
```
|
||||
-s SIZE
|
||||
--size SIZE
|
||||
-s HEIGHT WIDTH
|
||||
--size HEIGHT WIDTH
|
||||
```
|
||||
|
||||
Example for 1280
|
||||
|
||||
```
|
||||
-s 1280
|
||||
```
|
||||
|
||||
or
|
||||
|
||||
```
|
||||
-s 1280 1280
|
||||
```
|
||||
|
||||
#### 5. Copy generated files
|
||||
|
||||
Copy the generated `cfg` and `wts` files to the `DeepStream-Yolo` folder.
|
||||
|
||||
##
|
||||
|
||||
### Compile the lib
|
||||
|
||||
Open the `DeepStream-Yolo` folder and compile the lib
|
||||
|
||||
* DeepStream 6.1 on x86 platform
|
||||
|
||||
```
|
||||
CUDA_VER=11.6 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
* DeepStream 6.0.1 / 6.0 on x86 platform
|
||||
|
||||
```
|
||||
CUDA_VER=11.4 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
* DeepStream 6.1 on Jetson platform
|
||||
|
||||
```
|
||||
CUDA_VER=11.4 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
* DeepStream 6.0.1 / 6.0 on Jetson platform
|
||||
|
||||
```
|
||||
CUDA_VER=10.2 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
##
|
||||
|
||||
### Edit the config_infer_primary_yoloV5 file
|
||||
|
||||
Edit the `config_infer_primary_yoloV5.txt` file according to your model (example for YOLOv5s)
|
||||
|
||||
```
|
||||
[property]
|
||||
...
|
||||
custom-network-config=yolov5s.cfg
|
||||
model-file=yolov5s.wts
|
||||
...
|
||||
```
|
||||
|
||||
##
|
||||
|
||||
### Edit the deepstream_app_config.txt file
|
||||
|
||||
```
|
||||
...
|
||||
[primary-gie]
|
||||
...
|
||||
config-file=config_infer_primary_yoloV5.txt
|
||||
```
|
||||
|
||||
##
|
||||
|
||||
### Testing the model
|
||||
|
||||
```
|
||||
deepstream-app -c deepstream_app_config.txt
|
||||
```
|
||||
@@ -21,7 +21,7 @@ cd DeepStream-Yolo
|
||||
|
||||
#### 3. Copy the `cfg` and `weights`/`wts` files to DeepStream-Yolo folder
|
||||
|
||||
**NOTE**: It's important to keep the YOLO model reference (`yolov4_`, `yolov5_`, `yolor_`, etc) in you `cfg` and `weights`/`wts` files to generate the engine correctly.
|
||||
**NOTE**: It is important to keep the YOLO model reference (`yolov4_`, `yolov5_`, `yolor_`, etc) in you `cfg` and `weights`/`wts` filenames to generate the engine correctly.
|
||||
|
||||
##
|
||||
|
||||
|
||||
@@ -28,7 +28,7 @@ cd DeepStream-Yolo
|
||||
|
||||
#### 4. Copy the `cfg` and `weights`/`wts` files to each GIE folder
|
||||
|
||||
**NOTE**: It's important to keep the YOLO model reference (`yolov4_`, `yolov5_`, `yolor_`, etc) in you `cfg` and `weights`/`wts` files to generate the engine correctly.
|
||||
**NOTE**: It is important to keep the YOLO model reference (`yolov4_`, `yolov5_`, `yolor_`, etc) in you `cfg` and `weights`/`wts` filenames to generate the engine correctly.
|
||||
|
||||
##
|
||||
|
||||
|
||||
@@ -53,6 +53,7 @@ SRCFILES:= nvdsinfer_yolo_engine.cpp \
|
||||
nvdsparsebbox_Yolo.cpp \
|
||||
yoloPlugins.cpp \
|
||||
layers/convolutional_layer.cpp \
|
||||
layers/batchnorm_layer.cpp \
|
||||
layers/implicit_layer.cpp \
|
||||
layers/channels_layer.cpp \
|
||||
layers/shortcut_layer.cpp \
|
||||
|
||||
@@ -12,7 +12,8 @@ nvinfer1::ILayer* activationLayer(
|
||||
nvinfer1::ITensor* input,
|
||||
nvinfer1::INetworkDefinition* network)
|
||||
{
|
||||
if (activation == "linear") {
|
||||
if (activation == "linear")
|
||||
{
|
||||
// Pass
|
||||
}
|
||||
else if (activation == "relu")
|
||||
@@ -46,8 +47,8 @@ nvinfer1::ILayer* activationLayer(
|
||||
{
|
||||
nvinfer1::IActivationLayer* leaky = network->addActivation(
|
||||
*input, nvinfer1::ActivationType::kLEAKY_RELU);
|
||||
leaky->setAlpha(0.1);
|
||||
assert(leaky != nullptr);
|
||||
leaky->setAlpha(0.1);
|
||||
std::string leakyLayerName = "leaky_" + std::to_string(layerIdx);
|
||||
leaky->setName(leakyLayerName.c_str());
|
||||
output = leaky;
|
||||
@@ -74,7 +75,7 @@ nvinfer1::ILayer* activationLayer(
|
||||
std::string tanhLayerName = "tanh_" + std::to_string(layerIdx);
|
||||
tanh->setName(tanhLayerName.c_str());
|
||||
nvinfer1::IElementWiseLayer* mish = network->addElementWise(
|
||||
*tanh->getOutput(0), *input,
|
||||
*input, *tanh->getOutput(0),
|
||||
nvinfer1::ElementWiseOperation::kPROD);
|
||||
assert(mish != nullptr);
|
||||
std::string mishLayerName = "mish_" + std::to_string(layerIdx);
|
||||
@@ -89,14 +90,32 @@ nvinfer1::ILayer* activationLayer(
|
||||
std::string sigmoidLayerName = "sigmoid_" + std::to_string(layerIdx);
|
||||
sigmoid->setName(sigmoidLayerName.c_str());
|
||||
nvinfer1::IElementWiseLayer* silu = network->addElementWise(
|
||||
*sigmoid->getOutput(0), *input,
|
||||
*input, *sigmoid->getOutput(0),
|
||||
nvinfer1::ElementWiseOperation::kPROD);
|
||||
assert(silu != nullptr);
|
||||
std::string siluLayerName = "silu_" + std::to_string(layerIdx);
|
||||
silu->setName(siluLayerName.c_str());
|
||||
output = silu;
|
||||
}
|
||||
else {
|
||||
else if (activation == "hardswish")
|
||||
{
|
||||
nvinfer1::IActivationLayer* hard_sigmoid = network->addActivation(
|
||||
*input, nvinfer1::ActivationType::kHARD_SIGMOID);
|
||||
assert(hard_sigmoid != nullptr);
|
||||
hard_sigmoid->setAlpha(1.0 / 6.0);
|
||||
hard_sigmoid->setBeta(0.5);
|
||||
std::string hardSigmoidLayerName = "hard_sigmoid_" + std::to_string(layerIdx);
|
||||
hard_sigmoid->setName(hardSigmoidLayerName.c_str());
|
||||
nvinfer1::IElementWiseLayer* hard_swish = network->addElementWise(
|
||||
*input, *hard_sigmoid->getOutput(0),
|
||||
nvinfer1::ElementWiseOperation::kPROD);
|
||||
assert(hard_swish != nullptr);
|
||||
std::string hardSwishLayerName = "hard_swish_" + std::to_string(layerIdx);
|
||||
hard_swish->setName(hardSwishLayerName.c_str());
|
||||
output = hard_swish;
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cerr << "Activation not supported: " << activation << std::endl;
|
||||
std::abort();
|
||||
}
|
||||
|
||||
114
nvdsinfer_custom_impl_Yolo/layers/batchnorm_layer.cpp
Normal file
114
nvdsinfer_custom_impl_Yolo/layers/batchnorm_layer.cpp
Normal file
@@ -0,0 +1,114 @@
|
||||
/*
|
||||
* Created by Marcos Luciano
|
||||
* https://www.github.com/marcoslucianops
|
||||
*/
|
||||
|
||||
#include <math.h>
|
||||
#include "batchnorm_layer.h"
|
||||
|
||||
nvinfer1::ILayer* batchnormLayer(
|
||||
int layerIdx,
|
||||
std::map<std::string, std::string>& block,
|
||||
std::vector<float>& weights,
|
||||
std::vector<nvinfer1::Weights>& trtWeights,
|
||||
int& weightPtr,
|
||||
std::string weightsType,
|
||||
float eps,
|
||||
nvinfer1::ITensor* input,
|
||||
nvinfer1::INetworkDefinition* network)
|
||||
{
|
||||
assert(block.at("type") == "batchnorm");
|
||||
assert(block.find("filters") != block.end());
|
||||
|
||||
int filters = std::stoi(block.at("filters"));
|
||||
std::string activation = block.at("activation");
|
||||
|
||||
std::vector<float> bnBiases;
|
||||
std::vector<float> bnWeights;
|
||||
std::vector<float> bnRunningMean;
|
||||
std::vector<float> bnRunningVar;
|
||||
|
||||
if (weightsType == "weights") {
|
||||
for (int i = 0; i < filters; ++i)
|
||||
{
|
||||
bnBiases.push_back(weights[weightPtr]);
|
||||
weightPtr++;
|
||||
}
|
||||
for (int i = 0; i < filters; ++i)
|
||||
{
|
||||
bnWeights.push_back(weights[weightPtr]);
|
||||
weightPtr++;
|
||||
}
|
||||
for (int i = 0; i < filters; ++i)
|
||||
{
|
||||
bnRunningMean.push_back(weights[weightPtr]);
|
||||
weightPtr++;
|
||||
}
|
||||
for (int i = 0; i < filters; ++i)
|
||||
{
|
||||
bnRunningVar.push_back(sqrt(weights[weightPtr] + 1.0e-5));
|
||||
weightPtr++;
|
||||
}
|
||||
}
|
||||
else {
|
||||
for (int i = 0; i < filters; ++i)
|
||||
{
|
||||
bnWeights.push_back(weights[weightPtr]);
|
||||
weightPtr++;
|
||||
}
|
||||
for (int i = 0; i < filters; ++i)
|
||||
{
|
||||
bnBiases.push_back(weights[weightPtr]);
|
||||
weightPtr++;
|
||||
}
|
||||
for (int i = 0; i < filters; ++i)
|
||||
{
|
||||
bnRunningMean.push_back(weights[weightPtr]);
|
||||
weightPtr++;
|
||||
}
|
||||
for (int i = 0; i < filters; ++i)
|
||||
{
|
||||
bnRunningVar.push_back(sqrt(weights[weightPtr] + eps));
|
||||
weightPtr++;
|
||||
}
|
||||
}
|
||||
|
||||
int size = filters;
|
||||
nvinfer1::Weights shift{nvinfer1::DataType::kFLOAT, nullptr, size};
|
||||
nvinfer1::Weights scale{nvinfer1::DataType::kFLOAT, nullptr, size};
|
||||
nvinfer1::Weights power{nvinfer1::DataType::kFLOAT, nullptr, size};
|
||||
float* shiftWt = new float[size];
|
||||
for (int i = 0; i < size; ++i)
|
||||
{
|
||||
shiftWt[i]
|
||||
= bnBiases.at(i) - ((bnRunningMean.at(i) * bnWeights.at(i)) / bnRunningVar.at(i));
|
||||
}
|
||||
shift.values = shiftWt;
|
||||
float* scaleWt = new float[size];
|
||||
for (int i = 0; i < size; ++i)
|
||||
{
|
||||
scaleWt[i] = bnWeights.at(i) / bnRunningVar[i];
|
||||
}
|
||||
scale.values = scaleWt;
|
||||
float* powerWt = new float[size];
|
||||
for (int i = 0; i < size; ++i)
|
||||
{
|
||||
powerWt[i] = 1.0;
|
||||
}
|
||||
power.values = powerWt;
|
||||
trtWeights.push_back(shift);
|
||||
trtWeights.push_back(scale);
|
||||
trtWeights.push_back(power);
|
||||
|
||||
nvinfer1::IScaleLayer* bn = network->addScale(
|
||||
*input, nvinfer1::ScaleMode::kCHANNEL, shift, scale, power);
|
||||
assert(bn != nullptr);
|
||||
std::string bnLayerName = "batch_norm_" + std::to_string(layerIdx);
|
||||
bn->setName(bnLayerName.c_str());
|
||||
nvinfer1::ILayer* output = bn;
|
||||
|
||||
output = activationLayer(layerIdx, activation, output, output->getOutput(0), network);
|
||||
assert(output != nullptr);
|
||||
|
||||
return output;
|
||||
}
|
||||
27
nvdsinfer_custom_impl_Yolo/layers/batchnorm_layer.h
Normal file
27
nvdsinfer_custom_impl_Yolo/layers/batchnorm_layer.h
Normal file
@@ -0,0 +1,27 @@
|
||||
/*
|
||||
* Created by Marcos Luciano
|
||||
* https://www.github.com/marcoslucianops
|
||||
*/
|
||||
|
||||
#ifndef __BATCHNORM_LAYER_H__
|
||||
#define __BATCHNORM_LAYER_H__
|
||||
|
||||
#include <map>
|
||||
#include <vector>
|
||||
|
||||
#include "NvInfer.h"
|
||||
|
||||
#include "activation_layer.h"
|
||||
|
||||
nvinfer1::ILayer* batchnormLayer(
|
||||
int layerIdx,
|
||||
std::map<std::string, std::string>& block,
|
||||
std::vector<float>& weights,
|
||||
std::vector<nvinfer1::Weights>& trtWeights,
|
||||
int& weightPtr,
|
||||
std::string weightsType,
|
||||
float eps,
|
||||
nvinfer1::ITensor* input,
|
||||
nvinfer1::INetworkDefinition* network);
|
||||
|
||||
#endif
|
||||
@@ -44,6 +44,11 @@ nvinfer1::ILayer* convolutionalLayer(
|
||||
groups = std::stoi(block.at("groups"));
|
||||
}
|
||||
|
||||
if (block.find("bias") != block.end())
|
||||
{
|
||||
bias = std::stoi(block.at("bias"));
|
||||
}
|
||||
|
||||
int pad;
|
||||
if (padding)
|
||||
pad = (kernelSize - 1) / 2;
|
||||
@@ -61,7 +66,9 @@ nvinfer1::ILayer* convolutionalLayer(
|
||||
if (weightsType == "weights") {
|
||||
if (batchNormalize == false)
|
||||
{
|
||||
float* val = new float[filters];
|
||||
float* val;
|
||||
if (bias != 0) {
|
||||
val = new float[filters];
|
||||
for (int i = 0; i < filters; ++i)
|
||||
{
|
||||
val[i] = weights[weightPtr];
|
||||
@@ -69,6 +76,7 @@ nvinfer1::ILayer* convolutionalLayer(
|
||||
}
|
||||
convBias.values = val;
|
||||
trtWeights.push_back(convBias);
|
||||
}
|
||||
val = new float[size];
|
||||
for (int i = 0; i < size; ++i)
|
||||
{
|
||||
@@ -108,6 +116,7 @@ nvinfer1::ILayer* convolutionalLayer(
|
||||
}
|
||||
convWt.values = val;
|
||||
trtWeights.push_back(convWt);
|
||||
if (bias != 0)
|
||||
trtWeights.push_back(convBias);
|
||||
}
|
||||
}
|
||||
@@ -122,6 +131,7 @@ nvinfer1::ILayer* convolutionalLayer(
|
||||
}
|
||||
convWt.values = val;
|
||||
trtWeights.push_back(convWt);
|
||||
if (bias != 0) {
|
||||
val = new float[filters];
|
||||
for (int i = 0; i < filters; ++i)
|
||||
{
|
||||
@@ -131,6 +141,7 @@ nvinfer1::ILayer* convolutionalLayer(
|
||||
convBias.values = val;
|
||||
trtWeights.push_back(convBias);
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
float* val = new float[size];
|
||||
@@ -161,6 +172,7 @@ nvinfer1::ILayer* convolutionalLayer(
|
||||
weightPtr++;
|
||||
}
|
||||
trtWeights.push_back(convWt);
|
||||
if (bias != 0)
|
||||
trtWeights.push_back(convBias);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -19,11 +19,11 @@ nvinfer1::ILayer* maxpoolLayer(
|
||||
int stride = std::stoi(block.at("stride"));
|
||||
|
||||
nvinfer1::IPoolingLayer* pool
|
||||
= network->addPoolingNd(*input, nvinfer1::PoolingType::kMAX, nvinfer1::DimsHW{size, size});
|
||||
= network->addPoolingNd(*input, nvinfer1::PoolingType::kMAX, nvinfer1::Dims{2, {size, size}});
|
||||
assert(pool);
|
||||
std::string maxpoolLayerName = "maxpool_" + std::to_string(layerIdx);
|
||||
pool->setStrideNd(nvinfer1::DimsHW{stride, stride});
|
||||
pool->setPaddingMode(nvinfer1::PaddingMode::kSAME_UPPER);
|
||||
pool->setStrideNd(nvinfer1::Dims{2, {stride, stride}});
|
||||
pool->setPaddingNd(nvinfer1::Dims{2, {size / 2, size / 2}});
|
||||
pool->setName(maxpoolLayerName.c_str());
|
||||
|
||||
return pool;
|
||||
|
||||
@@ -207,6 +207,20 @@ NvDsInferStatus Yolo::buildYoloNetwork(std::vector<float>& weights, nvinfer1::IN
|
||||
printLayerInfo(layerIndex, layerType, inputVol, outputVol, std::to_string(weightPtr));
|
||||
}
|
||||
|
||||
else if (m_ConfigBlocks.at(i).at("type") == "batchnorm")
|
||||
{
|
||||
std::string inputVol = dimsToString(previous->getDimensions());
|
||||
nvinfer1::ILayer* out = batchnormLayer(
|
||||
i, m_ConfigBlocks.at(i), weights, m_TrtWeights, weightPtr, weightsType, eps, previous, &network);
|
||||
previous = out->getOutput(0);
|
||||
assert(previous != nullptr);
|
||||
channels = getNumChannels(previous);
|
||||
std::string outputVol = dimsToString(previous->getDimensions());
|
||||
tensorOutputs.push_back(previous);
|
||||
std::string layerType = "bn_" + m_ConfigBlocks.at(i).at("activation");
|
||||
printLayerInfo(layerIndex, layerType, inputVol, outputVol, std::to_string(weightPtr));
|
||||
}
|
||||
|
||||
else if (m_ConfigBlocks.at(i).at("type") == "implicit_add" || m_ConfigBlocks.at(i).at("type") == "implicit_mul")
|
||||
{
|
||||
std::string type;
|
||||
|
||||
@@ -27,6 +27,7 @@
|
||||
#define _YOLO_H_
|
||||
|
||||
#include "layers/convolutional_layer.h"
|
||||
#include "layers/batchnorm_layer.h"
|
||||
#include "layers/implicit_layer.h"
|
||||
#include "layers/channels_layer.h"
|
||||
#include "layers/shortcut_layer.h"
|
||||
|
||||
201
readme.md
201
readme.md
@@ -9,24 +9,25 @@ NVIDIA DeepStream SDK 6.1 / 6.0.1 / 6.0 configuration for YOLO models
|
||||
* YOLOX support
|
||||
* PP-YOLO support
|
||||
* YOLOv6 support
|
||||
* YOLOv7 support
|
||||
* Dynamic batch-size
|
||||
|
||||
### Improvements on this repository
|
||||
|
||||
* Darknet cfg params parser (no need to edit `nvdsparsebbox_Yolo.cpp` or other files)
|
||||
* Support for `new_coords`, `beta_nms` and `scale_x_y` params
|
||||
* Support for `new_coords` and `scale_x_y` params
|
||||
* Support for new models
|
||||
* Support for new layers
|
||||
* Support for new activations
|
||||
* Support for convolutional groups
|
||||
* Support for INT8 calibration
|
||||
* Support for non square models
|
||||
* Support for `reorg`, `implicit` and `channel` layers (YOLOR)
|
||||
* YOLOv5 4.0, 5.0, 6.0 and 6.1 support
|
||||
* YOLOR support
|
||||
* **GPU YOLO Decoder (moved from CPU to GPU to get better performance)** [#138](https://github.com/marcoslucianops/DeepStream-Yolo/issues/138)
|
||||
* New documentation for multiple models
|
||||
* **YOLOv5 >= 2.0 support**
|
||||
* **YOLOR support**
|
||||
* **GPU YOLO Decoder** [#138](https://github.com/marcoslucianops/DeepStream-Yolo/issues/138)
|
||||
* **GPU Batched NMS** [#142](https://github.com/marcoslucianops/DeepStream-Yolo/issues/142)
|
||||
* **New documentation for multiple models**
|
||||
* **New YOLOv5 conversion**
|
||||
|
||||
##
|
||||
|
||||
@@ -37,10 +38,10 @@ NVIDIA DeepStream SDK 6.1 / 6.0.1 / 6.0 configuration for YOLO models
|
||||
* [Benchmarks](#benchmarks)
|
||||
* [dGPU installation](#dgpu-installation)
|
||||
* [Basic usage](#basic-usage)
|
||||
* [YOLOv5 usage](#yolov5-usage)
|
||||
* [YOLOR usage](#yolor-usage)
|
||||
* [NMS configuration](#nms-configuration)
|
||||
* [INT8 calibration](#int8-calibration)
|
||||
* [YOLOv5 usage](#docs/YOLOv5.md)
|
||||
* [YOLOR usage](#docs/YOLOR.md)
|
||||
* [Using your custom model](docs/customModels.md)
|
||||
* [Multiple YOLO GIEs](docs/multipleGIEs.md)
|
||||
|
||||
@@ -95,7 +96,7 @@ NVIDIA DeepStream SDK 6.1 / 6.0.1 / 6.0 configuration for YOLO models
|
||||
### Tested models
|
||||
|
||||
* [Darknet YOLO](https://github.com/AlexeyAB/darknet)
|
||||
* [YOLOv5 4.0, 5.0, 6.0 and 6.1](https://github.com/ultralytics/yolov5)
|
||||
* [YOLOv5 >= 2.0](https://github.com/ultralytics/yolov5)
|
||||
* [YOLOR](https://github.com/WongKinYiu/yolor)
|
||||
* [MobileNet-YOLO](https://github.com/dog-qiuqiu/MobileNet-Yolo)
|
||||
* [YOLO-Fastest](https://github.com/dog-qiuqiu/Yolo-Fastest)
|
||||
@@ -448,188 +449,6 @@ config-file=config_infer_primary_yoloV2.txt
|
||||
|
||||
##
|
||||
|
||||
### YOLOv5 usage
|
||||
|
||||
**NOTE**: Make sure to change the YOLOv5 repo version according to your model version before the conversion.
|
||||
|
||||
#### 1. Copy the `gen_wts_yoloV5.py` file from `DeepStream-Yolo/utils` directory to the [YOLOv5](https://github.com/ultralytics/yolov5) folder
|
||||
|
||||
#### 2. Open the YOLOv5 folder
|
||||
|
||||
#### 3. Download the `pt` file from [YOLOv5](https://github.com/ultralytics/yolov5/releases/) repo (example for YOLOv5n 6.1)
|
||||
|
||||
```
|
||||
wget https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5n.pt
|
||||
```
|
||||
|
||||
#### 4. Generate the `cfg` and `wts` files (example for YOLOv5n)
|
||||
|
||||
```
|
||||
python3 gen_wts_yoloV5.py -w yolov5n.pt -c models/yolov5n.yaml
|
||||
```
|
||||
|
||||
#### 5. Copy the generated `cfg` and `wts` files to the DeepStream-Yolo folder
|
||||
|
||||
#### 6. Open the DeepStream-Yolo folder
|
||||
|
||||
#### 7. Compile the lib
|
||||
|
||||
* DeepStream 6.1 on x86 platform
|
||||
|
||||
```
|
||||
CUDA_VER=11.6 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
* DeepStream 6.0.1 / 6.0 on x86 platform
|
||||
|
||||
```
|
||||
CUDA_VER=11.4 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
* DeepStream 6.1 on Jetson platform
|
||||
|
||||
```
|
||||
CUDA_VER=11.4 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
* DeepStream 6.0.1 / 6.0 on Jetson platform
|
||||
|
||||
```
|
||||
CUDA_VER=10.2 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
#### 8. Edit the `config_infer_primary_yoloV5.txt` file according to your model (example for YOLOv5n)
|
||||
|
||||
```
|
||||
[property]
|
||||
...
|
||||
custom-network-config=yolov5n.cfg
|
||||
model-file=yolov5n.wts
|
||||
...
|
||||
```
|
||||
|
||||
#### 9. Edit the `deepstream_app_config.txt` file
|
||||
|
||||
```
|
||||
...
|
||||
[primary-gie]
|
||||
...
|
||||
config-file=config_infer_primary_yoloV5.txt
|
||||
```
|
||||
|
||||
#### 10. Run
|
||||
|
||||
```
|
||||
deepstream-app -c deepstream_app_config.txt
|
||||
```
|
||||
|
||||
**NOTE**: For YOLOv5 P6, check the `gen_wts_yoloV5.py` file args and set them according to your model.
|
||||
|
||||
* Input weights (.pt) file path
|
||||
|
||||
```
|
||||
-w or --weights
|
||||
```
|
||||
|
||||
* Input cfg (.yaml) file path
|
||||
|
||||
```
|
||||
-c or --yaml
|
||||
```
|
||||
|
||||
* Inference size [size] or [height , weight]
|
||||
|
||||
Default: 640 / 1280 (if --p6)
|
||||
|
||||
```
|
||||
-s or --size
|
||||
```
|
||||
|
||||
* Example for 1280
|
||||
|
||||
```
|
||||
-s 1280
|
||||
```
|
||||
|
||||
or
|
||||
|
||||
```
|
||||
-s 1280 1280
|
||||
```
|
||||
|
||||
##
|
||||
|
||||
### YOLOR usage
|
||||
|
||||
#### 1. Copy the `gen_wts_yolor.py` file from `DeepStream-Yolo/utils` directory to the [YOLOR](https://github.com/WongKinYiu/yolor) folder
|
||||
|
||||
#### 2. Open the YOLOR folder
|
||||
|
||||
#### 3. Download the `pt` file from [YOLOR](https://github.com/WongKinYiu/yolor) repo
|
||||
|
||||
#### 4. Generate the `cfg` and `wts` files (example for YOLOR-CSP)
|
||||
|
||||
```
|
||||
python3 gen_wts_yolor.py -w yolor_csp.pt -c cfg/yolor_csp.cfg
|
||||
```
|
||||
|
||||
#### 5. Copy the generated `cfg` and `wts` files to the DeepStream-Yolo folder
|
||||
|
||||
#### 6. Open the DeepStream-Yolo folder
|
||||
|
||||
#### 7. Compile the lib
|
||||
|
||||
* DeepStream 6.1 on x86 platform
|
||||
|
||||
```
|
||||
CUDA_VER=11.6 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
* DeepStream 6.0.1 / 6.0 on x86 platform
|
||||
|
||||
```
|
||||
CUDA_VER=11.4 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
* DeepStream 6.1 on Jetson platform
|
||||
|
||||
```
|
||||
CUDA_VER=11.4 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
* DeepStream 6.0.1 / 6.0 on Jetson platform
|
||||
|
||||
```
|
||||
CUDA_VER=10.2 make -C nvdsinfer_custom_impl_Yolo
|
||||
```
|
||||
|
||||
#### 8. Edit the `config_infer_primary_yolor.txt` file according to your model (example for YOLOR-CSP)
|
||||
|
||||
```
|
||||
[property]
|
||||
...
|
||||
custom-network-config=yolor_csp.cfg
|
||||
model-file=yolor_csp.wts
|
||||
...
|
||||
```
|
||||
|
||||
#### 9. Edit the `deepstream_app_config.txt` file
|
||||
|
||||
```
|
||||
...
|
||||
[primary-gie]
|
||||
...
|
||||
config-file=config_infer_primary_yolor.txt
|
||||
```
|
||||
|
||||
#### 10. Run
|
||||
|
||||
```
|
||||
deepstream-app -c deepstream_app_config.txt
|
||||
```
|
||||
|
||||
##
|
||||
|
||||
### NMS Configuration
|
||||
|
||||
To change the `iou-threshold`, `score-threshold` and `topk` values, modify the `config_nms.txt` file and regenerate the model engine file.
|
||||
|
||||
@@ -1,108 +1,305 @@
|
||||
import argparse
|
||||
import yaml
|
||||
import math
|
||||
import os
|
||||
import struct
|
||||
import torch
|
||||
from utils.torch_utils import select_device
|
||||
|
||||
|
||||
class YoloLayers():
|
||||
def get_route(self, n, layers):
|
||||
route = 0
|
||||
for i, layer in enumerate(layers):
|
||||
if i <= n:
|
||||
route += layer[1]
|
||||
else:
|
||||
break
|
||||
return route
|
||||
class Layers(object):
|
||||
def __init__(self, n, size, fw, fc):
|
||||
self.blocks = [0 for _ in range(n)]
|
||||
self.current = 0
|
||||
|
||||
def route(self, layers=''):
|
||||
return '\n[route]\n' + \
|
||||
'layers=%s\n' % layers
|
||||
self.width = size[0] if len(size) == 1 else size[1]
|
||||
self.height = size[0]
|
||||
|
||||
self.num = 0
|
||||
self.nc = 0
|
||||
self.anchors = ''
|
||||
self.masks = []
|
||||
|
||||
self.fw = fw
|
||||
self.fc = fc
|
||||
self.wc = 0
|
||||
|
||||
self.net()
|
||||
|
||||
def Focus(self, child):
|
||||
self.current = child.i
|
||||
self.fc.write('\n# Focus\n')
|
||||
|
||||
self.reorg()
|
||||
self.convolutional(child.conv)
|
||||
|
||||
def Conv(self, child):
|
||||
self.current = child.i
|
||||
self.fc.write('\n# Conv\n')
|
||||
|
||||
self.convolutional(child)
|
||||
|
||||
def BottleneckCSP(self, child):
|
||||
self.current = child.i
|
||||
self.fc.write('\n# BottleneckCSP\n')
|
||||
|
||||
self.convolutional(child.cv2)
|
||||
self.route('-2')
|
||||
self.convolutional(child.cv1)
|
||||
idx = -3
|
||||
for m in child.m:
|
||||
if m.add:
|
||||
self.convolutional(m.cv1)
|
||||
self.convolutional(m.cv2)
|
||||
self.shortcut(-3)
|
||||
idx -= 3
|
||||
else:
|
||||
self.convolutional(m.cv1)
|
||||
self.convolutional(m.cv2)
|
||||
idx -= 2
|
||||
self.convolutional(child.cv3)
|
||||
self.route('-1, %d' % (idx - 1))
|
||||
self.batchnorm(child.bn, child.act)
|
||||
self.convolutional(child.cv4)
|
||||
|
||||
def C3(self, child):
|
||||
self.current = child.i
|
||||
self.fc.write('\n# C3\n')
|
||||
|
||||
self.convolutional(child.cv2)
|
||||
self.route('-2')
|
||||
self.convolutional(child.cv1)
|
||||
idx = -3
|
||||
for m in child.m:
|
||||
if m.add:
|
||||
self.convolutional(m.cv1)
|
||||
self.convolutional(m.cv2)
|
||||
self.shortcut(-3)
|
||||
idx -= 3
|
||||
else:
|
||||
self.convolutional(m.cv1)
|
||||
self.convolutional(m.cv2)
|
||||
idx -= 2
|
||||
self.route('-1, %d' % idx)
|
||||
self.convolutional(child.cv3)
|
||||
|
||||
def SPP(self, child):
|
||||
self.current = child.i
|
||||
self.fc.write('\n# SPP\n')
|
||||
|
||||
self.convolutional(child.cv1)
|
||||
self.maxpool(child.m[0])
|
||||
self.route('-2')
|
||||
self.maxpool(child.m[1])
|
||||
self.route('-4')
|
||||
self.maxpool(child.m[2])
|
||||
self.route('-6, -5, -3, -1')
|
||||
self.convolutional(child.cv2)
|
||||
|
||||
def SPPF(self, child):
|
||||
self.current = child.i
|
||||
self.fc.write('\n# SPPF\n')
|
||||
|
||||
self.convolutional(child.cv1)
|
||||
self.maxpool(child.m)
|
||||
self.maxpool(child.m)
|
||||
self.maxpool(child.m)
|
||||
self.route('-4, -3, -2, -1')
|
||||
self.convolutional(child.cv2)
|
||||
|
||||
def Upsample(self, child):
|
||||
self.current = child.i
|
||||
self.fc.write('\n# Upsample\n')
|
||||
|
||||
self.upsample(child)
|
||||
|
||||
def Concat(self, child):
|
||||
self.current = child.i
|
||||
self.fc.write('\n# Concat\n')
|
||||
|
||||
r = self.get_route(child.f[1])
|
||||
self.route('-1, %d' % (r - 1))
|
||||
|
||||
def Detect(self, child):
|
||||
self.current = child.i
|
||||
self.fc.write('\n# Detect\n')
|
||||
|
||||
self.get_anchors(child.state_dict(), child.m[0].out_channels)
|
||||
|
||||
for i, m in enumerate(child.m):
|
||||
r = self.get_route(child.f[i])
|
||||
self.route('%d' % (r - 1))
|
||||
self.convolutional(m, detect=True)
|
||||
self.yolo(i)
|
||||
|
||||
def net(self):
|
||||
self.fc.write('[net]\n' +
|
||||
'width=%d\n' % self.width +
|
||||
'height=%d\n' % self.height +
|
||||
'channels=3\n' +
|
||||
'letter_box=1\n')
|
||||
|
||||
def reorg(self):
|
||||
return '\n[reorg]\n'
|
||||
self.blocks[self.current] += 1
|
||||
|
||||
def shortcut(self, route=-1, activation='linear'):
|
||||
return '\n[shortcut]\n' + \
|
||||
'from=%d\n' % route + \
|
||||
'activation=%s\n' % activation
|
||||
self.fc.write('\n[reorg]\n')
|
||||
|
||||
def maxpool(self, stride=1, size=1):
|
||||
return '\n[maxpool]\n' + \
|
||||
'stride=%d\n' % stride + \
|
||||
'size=%d\n' % size
|
||||
def convolutional(self, cv, detect=False):
|
||||
self.blocks[self.current] += 1
|
||||
|
||||
def upsample(self, stride=1):
|
||||
return '\n[upsample]\n' + \
|
||||
'stride=%d\n' % stride
|
||||
self.get_state_dict(cv.state_dict())
|
||||
|
||||
if cv._get_name() == 'Conv2d':
|
||||
filters = cv.out_channels
|
||||
size = cv.kernel_size
|
||||
stride = cv.stride
|
||||
pad = cv.padding
|
||||
groups = cv.groups
|
||||
bias = cv.bias
|
||||
bn = False
|
||||
act = 'linear' if not detect else 'logistic'
|
||||
else:
|
||||
filters = cv.conv.out_channels
|
||||
size = cv.conv.kernel_size
|
||||
stride = cv.conv.stride
|
||||
pad = cv.conv.padding
|
||||
groups = cv.conv.groups
|
||||
bias = cv.conv.bias
|
||||
bn = True if hasattr(cv, 'bn') else False
|
||||
act = self.get_activation(cv.act._get_name()) if hasattr(cv, 'act') else 'linear'
|
||||
|
||||
def convolutional(self, bn=False, size=1, stride=1, pad=1, filters=1, groups=1, activation='linear'):
|
||||
b = 'batch_normalize=1\n' if bn is True else ''
|
||||
g = 'groups=%d\n' % groups if groups > 1 else ''
|
||||
return '\n[convolutional]\n' + \
|
||||
b + \
|
||||
'filters=%d\n' % filters + \
|
||||
'size=%d\n' % size + \
|
||||
'stride=%d\n' % stride + \
|
||||
'pad=%d\n' % pad + \
|
||||
g + \
|
||||
'activation=%s\n' % activation
|
||||
w = 'bias=0\n' if bias is None and bn is False else ''
|
||||
|
||||
def yolo(self, mask='', anchors='', classes=80, num=3):
|
||||
return '\n[yolo]\n' + \
|
||||
'mask=%s\n' % mask + \
|
||||
'anchors=%s\n' % anchors + \
|
||||
'classes=%d\n' % classes + \
|
||||
'num=%d\n' % num + \
|
||||
'scale_x_y=2.0\n' + \
|
||||
'beta_nms=0.6\n' + \
|
||||
'new_coords=1\n'
|
||||
self.fc.write('\n[convolutional]\n' +
|
||||
b +
|
||||
'filters=%d\n' % filters +
|
||||
'size=%s\n' % (size[0] if len(size) == 2 and size[0] == size[1] else str(size)[1:-1]) +
|
||||
'stride=%s\n' % (stride[0] if len(stride) == 2 and stride[0] == stride[1] else str(stride)[1:-1]) +
|
||||
'pad=%s\n' % (pad[0] if len(pad) == 2 and pad[0] == pad[1] else str(pad)[1:-1]) +
|
||||
g +
|
||||
w +
|
||||
'activation=%s\n' % act)
|
||||
|
||||
def batchnorm(self, bn, act):
|
||||
self.blocks[self.current] += 1
|
||||
|
||||
self.get_state_dict(bn.state_dict())
|
||||
|
||||
filters = bn.num_features
|
||||
act = self.get_activation(act._get_name())
|
||||
|
||||
self.fc.write('\n[batchnorm]\n' +
|
||||
'filters=%d\n' % filters +
|
||||
'activation=%s\n' % act)
|
||||
|
||||
def route(self, layers):
|
||||
self.blocks[self.current] += 1
|
||||
|
||||
self.fc.write('\n[route]\n' +
|
||||
'layers=%s\n' % layers)
|
||||
|
||||
def shortcut(self, r, activation='linear'):
|
||||
self.blocks[self.current] += 1
|
||||
|
||||
self.fc.write('\n[shortcut]\n' +
|
||||
'from=%d\n' % r +
|
||||
'activation=%s\n' % activation)
|
||||
|
||||
def maxpool(self, m):
|
||||
self.blocks[self.current] += 1
|
||||
|
||||
stride = m.stride
|
||||
size = m.kernel_size
|
||||
mode = m.ceil_mode
|
||||
|
||||
m = 'maxpool_up' if mode else 'maxpool'
|
||||
|
||||
self.fc.write('\n[%s]\n' % m +
|
||||
'stride=%d\n' % stride +
|
||||
'size=%d\n' % size)
|
||||
|
||||
def upsample(self, child):
|
||||
self.blocks[self.current] += 1
|
||||
|
||||
stride = child.scale_factor
|
||||
|
||||
self.fc.write('\n[upsample]\n' +
|
||||
'stride=%d\n' % stride)
|
||||
|
||||
def yolo(self, i):
|
||||
self.blocks[self.current] += 1
|
||||
|
||||
self.fc.write('\n[yolo]\n' +
|
||||
'mask=%s\n' % self.masks[i] +
|
||||
'anchors=%s\n' % self.anchors +
|
||||
'classes=%d\n' % self.nc +
|
||||
'num=%d\n' % self.num +
|
||||
'scale_x_y=2.0\n' +
|
||||
'new_coords=1\n')
|
||||
|
||||
def get_state_dict(self, state_dict):
|
||||
for k, v in state_dict.items():
|
||||
if 'num_batches_tracked' not in k:
|
||||
vr = v.reshape(-1).numpy()
|
||||
self.fw.write('{} {} '.format(k, len(vr)))
|
||||
for vv in vr:
|
||||
self.fw.write(' ')
|
||||
self.fw.write(struct.pack('>f', float(vv)).hex())
|
||||
self.fw.write('\n')
|
||||
self.wc += 1
|
||||
|
||||
def get_anchors(self, state_dict, out_channels):
|
||||
anchor_grid = state_dict['anchor_grid']
|
||||
aa = anchor_grid.reshape(-1).tolist()
|
||||
am = anchor_grid.tolist()
|
||||
|
||||
self.num = (len(aa) / 2)
|
||||
self.nc = int((out_channels / (self.num / len(am))) - 5)
|
||||
self.anchors = str(aa)[1:-1]
|
||||
|
||||
n = 0
|
||||
for m in am:
|
||||
mask = []
|
||||
for _ in range(len(m)):
|
||||
mask.append(n)
|
||||
n += 1
|
||||
self.masks.append(str(mask)[1:-1])
|
||||
|
||||
def get_route(self, n):
|
||||
r = 0
|
||||
for i, b in enumerate(self.blocks):
|
||||
if i <= n:
|
||||
r += b
|
||||
else:
|
||||
break
|
||||
return r
|
||||
|
||||
def get_activation(self, act):
|
||||
if act == 'Hardswish':
|
||||
return 'hardswish'
|
||||
elif act == 'LeakyReLU':
|
||||
return 'leaky'
|
||||
elif act == 'SiLU':
|
||||
return 'silu'
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description='PyTorch YOLOv5 conversion')
|
||||
parser.add_argument('-w', '--weights', required=True, help='Input weights (.pt) file path (required)')
|
||||
parser.add_argument('-c', '--yaml', help='Input cfg (.yaml) file path')
|
||||
parser.add_argument(
|
||||
'-s', '--size', nargs='+', type=int, default=[640], help='Inference size [H,W] (default [640])')
|
||||
args = parser.parse_args()
|
||||
if not os.path.isfile(args.weights):
|
||||
raise SystemExit('Invalid weights file')
|
||||
if not args.yaml:
|
||||
args.yaml = ''
|
||||
return args.weights, args.yaml, args.size
|
||||
return args.weights, args.size
|
||||
|
||||
|
||||
def get_width(x, gw, divisor=8):
|
||||
return int(math.ceil((x * gw) / divisor)) * divisor
|
||||
|
||||
|
||||
def get_depth(x, gd):
|
||||
if x == 1:
|
||||
return 1
|
||||
r = int(round(x * gd))
|
||||
if x * gd - int(x * gd) == 0.5 and int(x * gd) % 2 == 0:
|
||||
r -= 1
|
||||
return max(r, 1)
|
||||
|
||||
|
||||
pt_file, yaml_file, inference_size = parse_args()
|
||||
pt_file, inference_size = parse_args()
|
||||
|
||||
model_name = os.path.basename(pt_file).split('.pt')[0]
|
||||
wts_file = model_name + '.wts' if 'yolov5' in model_name else 'yolov5_' + model_name + '.wts'
|
||||
cfg_file = model_name + '.cfg' if 'yolov5' in model_name else 'yolov5_' + model_name + '.cfg'
|
||||
|
||||
if yaml_file == '':
|
||||
yaml_file = 'models/' + model_name + '.yaml'
|
||||
if not os.path.isfile(yaml_file):
|
||||
yaml_file = 'models/hub/' + model_name + '.yaml'
|
||||
if not os.path.isfile(yaml_file):
|
||||
raise SystemExit('YAML file not found')
|
||||
elif not os.path.isfile(yaml_file):
|
||||
raise SystemExit('Invalid YAML file')
|
||||
|
||||
device = select_device('cpu')
|
||||
model = torch.load(pt_file, map_location=device)['model'].float()
|
||||
|
||||
@@ -112,217 +309,29 @@ model.model[-1].register_buffer('anchor_grid', anchor_grid)
|
||||
|
||||
model.to(device).eval()
|
||||
|
||||
nc = 0
|
||||
anchors = ''
|
||||
masks = []
|
||||
with open(wts_file, 'w') as fw, open(cfg_file, 'w') as fc:
|
||||
layers = Layers(len(model.model), inference_size, fw, fc)
|
||||
|
||||
yolo_idx = 0
|
||||
spp_idx = 0
|
||||
|
||||
for k, v in model.state_dict().items():
|
||||
if 'anchor_grid' in k:
|
||||
yolo_idx = int(k.split('.')[1])
|
||||
vr = v.cpu().numpy().tolist()
|
||||
a = v.reshape(-1).cpu().numpy().astype(float).tolist()
|
||||
anchors = str(a)[1:-1]
|
||||
num = 0
|
||||
for m in vr:
|
||||
mask = []
|
||||
for _ in range(len(m)):
|
||||
mask.append(num)
|
||||
num += 1
|
||||
masks.append(mask)
|
||||
elif '.%d.m.0.weight' % yolo_idx in k:
|
||||
vr = v.cpu().numpy().tolist()
|
||||
nc = int((len(vr) / len(masks[0])) - 5)
|
||||
|
||||
with open(cfg_file, 'w') as c:
|
||||
with open(yaml_file, 'r', encoding='utf-8') as f:
|
||||
c.write('[net]\n')
|
||||
c.write('width=%d\n' % (inference_size[0] if len(inference_size) == 1 else inference_size[1]))
|
||||
c.write('height=%d\n' % inference_size[0])
|
||||
c.write('channels=3\n')
|
||||
c.write('letter_box=1\n')
|
||||
depth_multiple = 0
|
||||
width_multiple = 0
|
||||
layers = []
|
||||
yoloLayers = YoloLayers()
|
||||
f = yaml.load(f, Loader=yaml.FullLoader)
|
||||
for topic in f:
|
||||
if topic == 'depth_multiple':
|
||||
depth_multiple = f[topic]
|
||||
elif topic == 'width_multiple':
|
||||
width_multiple = f[topic]
|
||||
elif topic == 'backbone' or topic == 'head':
|
||||
for v in f[topic]:
|
||||
if v[2] == 'Focus':
|
||||
layer = '\n# Focus\n'
|
||||
blocks = 0
|
||||
layer += yoloLayers.reorg()
|
||||
blocks += 1
|
||||
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), size=v[3][1],
|
||||
activation='silu')
|
||||
blocks += 1
|
||||
layers.append([layer, blocks])
|
||||
if v[2] == 'Conv':
|
||||
layer = '\n# Conv\n'
|
||||
blocks = 0
|
||||
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple), size=v[3][1],
|
||||
stride=v[3][2], activation='silu')
|
||||
blocks += 1
|
||||
layers.append([layer, blocks])
|
||||
elif v[2] == 'C3':
|
||||
layer = '\n# C3\n'
|
||||
blocks = 0
|
||||
# SPLIT
|
||||
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
|
||||
activation='silu')
|
||||
blocks += 1
|
||||
layer += yoloLayers.route(layers='-2')
|
||||
blocks += 1
|
||||
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
|
||||
activation='silu')
|
||||
blocks += 1
|
||||
# Residual Block
|
||||
if len(v[3]) == 1 or v[3][1] is True:
|
||||
for _ in range(get_depth(v[1], depth_multiple)):
|
||||
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
|
||||
activation='silu')
|
||||
blocks += 1
|
||||
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
|
||||
size=3, activation='silu')
|
||||
blocks += 1
|
||||
layer += yoloLayers.shortcut(route=-3)
|
||||
blocks += 1
|
||||
# Merge
|
||||
layer += yoloLayers.route(layers='-1, -%d' % (3 * get_depth(v[1], depth_multiple) + 3))
|
||||
blocks += 1
|
||||
for child in model.model.children():
|
||||
if child._get_name() == 'Focus':
|
||||
layers.Focus(child)
|
||||
elif child._get_name() == 'Conv':
|
||||
layers.Conv(child)
|
||||
elif child._get_name() == 'BottleneckCSP':
|
||||
layers.BottleneckCSP(child)
|
||||
elif child._get_name() == 'C3':
|
||||
layers.C3(child)
|
||||
elif child._get_name() == 'SPP':
|
||||
layers.SPP(child)
|
||||
elif child._get_name() == 'SPPF':
|
||||
layers.SPPF(child)
|
||||
elif child._get_name() == 'Upsample':
|
||||
layers.Upsample(child)
|
||||
elif child._get_name() == 'Concat':
|
||||
layers.Concat(child)
|
||||
elif child._get_name() == 'Detect':
|
||||
layers.Detect(child)
|
||||
else:
|
||||
for _ in range(get_depth(v[1], depth_multiple)):
|
||||
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
|
||||
activation='silu')
|
||||
blocks += 1
|
||||
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
|
||||
size=3, activation='silu')
|
||||
blocks += 1
|
||||
# Merge
|
||||
layer += yoloLayers.route(layers='-1, -%d' % (2 * get_depth(v[1], depth_multiple) + 3))
|
||||
blocks += 1
|
||||
# Transition
|
||||
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple),
|
||||
activation='silu')
|
||||
blocks += 1
|
||||
layers.append([layer, blocks])
|
||||
elif v[2] == 'SPP':
|
||||
spp_idx = len(layers)
|
||||
layer = '\n# SPP\n'
|
||||
blocks = 0
|
||||
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
|
||||
activation='silu')
|
||||
blocks += 1
|
||||
layer += yoloLayers.maxpool(size=v[3][1][0])
|
||||
blocks += 1
|
||||
layer += yoloLayers.route(layers='-2')
|
||||
blocks += 1
|
||||
layer += yoloLayers.maxpool(size=v[3][1][1])
|
||||
blocks += 1
|
||||
layer += yoloLayers.route(layers='-4')
|
||||
blocks += 1
|
||||
layer += yoloLayers.maxpool(size=v[3][1][2])
|
||||
blocks += 1
|
||||
layer += yoloLayers.route(layers='-6, -5, -3, -1')
|
||||
blocks += 1
|
||||
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple),
|
||||
activation='silu')
|
||||
blocks += 1
|
||||
layers.append([layer, blocks])
|
||||
elif v[2] == 'SPPF':
|
||||
spp_idx = len(layers)
|
||||
layer = '\n# SPPF\n'
|
||||
blocks = 0
|
||||
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple) / 2,
|
||||
activation='silu')
|
||||
blocks += 1
|
||||
layer += yoloLayers.maxpool(size=v[3][1])
|
||||
blocks += 1
|
||||
layer += yoloLayers.maxpool(size=v[3][1])
|
||||
blocks += 1
|
||||
layer += yoloLayers.maxpool(size=v[3][1])
|
||||
blocks += 1
|
||||
layer += yoloLayers.route(layers='-4, -3, -2, -1')
|
||||
blocks += 1
|
||||
layer += yoloLayers.convolutional(bn=True, filters=get_width(v[3][0], width_multiple),
|
||||
activation='silu')
|
||||
blocks += 1
|
||||
layers.append([layer, blocks])
|
||||
elif v[2] == 'nn.Upsample':
|
||||
layer = '\n# nn.Upsample\n'
|
||||
blocks = 0
|
||||
layer += yoloLayers.upsample(stride=v[3][1])
|
||||
blocks += 1
|
||||
layers.append([layer, blocks])
|
||||
elif v[2] == 'Concat':
|
||||
route = v[0][1]
|
||||
route = yoloLayers.get_route(route, layers) if route > 0 else \
|
||||
yoloLayers.get_route(len(layers) + route, layers)
|
||||
layer = '\n# Concat\n'
|
||||
blocks = 0
|
||||
layer += yoloLayers.route(layers='-1, %d' % (route - 1))
|
||||
blocks += 1
|
||||
layers.append([layer, blocks])
|
||||
elif v[2] == 'Detect':
|
||||
for i, n in enumerate(v[0]):
|
||||
route = yoloLayers.get_route(n, layers)
|
||||
layer = '\n# Detect\n'
|
||||
blocks = 0
|
||||
layer += yoloLayers.route(layers='%d' % (route - 1))
|
||||
blocks += 1
|
||||
layer += yoloLayers.convolutional(filters=((nc + 5) * len(masks[i])), activation='logistic')
|
||||
blocks += 1
|
||||
layer += yoloLayers.yolo(mask=str(masks[i])[1:-1], anchors=anchors, classes=nc, num=num)
|
||||
blocks += 1
|
||||
layers.append([layer, blocks])
|
||||
for layer in layers:
|
||||
c.write(layer[0])
|
||||
raise SystemExit('Model not supported')
|
||||
|
||||
with open(wts_file, 'w') as f:
|
||||
wts_write = ''
|
||||
conv_count = 0
|
||||
cv1 = ''
|
||||
cv3 = ''
|
||||
cv3_idx = 0
|
||||
for k, v in model.state_dict().items():
|
||||
if 'num_batches_tracked' not in k and 'anchors' not in k and 'anchor_grid' not in k:
|
||||
vr = v.reshape(-1).cpu().numpy()
|
||||
idx = int(k.split('.')[1])
|
||||
if '.cv1.' in k and '.m.' not in k and idx != spp_idx:
|
||||
cv1 += '{} {} '.format(k, len(vr))
|
||||
for vv in vr:
|
||||
cv1 += ' '
|
||||
cv1 += struct.pack('>f', float(vv)).hex()
|
||||
cv1 += '\n'
|
||||
conv_count += 1
|
||||
elif cv1 != '' and '.m.' in k:
|
||||
wts_write += cv1
|
||||
cv1 = ''
|
||||
if '.cv3.' in k:
|
||||
cv3 += '{} {} '.format(k, len(vr))
|
||||
for vv in vr:
|
||||
cv3 += ' '
|
||||
cv3 += struct.pack('>f', float(vv)).hex()
|
||||
cv3 += '\n'
|
||||
cv3_idx = idx
|
||||
conv_count += 1
|
||||
elif cv3 != '' and cv3_idx != idx:
|
||||
wts_write += cv3
|
||||
cv3 = ''
|
||||
cv3_idx = 0
|
||||
if '.cv3.' not in k and not ('.cv1.' in k and '.m.' not in k and idx != spp_idx):
|
||||
wts_write += '{} {} '.format(k, len(vr))
|
||||
for vv in vr:
|
||||
wts_write += ' '
|
||||
wts_write += struct.pack('>f', float(vv)).hex()
|
||||
wts_write += '\n'
|
||||
conv_count += 1
|
||||
f.write('{}\n'.format(conv_count))
|
||||
f.write(wts_write)
|
||||
os.system('echo "%d" | cat - %s > temp && mv temp %s' % (layers.wc, wts_file, wts_file))
|
||||
|
||||
Reference in New Issue
Block a user