tensorRT RefineDet (wang-xinyu#443)

Author: wuzuowuyou

refinedet/CMakeLists.txt

@@ -0,0 +1,48 @@
+cmake_minimum_required(VERSION 2.6)
+
+project(refinedet)
+
+add_definitions(-std=c++11)
+
+option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
+set(CMAKE_CXX_STANDARD 11)
+set(CMAKE_BUILD_TYPE Debug)
+
+# tensorrt
+include_directories(/data_2/tensorrt/TensorRT-7.0.0.11/include/) #include_directories(/usr/include/x86_64-linux-gnu/)
+link_directories(/data_2/tensorrt/TensorRT-7.0.0.11/lib/) #link_directories(/usr/lib/x86_64-linux-gnu/)
+
+
+find_package(CUDA REQUIRED)
+
+include_directories(${PROJECT_SOURCE_DIR}/include)
+# include and link dirs of cuda and tensorrt, you need adapt them if yours are different
+# cuda
+include_directories(/usr/local/cuda/include)
+link_directories(/usr/local/cuda/lib64)
+
+#find_package(OpenCV)
+#include_directories(OpenCV_INCLUDE_DIRS)
+
+include_directories(/home/software_install/opencv3.4.6/include)
+link_directories(/home/software_install/opencv3.4.6/lib)
+
+
+set(CMAKE_PREFIX_PATH "/data_1/torch1.1.0") ###torch1.1.0
+find_package(Torch REQUIRED)
+
+include_directories(/data_1/torch1.1.0/include)
+link_directories(/data_1/torch1.1.0/lib)
+
+
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")
+
+
+add_executable(refinedet ${PROJECT_SOURCE_DIR}/calibrator.cpp ${PROJECT_SOURCE_DIR}/refinedet.cpp)
+target_link_libraries(refinedet nvinfer)
+target_link_libraries(refinedet cudart)
+target_link_libraries(refinedet "${TORCH_LIBRARIES}")
+target_link_libraries(refinedet opencv_calib3d opencv_core opencv_dnn opencv_imgproc opencv_highgui opencv_imgcodecs caffe2)
+
+add_definitions(-O2 -pthread)
+

refinedet/README.md

@@ -0,0 +1,39 @@
+# RefineDet
+
+For the Pytorch implementation, you can refer to [luuuyi/RefineDet.PyTorch](https://github.com/luuuyi/RefineDet.PyTorch)
+
+## How to run
+```
+1. generate wts file. from pytorch
+python gen_wts_refinedet.py
+// a file 'refinedet.wts' will be generated.
+
+2. build tensorrtx/RefineDet and run or Using clion to open a project(recommend)
+Configuration file in configure.h
+You need configure your own paths and modes(SERIALIZE or INFER)
+Detailed information reference configure.h
+mkdir build
+cd build
+cmake ..
+make
+```
+
+## dependence
+```
+TensorRT7.0.0.11 
+OpenCV >= 3.4
+libtorch >=1.1.0
+```
+
+
+## feature
+1.tensorrt Multi output  
+2.L2norm  
+3.Postprocessing with libtorch
+
+
+## More Information
+
+See the readme in [home page.](https://github.com/wang-xinyu/tensorrtx)  
+[tensorrt tutorials](https://github.com/wang-xinyu/tensorrtx/tree/master/tutorials)  
+For more detailed guidance, see [yhl blog](https://www.cnblogs.com/yanghailin/p/14525128.html)

refinedet/calibrator.cpp

@@ -0,0 +1,81 @@
+#include <iostream>
+#include <iterator>
+#include <fstream>
+#include <opencv2/dnn/dnn.hpp>
+#include "calibrator.h"
+#include "cuda_runtime_api.h"
+#include "utils.h"
+
+Int8EntropyCalibrator2::Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache)
+    : batchsize_(batchsize)
+    , input_w_(input_w)
+    , input_h_(input_h)
+    , img_idx_(0)
+    , img_dir_(img_dir)
+    , calib_table_name_(calib_table_name)
+    , input_blob_name_(input_blob_name)
+    , read_cache_(read_cache)
+{
+    input_count_ = 3 * input_w * input_h * batchsize;
+    CUDA_CHECK(cudaMalloc(&device_input_, input_count_ * sizeof(float)));
+    read_files_in_dir(img_dir, img_files_);
+}
+
+Int8EntropyCalibrator2::~Int8EntropyCalibrator2()
+{
+    CUDA_CHECK(cudaFree(device_input_));
+}
+
+int Int8EntropyCalibrator2::getBatchSize() const
+{
+    return batchsize_;
+}
+
+bool Int8EntropyCalibrator2::getBatch(void* bindings[], const char* names[], int nbBindings)
+{
+    if (img_idx_ + batchsize_ > (int)img_files_.size()) {
+        return false;
+    }
+
+    std::vector<cv::Mat> input_imgs_;
+    for (int i = img_idx_; i < img_idx_ + batchsize_; i++) {
+        std::cout << img_files_[i] << "  " << i << std::endl;
+        cv::Mat temp = cv::imread(img_dir_ + img_files_[i]);
+        if (temp.empty()){
+            std::cerr << "Fatal error: image cannot open!" << std::endl;
+            return false;
+        }
+//        cv::Mat pr_img = preprocess_img(temp, input_w_, input_h_);
+        input_imgs_.push_back(temp);
+    }
+    img_idx_ += batchsize_;
+    cv::Mat blob = cv::dnn::blobFromImages(input_imgs_, 1.0, cv::Size(input_w_, input_h_), cv::Scalar(123.0, 117.0, 104.0), true, false);
+//    cv::Mat blob = cv::dnn::blobFromImages(input_imgs_, 1.0 / 255.0, cv::Size(input_w_, input_h_), cv::Scalar(0, 0, 0), true, false);
+
+    CUDA_CHECK(cudaMemcpy(device_input_, blob.ptr<float>(0), input_count_ * sizeof(float), cudaMemcpyHostToDevice));
+    assert(!strcmp(names[0], input_blob_name_));
+    bindings[0] = device_input_;
+    return true;
+}
+
+const void* Int8EntropyCalibrator2::readCalibrationCache(size_t& length)
+{
+    std::cout << "reading calib cache: " << calib_table_name_ << std::endl;
+    calib_cache_.clear();
+    std::ifstream input(calib_table_name_, std::ios::binary);
+    input >> std::noskipws;
+    if (read_cache_ && input.good())
+    {
+        std::copy(std::istream_iterator<char>(input), std::istream_iterator<char>(), std::back_inserter(calib_cache_));
+    }
+    length = calib_cache_.size();
+    return length ? calib_cache_.data() : nullptr;
+}
+
+void Int8EntropyCalibrator2::writeCalibrationCache(const void* cache, size_t length)
+{
+    std::cout << "writing calib cache: " << calib_table_name_ << " size: " << length << std::endl;
+    std::ofstream output(calib_table_name_, std::ios::binary);
+    output.write(reinterpret_cast<const char*>(cache), length);
+}
+

refinedet/calibrator.h

@@ -0,0 +1,39 @@
+#ifndef ENTROPY_CALIBRATOR_H
+#define ENTROPY_CALIBRATOR_H
+
+#include "NvInfer.h"
+#include <string>
+#include <vector>
+
+//! \class Int8EntropyCalibrator2
+//!
+//! \brief Implements Entropy calibrator 2.
+//!  CalibrationAlgoType is kENTROPY_CALIBRATION_2.
+//!
+class Int8EntropyCalibrator2 : public nvinfer1::IInt8EntropyCalibrator2
+{
+public:
+    Int8EntropyCalibrator2(int batchsize, int input_w, int input_h, const char* img_dir, const char* calib_table_name, const char* input_blob_name, bool read_cache = true);
+
+    virtual ~Int8EntropyCalibrator2();
+    int getBatchSize() const override;
+    bool getBatch(void* bindings[], const char* names[], int nbBindings) override;
+    const void* readCalibrationCache(size_t& length) override;
+    void writeCalibrationCache(const void* cache, size_t length) override;
+
+private:
+    int batchsize_;
+    int input_w_;
+    int input_h_;
+    int img_idx_;
+    std::string img_dir_;
+    std::vector<std::string> img_files_;
+    size_t input_count_;
+    std::string calib_table_name_;
+    const char* input_blob_name_;
+    bool read_cache_;
+    void* device_input_;
+    std::vector<char> calib_cache_;
+};
+
+#endif // ENTROPY_CALIBRATOR_H

refinedet/configure.h

@@ -0,0 +1,60 @@
+
+#define USE_FP32  // set USE_INT8 or USE_FP16 or USE_FP32
+
+const int num_class = 25; //num_class + 1     //Including background class
+
+//SERIALIZE: It indicates that to generate engin by serialization, the following path needs to be set,path_wts_ and path_save_engine
+//INFER: It shows that it is a reasoning mode,the following path needs to be set,path_engine
+#define INFER    //SERIALIZE   INFER
+
+const std::string path_engine = "/data_2//cmake-build-debug/refinedet_0312-now.engine";
+const std::string path_wts = "/data_1/refinedet/pytorch_refinedet-master/refinedet0312.wts";
+const std::string path_save_engine = "./refinedet_0312-now.engine";
+
+//Picture folder to be detected
+const char *p_dir_name = "/data_1/img/";
+
+const float TH = 0.2;  //Confidence threshold
+const int T_show = 1; //1:Show the effect      0:Test map to generate TXT
+//The path to save the generated TXT when testing the map
+std::string save_path_txt = "/data_1/txt/";
+
+#define DEVICE 0  // GPU id
+
+// stuff we know about the network and the input/output blobs
+static const int INPUT_H = 320;
+static const int INPUT_W = 320;
+const char* INPUT_BLOB_NAME = "data";
+const char* OUTPUT_BLOB_NAME_arm_loc = "arm_loc";
+const char* OUTPUT_BLOB_NAME_arm_conf = "arm_conf";
+const char* OUTPUT_BLOB_NAME_odm_loc = "odm_loc";
+const char* OUTPUT_BLOB_NAME_odm_conf = "odm_conf";
+
+std::string label_map[] =
+        {
+                "background",
+                "aa",
+                "bb",
+                "cc",
+                "dd",
+                "ee",
+                "ff",
+                "gg",
+                "hh",
+                "ii",
+                "jj",
+                "kk",
+                "ll",
+                "mm",
+                "nn",
+                "oo",
+                "pp",
+                "qq",
+                "rr",
+                "ss",
+                "tt",
+                "uu",
+                "vv",
+                "ww",
+                "xx"
+        };

refinedet/gen_wts_refinedet.py

@@ -0,0 +1,28 @@
+import torch
+import torch.nn as nn
+import struct
+from models.refinedet import build_refinedet
+
+
+
+num_classes = 25
+path_model = "/data_2/project_2021/pytorch_refinedet/2021/20210308.pth"
+path_save_wts = "./refinedet0312.wts"
+input_size = 320
+
+net = build_refinedet('test', input_size, num_classes)  # initialize net
+net.load_state_dict(torch.load(path_model))
+net.eval()
+
+
+f = open(path_save_wts, 'w')
+f.write('{}\n'.format(len(net.state_dict().keys())))
+for k, v in net.state_dict().items():
+    vr = v.reshape(-1).cpu().numpy()
+    f.write('{} {} '.format(k, len(vr)))
+    for vv in vr:
+        f.write(' ')
+        f.write(struct.pack('>f',float(vv)).hex())
+    f.write('\n')
+
+print("success generate wts!")