New YOLOv5 conversion and support (>= v2.0)

config_infer_primary_yoloV5.txt

@@ -2,8 +2,8 @@
 gpu-id=0
 net-scale-factor=0.0039215697906911373
 model-color-format=0
-custom-network-config=yolov5n.cfg
+custom-network-config=yolov5s.cfg
-model-file=yolov5n.wts
+model-file=yolov5s.wts
 model-engine-file=model_b1_gpu0_fp32.engine
 #int8-calib-file=calib.table
 labelfile-path=labels.txt

docs/YOLOR.md

@@ -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
+```

docs/YOLOv5.md

@@ -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
+```

docs/customModels.md

@@ -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.
 
 ##
 

docs/multipleGIEs.md

@@ -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.
 
 ##
 

nvdsinfer_custom_impl_Yolo/Makefile

@@ -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 \

nvdsinfer_custom_impl_Yolo/layers/activation_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();
     }

nvdsinfer_custom_impl_Yolo/layers/batchnorm_layer.cpp

@@ -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;
+}

nvdsinfer_custom_impl_Yolo/layers/batchnorm_layer.h

@@ -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

nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.cpp

@@ -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,14 +66,17 @@ nvinfer1::ILayer* convolutionalLayer(
     if (weightsType == "weights") {
         if (batchNormalize == false)
         {
-            float* val = new float[filters];
+            float* val;
-            for (int i = 0; i < filters; ++i)
+            if (bias != 0) {
-            {
+                val = new float[filters];
-                val[i] = weights[weightPtr];
+                for (int i = 0; i < filters; ++i)
-                weightPtr++;
+                {
+                    val[i] = weights[weightPtr];
+                    weightPtr++;
+                }
+                convBias.values = val;
+                trtWeights.push_back(convBias);
             }
-            convBias.values = val;
-            trtWeights.push_back(convBias);
             val = new float[size];
             for (int i = 0; i < size; ++i)
             {
@@ -108,7 +116,8 @@ nvinfer1::ILayer* convolutionalLayer(
             }
             convWt.values = val;
             trtWeights.push_back(convWt);
-            trtWeights.push_back(convBias);
+            if (bias != 0)
+                trtWeights.push_back(convBias);
         }
     }
     else {
@@ -122,14 +131,16 @@ nvinfer1::ILayer* convolutionalLayer(
             }
             convWt.values = val;
             trtWeights.push_back(convWt);
-            val = new float[filters];
+            if (bias != 0) {
-            for (int i = 0; i < filters; ++i)
+                val = new float[filters];
-            {
+                for (int i = 0; i < filters; ++i)
-                val[i] = weights[weightPtr];
+                {
-                weightPtr++;
+                    val[i] = weights[weightPtr];
+                    weightPtr++;
+                }
+                convBias.values = val;
+                trtWeights.push_back(convBias);
             }
-            convBias.values = val;
-            trtWeights.push_back(convBias);
         }
         else
         {
@@ -161,7 +172,8 @@ nvinfer1::ILayer* convolutionalLayer(
                 weightPtr++;
             }
             trtWeights.push_back(convWt);
-            trtWeights.push_back(convBias);
+            if (bias != 0)
+                trtWeights.push_back(convBias);
         }
     }
 

nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp

@@ -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->setStrideNd(nvinfer1::Dims{2, {stride, stride}});
-    pool->setPaddingMode(nvinfer1::PaddingMode::kSAME_UPPER);
+    pool->setPaddingNd(nvinfer1::Dims{2, {size / 2, size / 2}});
     pool->setName(maxpoolLayerName.c_str());
 
     return pool;

nvdsinfer_custom_impl_Yolo/yolo.cpp

@@ -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;

nvdsinfer_custom_impl_Yolo/yolo.h

@@ -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"

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)
+* New documentation for multiple models
-* YOLOv5 4.0, 5.0, 6.0 and 6.1 support
+* **YOLOv5 >= 2.0 support**
-* YOLOR support
+* **YOLOR support**
-* **GPU YOLO Decoder (moved from CPU to GPU to get better performance)** [#138](https://github.com/marcoslucianops/DeepStream-Yolo/issues/138)
+* **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.

utils/gen_wts_yoloV5.py

@@ -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():
+class Layers(object):
-    def get_route(self, n, layers):
+    def __init__(self, n, size, fw, fc):
-        route = 0
+        self.blocks = [0 for _ in range(n)]
-        for i, layer in enumerate(layers):
+        self.current = 0
-            if i <= n:
-                route += layer[1]
-            else:
-                break
-        return route
 
-    def route(self, layers=''):
+        self.width = size[0] if len(size) == 1 else size[1]
-        return '\n[route]\n' + \
+        self.height = size[0]
-               'layers=%s\n' % layers
+
+        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'):
+        self.fc.write('\n[reorg]\n')
-        return '\n[shortcut]\n' + \
-               'from=%d\n' % route + \
-               'activation=%s\n' % activation
 
-    def maxpool(self, stride=1, size=1):
+    def convolutional(self, cv, detect=False):
-        return '\n[maxpool]\n' + \
+        self.blocks[self.current] += 1
-               'stride=%d\n' % stride + \
-               'size=%d\n' % size
 
-    def upsample(self, stride=1):
+        self.get_state_dict(cv.state_dict())
-        return '\n[upsample]\n' + \
+
-               'stride=%d\n' % stride
+        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' + \
+        w = 'bias=0\n' if bias is None and bn is False else ''
-               b + \
-               'filters=%d\n' % filters + \
-               'size=%d\n' % size + \
-               'stride=%d\n' % stride + \
-               'pad=%d\n' % pad + \
-               g + \
-               'activation=%s\n' % activation
 
-    def yolo(self, mask='', anchors='', classes=80, num=3):
+        self.fc.write('\n[convolutional]\n' +
-        return '\n[yolo]\n' + \
+                      b +
-               'mask=%s\n' % mask + \
+                      'filters=%d\n' % filters +
-               'anchors=%s\n' % anchors + \
+                      'size=%s\n' % (size[0] if len(size) == 2 and size[0] == size[1] else str(size)[1:-1]) +
-               'classes=%d\n' % classes + \
+                      'stride=%s\n' % (stride[0] if len(stride) == 2 and stride[0] == stride[1] else str(stride)[1:-1]) +
-               'num=%d\n' % num + \
+                      'pad=%s\n' % (pad[0] if len(pad) == 2 and pad[0] == pad[1] else str(pad)[1:-1]) +
-               'scale_x_y=2.0\n' + \
+                      g +
-               'beta_nms=0.6\n' + \
+                      w +
-               'new_coords=1\n'
+                      '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:
+    return args.weights, args.size
-        args.yaml = ''
-    return args.weights, args.yaml, args.size
 
 
-def get_width(x, gw, divisor=8):
+pt_file, inference_size = parse_args()
-    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()
 
 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
+with open(wts_file, 'w') as fw, open(cfg_file, 'w') as fc:
-anchors = ''
+    layers = Layers(len(model.model), inference_size, fw, fc)
-masks = []
 
-yolo_idx = 0
+    for child in model.model.children():
-spp_idx = 0
+        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:
+            raise SystemExit('Model not supported')
 
-for k, v in model.state_dict().items():
+os.system('echo "%d" | cat - %s > temp && mv temp %s' % (layers.wc, wts_file, wts_file))
-    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
-                        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])
-
-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)