diff --git a/config_infer_primary_yoloV5.txt b/config_infer_primary_yoloV5.txt index 72bb8ed..1b0c77e 100644 --- a/config_infer_primary_yoloV5.txt +++ b/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 -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 diff --git a/docs/YOLOR.md b/docs/YOLOR.md new file mode 100644 index 0000000..e111a80 --- /dev/null +++ b/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 +``` diff --git a/docs/YOLOv5.md b/docs/YOLOv5.md new file mode 100644 index 0000000..7a5da06 --- /dev/null +++ b/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 +``` diff --git a/docs/customModels.md b/docs/customModels.md index 25049c1..bc8764d 100644 --- a/docs/customModels.md +++ b/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. ## diff --git a/docs/multipleGIEs.md b/docs/multipleGIEs.md index bee0cd9..2aab074 100644 --- a/docs/multipleGIEs.md +++ b/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. ## diff --git a/nvdsinfer_custom_impl_Yolo/Makefile b/nvdsinfer_custom_impl_Yolo/Makefile index a2cd96c..c5bedbc 100644 --- a/nvdsinfer_custom_impl_Yolo/Makefile +++ b/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 \ diff --git a/nvdsinfer_custom_impl_Yolo/layers/activation_layer.cpp b/nvdsinfer_custom_impl_Yolo/layers/activation_layer.cpp index b42b4de..500789c 100644 --- a/nvdsinfer_custom_impl_Yolo/layers/activation_layer.cpp +++ b/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(); } diff --git a/nvdsinfer_custom_impl_Yolo/layers/batchnorm_layer.cpp b/nvdsinfer_custom_impl_Yolo/layers/batchnorm_layer.cpp new file mode 100644 index 0000000..063f3d1 --- /dev/null +++ b/nvdsinfer_custom_impl_Yolo/layers/batchnorm_layer.cpp @@ -0,0 +1,114 @@ +/* + * Created by Marcos Luciano + * https://www.github.com/marcoslucianops + */ + +#include +#include "batchnorm_layer.h" + +nvinfer1::ILayer* batchnormLayer( + int layerIdx, + std::map& block, + std::vector& weights, + std::vector& 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 bnBiases; + std::vector bnWeights; + std::vector bnRunningMean; + std::vector 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; +} diff --git a/nvdsinfer_custom_impl_Yolo/layers/batchnorm_layer.h b/nvdsinfer_custom_impl_Yolo/layers/batchnorm_layer.h new file mode 100644 index 0000000..514b456 --- /dev/null +++ b/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 +#include + +#include "NvInfer.h" + +#include "activation_layer.h" + +nvinfer1::ILayer* batchnormLayer( + int layerIdx, + std::map& block, + std::vector& weights, + std::vector& trtWeights, + int& weightPtr, + std::string weightsType, + float eps, + nvinfer1::ITensor* input, + nvinfer1::INetworkDefinition* network); + +#endif diff --git a/nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.cpp b/nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.cpp index 2bcb00a..6b502ad 100644 --- a/nvdsinfer_custom_impl_Yolo/layers/convolutional_layer.cpp +++ b/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]; - for (int i = 0; i < filters; ++i) - { - val[i] = weights[weightPtr]; - weightPtr++; + float* val; + if (bias != 0) { + val = new float[filters]; + for (int i = 0; i < filters; ++i) + { + 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]; - for (int i = 0; i < filters; ++i) - { - val[i] = weights[weightPtr]; - weightPtr++; + if (bias != 0) { + val = new float[filters]; + for (int i = 0; i < filters; ++i) + { + 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); } } diff --git a/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp b/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp index 38efa06..8a9a31b 100644 --- a/nvdsinfer_custom_impl_Yolo/layers/maxpool_layer.cpp +++ b/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->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; diff --git a/nvdsinfer_custom_impl_Yolo/yolo.cpp b/nvdsinfer_custom_impl_Yolo/yolo.cpp index 3099f65..2e7fb9f 100644 --- a/nvdsinfer_custom_impl_Yolo/yolo.cpp +++ b/nvdsinfer_custom_impl_Yolo/yolo.cpp @@ -207,6 +207,20 @@ NvDsInferStatus Yolo::buildYoloNetwork(std::vector& 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; diff --git a/nvdsinfer_custom_impl_Yolo/yolo.h b/nvdsinfer_custom_impl_Yolo/yolo.h index 219ccac..038a2ec 100644 --- a/nvdsinfer_custom_impl_Yolo/yolo.h +++ b/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" diff --git a/readme.md b/readme.md index 7ac5f42..b21dec0 100644 --- a/readme.md +++ b/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. diff --git a/utils/gen_wts_yoloV5.py b/utils/gen_wts_yoloV5.py index f979a40..9d20859 100644 --- a/utils/gen_wts_yoloV5.py +++ b/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(): - 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 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: + raise SystemExit('Model not supported') -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 - 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) +os.system('echo "%d" | cat - %s > temp && mv temp %s' % (layers.wc, wts_file, wts_file))