diff --git a/refinedet/CMakeLists.txt b/refinedet/CMakeLists.txt new file mode 100755 index 0000000..43e36e9 --- /dev/null +++ b/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) + diff --git a/refinedet/README.md b/refinedet/README.md new file mode 100755 index 0000000..aaafdb3 --- /dev/null +++ b/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) diff --git a/refinedet/calibrator.cpp b/refinedet/calibrator.cpp new file mode 100755 index 0000000..6e3ed4d --- /dev/null +++ b/refinedet/calibrator.cpp @@ -0,0 +1,81 @@ +#include +#include +#include +#include +#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 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(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(input), std::istream_iterator(), 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(cache), length); +} + diff --git a/refinedet/calibrator.h b/refinedet/calibrator.h new file mode 100755 index 0000000..1cc9dbb --- /dev/null +++ b/refinedet/calibrator.h @@ -0,0 +1,39 @@ +#ifndef ENTROPY_CALIBRATOR_H +#define ENTROPY_CALIBRATOR_H + +#include "NvInfer.h" +#include +#include + +//! \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 img_files_; + size_t input_count_; + std::string calib_table_name_; + const char* input_blob_name_; + bool read_cache_; + void* device_input_; + std::vector calib_cache_; +}; + +#endif // ENTROPY_CALIBRATOR_H diff --git a/refinedet/configure.h b/refinedet/configure.h new file mode 100644 index 0000000..87f3f60 --- /dev/null +++ b/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" + }; \ No newline at end of file diff --git a/refinedet/gen_wts_refinedet.py b/refinedet/gen_wts_refinedet.py new file mode 100644 index 0000000..500203f --- /dev/null +++ b/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!") \ No newline at end of file diff --git a/refinedet/logging.h b/refinedet/logging.h new file mode 100755 index 0000000..2dd5847 --- /dev/null +++ b/refinedet/logging.h @@ -0,0 +1,504 @@ +/* + * Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#ifndef TENSORRT_LOGGING_H +#define TENSORRT_LOGGING_H + +#include "NvInferRuntimeCommon.h" +#include +#include +#include +#include +#include +#include +#include + +using Severity = nvinfer1::ILogger::Severity; + +class LogStreamConsumerBuffer : public std::stringbuf +{ +public: + LogStreamConsumerBuffer(std::ostream& stream, const std::string& prefix, bool shouldLog) + : mOutput(stream) + , mPrefix(prefix) + , mShouldLog(shouldLog) + { + } + + LogStreamConsumerBuffer(LogStreamConsumerBuffer&& other) + : mOutput(other.mOutput) + { + } + + ~LogStreamConsumerBuffer() + { + // std::streambuf::pbase() gives a pointer to the beginning of the buffered part of the output sequence + // std::streambuf::pptr() gives a pointer to the current position of the output sequence + // if the pointer to the beginning is not equal to the pointer to the current position, + // call putOutput() to log the output to the stream + if (pbase() != pptr()) + { + putOutput(); + } + } + + // synchronizes the stream buffer and returns 0 on success + // synchronizing the stream buffer consists of inserting the buffer contents into the stream, + // resetting the buffer and flushing the stream + virtual int sync() + { + putOutput(); + return 0; + } + + void putOutput() + { + if (mShouldLog) + { + // prepend timestamp + std::time_t timestamp = std::time(nullptr); + tm* tm_local = std::localtime(×tamp); + std::cout << "["; + std::cout << std::setw(2) << std::setfill('0') << 1 + tm_local->tm_mon << "/"; + std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_mday << "/"; + std::cout << std::setw(4) << std::setfill('0') << 1900 + tm_local->tm_year << "-"; + std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_hour << ":"; + std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_min << ":"; + std::cout << std::setw(2) << std::setfill('0') << tm_local->tm_sec << "] "; + // std::stringbuf::str() gets the string contents of the buffer + // insert the buffer contents pre-appended by the appropriate prefix into the stream + mOutput << mPrefix << str(); + // set the buffer to empty + str(""); + // flush the stream + mOutput.flush(); + } + } + + void setShouldLog(bool shouldLog) + { + mShouldLog = shouldLog; + } + +private: + std::ostream& mOutput; + std::string mPrefix; + bool mShouldLog; +}; + +//! +//! \class LogStreamConsumerBase +//! \brief Convenience object used to initialize LogStreamConsumerBuffer before std::ostream in LogStreamConsumer +//! +class LogStreamConsumerBase +{ +public: + LogStreamConsumerBase(std::ostream& stream, const std::string& prefix, bool shouldLog) + : mBuffer(stream, prefix, shouldLog) + { + } + +protected: + LogStreamConsumerBuffer mBuffer; +}; + +//! +//! \class LogStreamConsumer +//! \brief Convenience object used to facilitate use of C++ stream syntax when logging messages. +//! Order of base classes is LogStreamConsumerBase and then std::ostream. +//! This is because the LogStreamConsumerBase class is used to initialize the LogStreamConsumerBuffer member field +//! in LogStreamConsumer and then the address of the buffer is passed to std::ostream. +//! This is necessary to prevent the address of an uninitialized buffer from being passed to std::ostream. +//! Please do not change the order of the parent classes. +//! +class LogStreamConsumer : protected LogStreamConsumerBase, public std::ostream +{ +public: + //! \brief Creates a LogStreamConsumer which logs messages with level severity. + //! Reportable severity determines if the messages are severe enough to be logged. + LogStreamConsumer(Severity reportableSeverity, Severity severity) + : LogStreamConsumerBase(severityOstream(severity), severityPrefix(severity), severity <= reportableSeverity) + , std::ostream(&mBuffer) // links the stream buffer with the stream + , mShouldLog(severity <= reportableSeverity) + , mSeverity(severity) + { + } + + LogStreamConsumer(LogStreamConsumer&& other) + : LogStreamConsumerBase(severityOstream(other.mSeverity), severityPrefix(other.mSeverity), other.mShouldLog) + , std::ostream(&mBuffer) // links the stream buffer with the stream + , mShouldLog(other.mShouldLog) + , mSeverity(other.mSeverity) + { + } + + void setReportableSeverity(Severity reportableSeverity) + { + mShouldLog = mSeverity <= reportableSeverity; + mBuffer.setShouldLog(mShouldLog); + } + +private: + static std::ostream& severityOstream(Severity severity) + { + return severity >= Severity::kINFO ? std::cout : std::cerr; + } + + static std::string severityPrefix(Severity severity) + { + switch (severity) + { + case Severity::kINTERNAL_ERROR: return "[F] "; + case Severity::kERROR: return "[E] "; + case Severity::kWARNING: return "[W] "; + case Severity::kINFO: return "[I] "; + case Severity::kVERBOSE: return "[V] "; + default: assert(0); return ""; + } + } + + bool mShouldLog; + Severity mSeverity; +}; + +//! \class Logger +//! +//! \brief Class which manages logging of TensorRT tools and samples +//! +//! \details This class provides a common interface for TensorRT tools and samples to log information to the console, +//! and supports logging two types of messages: +//! +//! - Debugging messages with an associated severity (info, warning, error, or internal error/fatal) +//! - Test pass/fail messages +//! +//! The advantage of having all samples use this class for logging as opposed to emitting directly to stdout/stderr is +//! that the logic for controlling the verbosity and formatting of sample output is centralized in one location. +//! +//! In the future, this class could be extended to support dumping test results to a file in some standard format +//! (for example, JUnit XML), and providing additional metadata (e.g. timing the duration of a test run). +//! +//! TODO: For backwards compatibility with existing samples, this class inherits directly from the nvinfer1::ILogger +//! interface, which is problematic since there isn't a clean separation between messages coming from the TensorRT +//! library and messages coming from the sample. +//! +//! In the future (once all samples are updated to use Logger::getTRTLogger() to access the ILogger) we can refactor the +//! class to eliminate the inheritance and instead make the nvinfer1::ILogger implementation a member of the Logger +//! object. + +class Logger : public nvinfer1::ILogger +{ +public: + Logger(Severity severity = Severity::kWARNING) + : mReportableSeverity(severity) + { + } + + //! + //! \enum TestResult + //! \brief Represents the state of a given test + //! + enum class TestResult + { + kRUNNING, //!< The test is running + kPASSED, //!< The test passed + kFAILED, //!< The test failed + kWAIVED //!< The test was waived + }; + + //! + //! \brief Forward-compatible method for retrieving the nvinfer::ILogger associated with this Logger + //! \return The nvinfer1::ILogger associated with this Logger + //! + //! TODO Once all samples are updated to use this method to register the logger with TensorRT, + //! we can eliminate the inheritance of Logger from ILogger + //! + nvinfer1::ILogger& getTRTLogger() + { + return *this; + } + + //! + //! \brief Implementation of the nvinfer1::ILogger::log() virtual method + //! + //! Note samples should not be calling this function directly; it will eventually go away once we eliminate the + //! inheritance from nvinfer1::ILogger + //! + void log(Severity severity, const char* msg) override + { + LogStreamConsumer(mReportableSeverity, severity) << "[TRT] " << std::string(msg) << std::endl; + } + + //! + //! \brief Method for controlling the verbosity of logging output + //! + //! \param severity The logger will only emit messages that have severity of this level or higher. + //! + void setReportableSeverity(Severity severity) + { + mReportableSeverity = severity; + } + + //! + //! \brief Opaque handle that holds logging information for a particular test + //! + //! This object is an opaque handle to information used by the Logger to print test results. + //! The sample must call Logger::defineTest() in order to obtain a TestAtom that can be used + //! with Logger::reportTest{Start,End}(). + //! + class TestAtom + { + public: + TestAtom(TestAtom&&) = default; + + private: + friend class Logger; + + TestAtom(bool started, const std::string& name, const std::string& cmdline) + : mStarted(started) + , mName(name) + , mCmdline(cmdline) + { + } + + bool mStarted; + std::string mName; + std::string mCmdline; + }; + + //! + //! \brief Define a test for logging + //! + //! \param[in] name The name of the test. This should be a string starting with + //! "TensorRT" and containing dot-separated strings containing + //! the characters [A-Za-z0-9_]. + //! For example, "TensorRT.sample_googlenet" + //! \param[in] cmdline The command line used to reproduce the test + // + //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). + //! + static TestAtom defineTest(const std::string& name, const std::string& cmdline) + { + return TestAtom(false, name, cmdline); + } + + //! + //! \brief A convenience overloaded version of defineTest() that accepts an array of command-line arguments + //! as input + //! + //! \param[in] name The name of the test + //! \param[in] argc The number of command-line arguments + //! \param[in] argv The array of command-line arguments (given as C strings) + //! + //! \return a TestAtom that can be used in Logger::reportTest{Start,End}(). + static TestAtom defineTest(const std::string& name, int argc, char const* const* argv) + { + auto cmdline = genCmdlineString(argc, argv); + return defineTest(name, cmdline); + } + + //! + //! \brief Report that a test has started. + //! + //! \pre reportTestStart() has not been called yet for the given testAtom + //! + //! \param[in] testAtom The handle to the test that has started + //! + static void reportTestStart(TestAtom& testAtom) + { + reportTestResult(testAtom, TestResult::kRUNNING); + assert(!testAtom.mStarted); + testAtom.mStarted = true; + } + + //! + //! \brief Report that a test has ended. + //! + //! \pre reportTestStart() has been called for the given testAtom + //! + //! \param[in] testAtom The handle to the test that has ended + //! \param[in] result The result of the test. Should be one of TestResult::kPASSED, + //! TestResult::kFAILED, TestResult::kWAIVED + //! + static void reportTestEnd(const TestAtom& testAtom, TestResult result) + { + assert(result != TestResult::kRUNNING); + assert(testAtom.mStarted); + reportTestResult(testAtom, result); + } + + static int reportPass(const TestAtom& testAtom) + { + reportTestEnd(testAtom, TestResult::kPASSED); + return EXIT_SUCCESS; + } + + static int reportFail(const TestAtom& testAtom) + { + reportTestEnd(testAtom, TestResult::kFAILED); + return EXIT_FAILURE; + } + + static int reportWaive(const TestAtom& testAtom) + { + reportTestEnd(testAtom, TestResult::kWAIVED); + return EXIT_SUCCESS; + } + + static int reportTest(const TestAtom& testAtom, bool pass) + { + return pass ? reportPass(testAtom) : reportFail(testAtom); + } + + Severity getReportableSeverity() const + { + return mReportableSeverity; + } + +private: + //! + //! \brief returns an appropriate string for prefixing a log message with the given severity + //! + static const char* severityPrefix(Severity severity) + { + switch (severity) + { + case Severity::kINTERNAL_ERROR: return "[F] "; + case Severity::kERROR: return "[E] "; + case Severity::kWARNING: return "[W] "; + case Severity::kINFO: return "[I] "; + case Severity::kVERBOSE: return "[V] "; + default: assert(0); return ""; + } + } + + //! + //! \brief returns an appropriate string for prefixing a test result message with the given result + //! + static const char* testResultString(TestResult result) + { + switch (result) + { + case TestResult::kRUNNING: return "RUNNING"; + case TestResult::kPASSED: return "PASSED"; + case TestResult::kFAILED: return "FAILED"; + case TestResult::kWAIVED: return "WAIVED"; + default: assert(0); return ""; + } + } + + //! + //! \brief returns an appropriate output stream (cout or cerr) to use with the given severity + //! + static std::ostream& severityOstream(Severity severity) + { + return severity >= Severity::kINFO ? std::cout : std::cerr; + } + + //! + //! \brief method that implements logging test results + //! + static void reportTestResult(const TestAtom& testAtom, TestResult result) + { + severityOstream(Severity::kINFO) << "&&&& " << testResultString(result) << " " << testAtom.mName << " # " + << testAtom.mCmdline << std::endl; + } + + //! + //! \brief generate a command line string from the given (argc, argv) values + //! + static std::string genCmdlineString(int argc, char const* const* argv) + { + std::stringstream ss; + for (int i = 0; i < argc; i++) + { + if (i > 0) + ss << " "; + ss << argv[i]; + } + return ss.str(); + } + + Severity mReportableSeverity; +}; + + +namespace +{ + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kVERBOSE +//! +//! Example usage: +//! +//! LOG_VERBOSE(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_VERBOSE(const Logger& logger) +{ + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kVERBOSE); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINFO +//! +//! Example usage: +//! +//! LOG_INFO(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_INFO(const Logger& logger) +{ + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINFO); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kWARNING +//! +//! Example usage: +//! +//! LOG_WARN(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_WARN(const Logger& logger) +{ + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kWARNING); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kERROR +//! +//! Example usage: +//! +//! LOG_ERROR(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_ERROR(const Logger& logger) +{ + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kERROR); +} + +//! +//! \brief produces a LogStreamConsumer object that can be used to log messages of severity kINTERNAL_ERROR +// ("fatal" severity) +//! +//! Example usage: +//! +//! LOG_FATAL(logger) << "hello world" << std::endl; +//! +inline LogStreamConsumer LOG_FATAL(const Logger& logger) +{ + return LogStreamConsumer(logger.getReportableSeverity(), Severity::kINTERNAL_ERROR); +} + +} // anonymous namespace + +#endif // TENSORRT_LOGGING_H diff --git a/refinedet/refinedet.cpp b/refinedet/refinedet.cpp new file mode 100755 index 0000000..00dff34 --- /dev/null +++ b/refinedet/refinedet.cpp @@ -0,0 +1,1023 @@ +#include +#include +#include +#include +#include +#include +#include "NvInfer.h" +#include "cuda_runtime_api.h" +#include "utils.h" +#include "logging.h" +#include "calibrator.h" +#include "configure.h" + +#include // One-stop header. +#include "torch/torch.h" +#include "torch/jit.h" + +using namespace nvinfer1; +static Logger gLogger; + +//Correct the rectangle area to prevent the image from crossing the boundary +void RoiCorrect(const cv::Mat &m, cv::Rect &r) +{ + if (r.x < 0) r.x = 0; + if (r.y < 0) r.y = 0; + + if(r.x >= m.cols-1) r.x=0; + if(r.y >= m.rows-1) r.y=0; + + if(r.width <= 0) r.width = 1; + if(r.height <= 0) r.height = 1; + + if(r.x + r.width > m.cols - 1) r.width = m.cols - 1 - r.x; + if(r.y + r.height > m.rows - 1) r.height = m.rows - 1 - r.y; +} + +// TensorRT weight files have a simple space delimited format: +// [type] [size] +std::map loadWeights(const std::string file) { + std::cout << "Loading weights: " << file << std::endl; + std::map weightMap; + + // Open weights file + std::ifstream input(file); + assert(input.is_open() && "Unable to load weight file."); + + // Read number of weight blobs + int32_t count; + input >> count; + assert(count > 0 && "Invalid weight map file."); + + while (count--) + { + Weights wt{DataType::kFLOAT, nullptr, 0}; + uint32_t size; + + // Read name and type of blob + std::string name; + input >> name >> std::dec >> size; + wt.type = DataType::kFLOAT; + + // Load blob + uint32_t* val = reinterpret_cast(malloc(sizeof(val) * size)); + for (uint32_t x = 0, y = size; x < y; ++x) + { + input >> std::hex >> val[x]; + } + wt.values = val; + + wt.count = size; + weightMap[name] = wt; + } + + return weightMap; +} + +IScaleLayer* addBatchNorm2d(INetworkDefinition *network, std::map& weightMap, ITensor& input, std::string lname, float eps) { + float *gamma = (float*)weightMap[lname + ".weight"].values; + float *beta = (float*)weightMap[lname + ".bias"].values; + float *mean = (float*)weightMap[lname + ".running_mean"].values; + float *var = (float*)weightMap[lname + ".running_var"].values; + int len = weightMap[lname + ".running_var"].count; + + float *scval = reinterpret_cast(malloc(sizeof(float) * len)); + for (int i = 0; i < len; i++) { + scval[i] = gamma[i] / sqrt(var[i] + eps); + } + Weights scale{DataType::kFLOAT, scval, len}; + + float *shval = reinterpret_cast(malloc(sizeof(float) * len)); + for (int i = 0; i < len; i++) { + shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps); + } + Weights shift{DataType::kFLOAT, shval, len}; + + float *pval = reinterpret_cast(malloc(sizeof(float) * len)); + for (int i = 0; i < len; i++) { + pval[i] = 1.0; + } + Weights power{DataType::kFLOAT, pval, len}; + + weightMap[lname + ".scale"] = scale; + weightMap[lname + ".shift"] = shift; + weightMap[lname + ".power"] = power; + IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power); + assert(scale_1); + return scale_1; +} + +//convBnLeaky(network, weightMap, *data, 32, 3, 1, 1, 0); +ILayer* convRelu(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int ksize, int s, int p,\ + int linx, const std::string pre_name = "vgg.", bool b_dilate = false) { + Weights emptywts{DataType::kFLOAT, nullptr, 0}; + + if (weightMap.count(pre_name + std::to_string(linx) + ".weight") == 0) + std::cout << "no key: " <addConvolutionNd(input, outch, DimsHW{ksize, ksize}, weightMap[pre_name + std::to_string(linx) + ".weight"], weightMap[pre_name + std::to_string(linx) + ".bias"]); + assert(conv1); + conv1->setStrideNd(DimsHW{s, s}); + conv1->setPaddingNd(DimsHW{p, p}); + if(true == b_dilate) + { + conv1->setDilation(DimsHW{3, 3}); + } + + auto lr = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU); + + return lr; +} + +//convBnLeaky(network, weightMap, *data, 32, 3, 1, 1, 0); +ILayer* convRelu_extras(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int ksize, int s, int p, const std::string weight_name, const std::string bias_name){ + + if (weightMap.count(weight_name) == 0) + std::cout << "no key: " <addConvolutionNd(input, outch, DimsHW{ksize, ksize}, weightMap[weight_name], weightMap[bias_name]); + assert(conv1); + conv1->setStrideNd(DimsHW{s, s}); + conv1->setPaddingNd(DimsHW{p, p}); + + auto lr = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU); + + return lr; +} + +//convBnLeaky(network, weightMap, *data, 32, 3, 1, 1, 0); +IConvolutionLayer* convReluconv_tcb0(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int ksize, int s, int p, int indx_0, int indx_1){ + + std::string name_w0 = "tcb0." + (std::string)std::to_string(indx_0) + ".weight"; + std::string name_b0 = "tcb0." + (std::string)std::to_string(indx_0) + ".bias"; + + std::string name_w1 = "tcb0." + (std::string)std::to_string(indx_1) + ".weight"; + std::string name_b1 = "tcb0." + (std::string)std::to_string(indx_1) + ".bias"; + + if (weightMap.count(name_w0) == 0) + std::cout << "no key: " <addConvolutionNd(input, outch, DimsHW{ksize, ksize}, weightMap[name_w0], weightMap[name_b0]); + assert(conv1); + conv1->setStrideNd(DimsHW{s, s}); + conv1->setPaddingNd(DimsHW{p, p}); + + auto lr = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU); + + IConvolutionLayer* conv2 = network->addConvolutionNd(*lr->getOutput(0), 256, DimsHW{3, 3}, weightMap[name_w1], weightMap[name_b1]); + assert(conv2); + conv2->setStrideNd(DimsHW{1, 1}); + conv2->setPaddingNd(DimsHW{1, 1}); + + return conv2; +} + +ILayer* ReluconvRelu_tcb2(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int ksize, int s, int p, int indx_0){ + auto lr = network->addActivation(input, ActivationType::kRELU); + + std::string name_w0 = "tcb2." + (std::string)std::to_string(indx_0) + ".weight"; + std::string name_b0 = "tcb2." + (std::string)std::to_string(indx_0) + ".bias"; + + if (weightMap.count(name_w0) == 0) + std::cout << "no key: " <addConvolutionNd(*lr->getOutput(0), outch, DimsHW{ksize, ksize}, weightMap[name_w0], weightMap[name_b0]); + assert(conv1); + conv1->setStrideNd(DimsHW{s, s}); + conv1->setPaddingNd(DimsHW{p, p}); + + auto lr1 = network->addActivation(*conv1->getOutput(0), ActivationType::kRELU); + return lr1; +} + +ILayer* conv_permutation(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int ksize, int s, int p, const std::string weight_name, const std::string bias_name) +{ + if (weightMap.count(weight_name) == 0) + std::cout << "no key: " <addConvolutionNd(input, outch, DimsHW{ksize, ksize}, weightMap[weight_name], weightMap[bias_name]); + assert(a0); + a0->setStrideNd(DimsHW{s, s}); + a0->setPaddingNd(DimsHW{p, p}); + + auto sfl = network->addShuffle(*a0->getOutput(0)); + sfl->setFirstTranspose(Permutation{1, 2, 0}); + + return sfl; +} + +ILayer* cat_4_tensor(INetworkDefinition *network, ILayer*tensor_0, ILayer*tensor_1, ILayer*tensor_2, ILayer*tensor_3) +{ + Dims dim_; + dim_.nbDims=1; + dim_.d[0]=-1; + //40 40 12 --->>40*40*12 + auto arm_loc_00 = network->addShuffle(*tensor_0->getOutput(0)); + assert(arm_loc_00); + arm_loc_00->setReshapeDimensions(dim_); + + //20 20 12 --->>20*20*12 + auto arm_loc_11 = network->addShuffle(*tensor_1->getOutput(0)); + assert(arm_loc_11); + arm_loc_11->setReshapeDimensions(dim_); //Dims2(-1, 1) + + //10 10 12 --->>10*10*12 + auto arm_loc_22 = network->addShuffle(*tensor_2->getOutput(0)); + assert(arm_loc_22); + arm_loc_22->setReshapeDimensions(dim_); + + //5 5 12 --->>5*5*12 + auto arm_loc_33 = network->addShuffle(*tensor_3->getOutput(0)); + assert(arm_loc_33); + arm_loc_33->setReshapeDimensions(dim_); + +// +// Dims dim0 = arm_loc_00->getOutput(0)->getDimensions(); +// std::cout <<"debug arm_loc_0 dim==" << dim0.d[0] << " " << dim0.d[1] << " " << dim0.d[2] << " " << dim0.d[3] << std::endl; +// Dims dim1 = arm_loc_11->getOutput(0)->getDimensions(); +// std::cout <<"debug arm_loc_1 dim==" << dim1.d[0] << " " << dim1.d[1] << " " << dim1.d[2] << " " << dim1.d[3] << std::endl; +// Dims dim2 = arm_loc_22->getOutput(0)->getDimensions(); +// std::cout <<"debug arm_loc_2 dim==" << dim2.d[0] << " " << dim2.d[1] << " " << dim2.d[2] << " " << dim2.d[3] << std::endl; +// Dims dim3 = arm_loc_33->getOutput(0)->getDimensions(); +// std::cout <<"debug arm_loc_3 dim==" << dim3.d[0] << " " << dim3.d[1] << " " << dim3.d[2] << " " << dim3.d[3] << std::endl; + + ITensor* arm_loc_t[] = {arm_loc_00->getOutput(0), arm_loc_11->getOutput(0), arm_loc_22->getOutput(0), arm_loc_33->getOutput(0)}; + auto arm_loc = network->addConcatenation(arm_loc_t, 4); + //[25500] + return arm_loc; +} + + +ILayer* reshapeSoftmax(INetworkDefinition *network, ITensor& input, int ch) { + //The input is one-dimensional[12750] + //reshape[XX,ch] + auto re1 = network->addShuffle(input); + assert(re1); + re1->setReshapeDimensions(Dims3(1, -1, ch)); //[1,6375,2]; +// re1->setReshapeDimensions(Dims2(-1, ch)); //[6375,2]; + + Dims dim0 = re1->getOutput(0)->getDimensions(); + std::cout <<"debug re1 dim==" << dim0.d[0] << " " << dim0.d[1] << " " << dim0.d[2] << " " << dim0.d[3] << std::endl; + +// return re1;///////////////////////////////////////// + + auto sm = network->addSoftMax(*re1->getOutput(0)); + sm->setAxes(1<<2); + assert(sm); + //And then reshape one-dimensional again, and it's the same shape as it came in + Dims dim_; + dim_.nbDims=1; + dim_.d[0]=-1; + auto re2 = network->addShuffle(*sm->getOutput(0)); + assert(re2); + re2->setReshapeDimensions(dim_); + + return re2; +} + +IScaleLayer* L2norm(INetworkDefinition *network, std::map& weightMap, ITensor& input, const std::string pre_name = "conv4_3_L2Norm.weight") +{ + //aa = x.pow(2) ## [1,512,40,40] + const static float pval1[3]{0.0, 1.0, 2.0}; + Weights wshift1{DataType::kFLOAT, pval1, 1}; + Weights wscale1{DataType::kFLOAT, pval1+1, 1}; + Weights wpower1{DataType::kFLOAT, pval1+2, 1}; + IScaleLayer* scale1 = network->addScale( + input, + ScaleMode::kUNIFORM, + wshift1, + wscale1, + wpower1); + assert(scale1); + + //bb = x.pow(2).sum(dim=1, keepdim=True) ## [1,1,40,40] + IReduceLayer* reduce1 = network->addReduce(*scale1->getOutput(0), + ReduceOperation::kSUM, + 1, + true); + assert(reduce1); + + //norm = x.pow(2).sum(dim=1, keepdim=True).sqrt()+self.eps # [1,1,40,40] + const static float pval2[3]{0.0, 1.0, 0.5}; + Weights wshift2{DataType::kFLOAT, pval2, 1}; + Weights wscale2{DataType::kFLOAT, pval2+1, 1}; + Weights wpower2{DataType::kFLOAT, pval2+2, 1}; + IScaleLayer* scale2 = network->addScale( + *reduce1->getOutput(0), + ScaleMode::kUNIFORM, + wshift2, + wscale2, + wpower2); + assert(scale2); + + // x = torch.div(x,norm) + IElementWiseLayer* ew2 = network->addElementWise(input, + *scale2->getOutput(0), + ElementWiseOperation::kDIV); + assert(ew2); + + //out = self.weight.unsqueeze(0).unsqueeze(2).unsqueeze(3).expand_as(x) * x + int len = weightMap[pre_name].count; + float* pval3 = reinterpret_cast(malloc(sizeof(float) * len)); + std::fill_n(pval3, len, 1.0); + Weights wpower3{DataType::kFLOAT, pval3, len}; + weightMap[pre_name + ".power3"] = wpower3; + + float* pval4 = reinterpret_cast(malloc(sizeof(float) * len)); + std::fill_n(pval4, len, 0.0); + Weights wpower4{DataType::kFLOAT, pval4, len}; + weightMap[pre_name + ".power4"] = wpower4; + + IScaleLayer* scale3 = network->addScale( + *ew2->getOutput(0), + ScaleMode::kCHANNEL, + wpower4, + weightMap[pre_name], + wpower3); + assert(scale3); + return scale3; +} + + +//convBnLeaky(network, weightMap, *data, 32, 3, 1, 1, 0); +ILayer* convBnLeaky(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int ksize, int s, int p, int linx) { + Weights emptywts{DataType::kFLOAT, nullptr, 0}; + IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ksize, ksize}, weightMap["module_list." + std::to_string(linx) + ".Conv2d.weight"], emptywts); + assert(conv1); + conv1->setStrideNd(DimsHW{s, s}); + conv1->setPaddingNd(DimsHW{p, p}); + + IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), "module_list." + std::to_string(linx) + ".BatchNorm2d", 1e-5); + + auto lr = network->addActivation(*bn1->getOutput(0), ActivationType::kLEAKY_RELU); + lr->setAlpha(0.1); + + return lr; +} + +// Creat the engine using only the API and not any parser. +ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) { + INetworkDefinition* network = builder->createNetworkV2(0U); + + ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{3, INPUT_H, INPUT_W}); + assert(data); + + std::map weightMap = loadWeights(path_wts); + Weights emptywts{DataType::kFLOAT, nullptr, 0}; + DimsHW maxpool_hw = DimsHW(2,2); + + auto lr0 = convRelu(network, weightMap, *data, 64, 3, 1, 1, 0); + auto lr1 = convRelu(network, weightMap, *lr0->getOutput(0), 64, 3, 1, 1, 2); + IPoolingLayer* pool1 = network->addPoolingNd(*lr1->getOutput(0), PoolingType::kMAX, DimsHW{2, 2}); + assert(pool1); + pool1->setStrideNd(DimsHW{2, 2}); + + auto lr2 = convRelu(network, weightMap, *pool1->getOutput(0), 128, 3, 1, 1, 5); + auto lr3 = convRelu(network, weightMap, *lr2->getOutput(0), 128, 3, 1, 1, 7); + IPoolingLayer* pool2 = network->addPoolingNd(*lr3->getOutput(0), PoolingType::kMAX, DimsHW{2, 2}); + assert(pool2); + pool2->setStrideNd(DimsHW{2, 2}); + + auto lr4 = convRelu(network, weightMap, *pool2->getOutput(0), 256, 3, 1, 1, 10); + auto lr5 = convRelu(network, weightMap, *lr4->getOutput(0), 256, 3, 1, 1, 12); + auto lr6 = convRelu(network, weightMap, *lr5->getOutput(0), 256, 3, 1, 1, 14); + IPoolingLayer* pool3 = network->addPoolingNd(*lr6->getOutput(0), PoolingType::kMAX, DimsHW{2, 2}); + assert(pool3); + pool3->setStrideNd(DimsHW{2, 2}); + + auto lr7 = convRelu(network, weightMap, *pool3->getOutput(0), 512, 3, 1, 1, 17); + auto lr8 = convRelu(network, weightMap, *lr7->getOutput(0), 512, 3, 1, 1, 19); + auto lr9 = convRelu(network, weightMap, *lr8->getOutput(0), 512, 3, 1, 1, 21); + IPoolingLayer* pool4 = network->addPoolingNd(*lr9->getOutput(0), PoolingType::kMAX, DimsHW{2, 2}); + assert(pool4); + pool4->setStrideNd(DimsHW{2, 2}); + + auto lr24 = convRelu(network, weightMap, *pool4->getOutput(0), 512, 3, 1, 1, 24); + auto lr26 = convRelu(network, weightMap, *lr24->getOutput(0), 512, 3, 1, 1, 26); + auto lr28 = convRelu(network, weightMap, *lr26->getOutput(0), 512, 3, 1, 1, 28); + IPoolingLayer* pool5 = network->addPoolingNd(*lr28->getOutput(0), PoolingType::kMAX, DimsHW{2, 2}); + assert(pool5); + pool5->setStrideNd(DimsHW{2, 2}); + + auto lr31 = convRelu(network, weightMap, *pool5->getOutput(0), 1024, 3, 1, 3, 31,"vgg.",true); + + //s_0 + auto out_conv4_3_L2Norm = L2norm(network, weightMap, *lr9->getOutput(0),"conv4_3_L2Norm.weight"); + //s_1 + auto out_conv5_3_L2Norm = L2norm(network, weightMap, *lr28->getOutput(0),"conv5_3_L2Norm.weight"); + + //s_2 + auto lr33 = convRelu(network, weightMap, *lr31->getOutput(0), 1024, 1, 1, 0, 33); + + auto extras0 = convRelu_extras(network, weightMap, *lr33->getOutput(0), 256, 1, 1, 0, "extras.0.weight", "extras.0.bias"); + //s_3 + auto extras1 = convRelu_extras(network, weightMap, *extras0->getOutput(0), 512, 3, 2, 1, "extras.1.weight", "extras.1.bias"); + + auto arm_loc_0 = conv_permutation(network, weightMap, *out_conv4_3_L2Norm->getOutput(0), 12, 3, 1, 1, "arm_loc.0.weight", "arm_loc.0.bias"); + auto arm_loc_1 = conv_permutation(network, weightMap, *out_conv5_3_L2Norm->getOutput(0), 12, 3, 1, 1, "arm_loc.1.weight", "arm_loc.1.bias"); + auto arm_loc_2 = conv_permutation(network, weightMap, *lr33->getOutput(0), 12, 3, 1, 1, "arm_loc.2.weight", "arm_loc.2.bias"); + auto arm_loc_3 = conv_permutation(network, weightMap, *extras1->getOutput(0), 12, 3, 1, 1, "arm_loc.3.weight", "arm_loc.3.bias"); + + auto arm_conf_0 = conv_permutation(network, weightMap, *out_conv4_3_L2Norm->getOutput(0), 6, 3, 1, 1, "arm_conf.0.weight", "arm_conf.0.bias"); + auto arm_conf_1 = conv_permutation(network, weightMap, *out_conv5_3_L2Norm->getOutput(0), 6, 3, 1, 1, "arm_conf.1.weight", "arm_conf.1.bias"); + auto arm_conf_2 = conv_permutation(network, weightMap, *lr33->getOutput(0), 6, 3, 1, 1, "arm_conf.2.weight", "arm_conf.2.bias"); + auto arm_conf_3 = conv_permutation(network, weightMap, *extras1->getOutput(0), 6, 3, 1, 1, "arm_conf.3.weight", "arm_conf.3.bias"); + + auto arm_loc = cat_4_tensor(network, arm_loc_0, arm_loc_1, arm_loc_2, arm_loc_3); + auto arm_conf = cat_4_tensor(network, arm_conf_0, arm_conf_1, arm_conf_2, arm_conf_3); + + auto ss_0 = convReluconv_tcb0(network, weightMap, *extras1->getOutput(0), 256, 3, 1, 1, 9, 11); + auto ss_00 = ReluconvRelu_tcb2(network, weightMap, *ss_0->getOutput(0), 256, 3, 1, 1, 10); + auto ss_1 = convReluconv_tcb0(network, weightMap, *lr33->getOutput(0), 256, 3, 1, 1, 6, 8); + + IDeconvolutionLayer* tcb1_2 = network->addDeconvolutionNd(*ss_00->getOutput(0), 256, DimsHW{2, 2}, weightMap["tcb1.2.weight"], weightMap["tcb1.2.bias"]); //nn.ConvTranspose2d(256, 256, 2, 2) + tcb1_2->setStrideNd(DimsHW{2, 2}); + assert(tcb1_2); + auto ss_1_add = network->addElementWise(*ss_1->getOutput(0), *tcb1_2->getOutput(0), ElementWiseOperation::kSUM); + auto ss_11 = ReluconvRelu_tcb2(network, weightMap, *ss_1_add->getOutput(0), 256, 3, 1, 1, 7); + + auto ss_2 = convReluconv_tcb0(network, weightMap, *out_conv5_3_L2Norm->getOutput(0), 256, 3, 1, 1, 3, 5); + IDeconvolutionLayer* tcb1_1 = network->addDeconvolutionNd(*ss_11->getOutput(0), 256, DimsHW{2, 2}, weightMap["tcb1.1.weight"], weightMap["tcb1.1.bias"]); //nn.ConvTranspose2d(256, 256, 2, 2) + tcb1_1->setStrideNd(DimsHW{2, 2}); + assert(tcb1_1); + auto ss_2_add = network->addElementWise(*ss_2->getOutput(0), *tcb1_1->getOutput(0), ElementWiseOperation::kSUM); + auto ss_22 = ReluconvRelu_tcb2(network, weightMap, *ss_2_add->getOutput(0), 256, 3, 1, 1, 4); + + auto ss_3 = convReluconv_tcb0(network, weightMap, *out_conv4_3_L2Norm->getOutput(0), 256, 3, 1, 1, 0, 2); + IDeconvolutionLayer* tcb1_0 = network->addDeconvolutionNd(*ss_22->getOutput(0), 256, DimsHW{2, 2}, weightMap["tcb1.0.weight"], weightMap["tcb1.0.bias"]); //nn.ConvTranspose2d(256, 256, 2, 2) + tcb1_0->setStrideNd(DimsHW{2, 2}); + assert(tcb1_0); + auto ss_3_add = network->addElementWise(*ss_3->getOutput(0), *tcb1_0->getOutput(0), ElementWiseOperation::kSUM); + auto ss_33 = ReluconvRelu_tcb2(network, weightMap, *ss_3_add->getOutput(0), 256, 3, 1, 1, 1); + + auto odm_loc_0 = conv_permutation(network, weightMap, *ss_33->getOutput(0), 12, 3, 1, 1, "odm_loc.0.weight", "odm_loc.0.bias"); + auto odm_loc_1 = conv_permutation(network, weightMap, *ss_22->getOutput(0), 12, 3, 1, 1, "odm_loc.1.weight", "odm_loc.1.bias"); + auto odm_loc_2 = conv_permutation(network, weightMap, *ss_11->getOutput(0), 12, 3, 1, 1, "odm_loc.2.weight", "odm_loc.2.bias"); + auto odm_loc_3 = conv_permutation(network, weightMap, *ss_00->getOutput(0), 12, 3, 1, 1, "odm_loc.3.weight", "odm_loc.3.bias"); + + auto odm_conf_0 = conv_permutation(network, weightMap, *ss_33->getOutput(0), 3 * num_class, 3, 1, 1, "odm_conf.0.weight", "odm_conf.0.bias"); + auto odm_conf_1 = conv_permutation(network, weightMap, *ss_22->getOutput(0), 3 * num_class, 3, 1, 1, "odm_conf.1.weight", "odm_conf.1.bias"); + auto odm_conf_2 = conv_permutation(network, weightMap, *ss_11->getOutput(0), 3 * num_class, 3, 1, 1, "odm_conf.2.weight", "odm_conf.2.bias"); + auto odm_conf_3 = conv_permutation(network, weightMap, *ss_00->getOutput(0), 3 * num_class, 3, 1, 1, "odm_conf.3.weight", "odm_conf.3.bias"); + + auto odm_loc = cat_4_tensor(network, odm_loc_0, odm_loc_1, odm_loc_2, odm_loc_3); + auto odm_conf = cat_4_tensor(network, odm_conf_0, odm_conf_1, odm_conf_2, odm_conf_3); + + //25500 + Dims dim = arm_loc->getOutput(0)->getDimensions(); + std::cout <<"debug arm_loc dim==" << dim.d[0] << " " << dim.d[1] << " " << dim.d[2] << " " << dim.d[3] << std::endl; + arm_loc->getOutput(0)->setName(OUTPUT_BLOB_NAME_arm_loc); + network->markOutput(*arm_loc->getOutput(0)); + + auto arm_conf_111 = reshapeSoftmax(network, *arm_conf->getOutput(0), 2); + //12750 + Dims dim2 = arm_conf_111->getOutput(0)->getDimensions(); + std::cout <<"debug arm_conf dim==" << dim2.d[0] << " " << dim2.d[1] << " " << dim2.d[2] << " " << dim2.d[3] << std::endl; + arm_conf_111->getOutput(0)->setName(OUTPUT_BLOB_NAME_arm_conf); + network->markOutput(*arm_conf_111->getOutput(0)); + + //25500 + Dims dim3 = odm_loc->getOutput(0)->getDimensions(); + std::cout <<"debug odm_loc dim==" << dim3.d[0] << " " << dim3.d[1] << " " << dim3.d[2] << " " << dim3.d[3] << std::endl; + odm_loc->getOutput(0)->setName(OUTPUT_BLOB_NAME_odm_loc); + network->markOutput(*odm_loc->getOutput(0)); + + //159375 + Dims dim4 = odm_conf->getOutput(0)->getDimensions(); + odm_conf = reshapeSoftmax(network, *odm_conf->getOutput(0), 25); + std::cout <<"debug odm_conf dim==" << dim4.d[0] << " " << dim4.d[1] << " " << dim4.d[2] << " " << dim4.d[3] << std::endl; + odm_conf->getOutput(0)->setName(OUTPUT_BLOB_NAME_odm_conf); + network->markOutput(*odm_conf->getOutput(0)); + + builder->setMaxBatchSize(maxBatchSize); + config->setMaxWorkspaceSize(16 * (1 << 20)); // 16MB + +#if defined(USE_FP16) + config->setFlag(BuilderFlag::kFP16); +#elif defined(USE_INT8) + std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl; + assert(builder->platformHasFastInt8()); + config->setFlag(BuilderFlag::kINT8); + Int8EntropyCalibrator2 *calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME); + config->setInt8Calibrator(calibrator); +#endif + + std::cout << "Building engine, please wait for a while..." << std::endl; + ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); + std::cout << "Build engine successfully!" << std::endl; + // Don't need the network any more + network->destroy(); + + // Release host memory + for (auto& mem : weightMap) + { + free((void*) (mem.second.values)); + } + + return engine; +} + +void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) { + // Create builder + IBuilder* builder = createInferBuilder(gLogger); + IBuilderConfig* config = builder->createBuilderConfig(); + + // Create model to populate the network, then set the outputs and create an engine + ICudaEngine* engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT); + assert(engine != nullptr); + + // Serialize the engine + (*modelStream) = engine->serialize(); + + // Close everything down + engine->destroy(); + builder->destroy(); +} + +torch::Tensor PriorBox() +{ + std::vector mean; + std::vector feature_maps = {40,20,10,5}; + int image_size = 320; + std::vector steps = {8,16,32,64}; + std::vector min_sizes = {32,64,128,256}; + std::vector aspect_ratios = {2,2,2,2}; + for(int k=0;k variance({0.1,0.2}); + torch::Tensor top_2 = torch::tensor({0,1}).cuda().to(torch::kLong); + torch::Tensor bottom_2 = torch::tensor({2,3}).cuda().to(torch::kLong); + + auto c1 = _prior.index_select(1,top_2)+_loc.index_select(1,top_2).mul(variance[0])*_prior.index_select(1,bottom_2); + auto c2 = _prior.index_select(1,bottom_2)*torch::exp(_loc.index_select(1,bottom_2)*variance[1]); + auto _retv = torch::cat({c1,c2},1); + if(b_form_pt) + { + auto c3 = _retv.index_select(1,top_2)-_retv.index_select(1,bottom_2).div(2); + auto c4 = c3 + _retv.index_select(1,bottom_2); + return torch::cat({c3,c4},1); + } else + { + return _retv; + } + +} + +torch::Tensor center(torch::Tensor retv) +{ + auto c1 = retv.select(1,0).unsqueeze(1); + auto c2 = retv.select(1,1).unsqueeze(1); + auto c3 = retv.select(1,2).unsqueeze(1); + auto c4 = retv.select(1,3).unsqueeze(1); + + auto _retv = torch::cat({(c1+c3).div(2),(c2+c4).div(2),c3-c1,c4-c2},1); + return _retv; +} + +bool nms(const torch::Tensor& boxes, const torch::Tensor& scores, torch::Tensor &keep, int &count,float overlap, int top_k) +{ + count =0; + keep = torch::zeros({scores.size(0)}).to(torch::kLong).to(scores.device()); + if(0 == boxes.numel()) + { + return false; + } + + torch::Tensor x1 = boxes.select(1,0).clone(); + torch::Tensor y1 = boxes.select(1,1).clone(); + torch::Tensor x2 = boxes.select(1,2).clone(); + torch::Tensor y2 = boxes.select(1,3).clone(); + torch::Tensor area = (x2-x1)*(y2-y1); + // std::cout< sort_ret = torch::sort(scores.unsqueeze(1), 0, 0); + torch::Tensor v = std::get<0>(sort_ret).squeeze(1).to(scores.device()); + torch::Tensor idx = std::get<1>(sort_ret).squeeze(1).to(scores.device()); + + int num_ = idx.size(0); + if(num_ > top_k) //python:idx = idx[-top_k:] + { + idx = idx.slice(0,num_-top_k,num_).clone(); + } + torch::Tensor xx1,yy1,xx2,yy2,w,h; + while(idx.numel() > 0) + { + auto i = idx[-1]; + keep[count] = i; + count += 1; + if(1 == idx.size(0)) + { + break; + } + idx = idx.slice(0,0,idx.size(0)-1).clone(); + + xx1 = x1.index_select(0,idx); + yy1 = y1.index_select(0,idx); + xx2 = x2.index_select(0,idx); + yy2 = y2.index_select(0,idx); + + xx1 = xx1.clamp(x1[i].item().toFloat(),INT_MAX*1.0); + yy1 = yy1.clamp(y1[i].item().toFloat(),INT_MAX*1.0); + xx2 = xx2.clamp(INT_MIN*1.0,x2[i].item().toFloat()); + yy2 = yy2.clamp(INT_MIN*1.0,y2[i].item().toFloat()); + + w = xx2 - xx1; + h = yy2 - yy1; + + w = w.clamp(0,INT_MAX); + h = h.clamp(0,INT_MAX); + + torch::Tensor inter = w * h; + torch::Tensor rem_areas = area.index_select(0,idx); + + torch::Tensor union_ = (rem_areas - inter) + area[i]; + torch::Tensor Iou = inter * 1.0 / union_; + torch::Tensor index_small = Iou < overlap; + auto mask_idx = torch::nonzero(index_small).squeeze(); + idx = idx.index_select(0,mask_idx);//pthon: idx = idx[IoU.le(overlap)] + } + return true; +} + +void doInference(IExecutionContext& context, void* buffers[], cudaStream_t &stream, float* input, std::vector> &detections) { + auto start_infer = std::chrono::system_clock::now(); + detections.clear(); + int batchSize = 1; + const ICudaEngine& engine = context.getEngine(); + + // Pointers to input and output device buffers to pass to engine. + // Engine requires exactly IEngine::getNbBindings() number of buffers. +// std::cout<<"engine.getNbBindings()==="<(end_infer - start_infer).count(); + std::cout <<"time consume context.enqueue===" << during_time << "ms" << std::endl; + + auto start_houchuli = std::chrono::system_clock::now(); + int m_prior_size = 6375; + torch::Tensor m_prior = PriorBox(); + torch::Tensor arm_loc = torch::from_blob(buffers[outputIndex_arm_loc],{m_prior_size,4}).cuda().toType(torch::kFloat64).unsqueeze(0); + torch::Tensor arm_conf = torch::from_blob(buffers[outputIndex_arm_conf],{m_prior_size,2}).cuda().toType(torch::kFloat64).unsqueeze(0); + torch::Tensor odm_loc = torch::from_blob(buffers[outputIndex_odm_loc],{m_prior_size,4}).cuda().toType(torch::kFloat64).unsqueeze(0); + torch::Tensor odm_conf = torch::from_blob(buffers[outputIndex_odm_conf],{m_prior_size,25}).cuda().toType(torch::kFloat64).unsqueeze(0); + + float obj_threshed = 0.01; + torch::Tensor arm_object_conf = arm_conf.squeeze(0).select(1,1); + torch::Tensor object_index = arm_object_conf > obj_threshed; + object_index=object_index.unsqueeze(1); + + torch::Tensor object_index_1 = object_index.expand_as(odm_conf.squeeze(0)).toType(torch::kFloat64); + auto filter_odm_conf = odm_conf.squeeze(0).toType(torch::kFloat64) * object_index_1; + torch::Tensor conf_preds_ = filter_odm_conf.clone().toType(torch::kFloat64); + torch::Tensor conf_preds = conf_preds_.transpose(1,0).toType(torch::kFloat64); + torch::Tensor default_m = decode(arm_loc[0],m_prior); +// default_m = center(default_m); + bool b_form_pt = true; + torch::Tensor decode_boxes_m = decode(odm_loc[0],default_m,b_form_pt);//6375,4 + + float conf_thresh = 0.01; + float mask_thresh = 0.01; + + torch::Tensor result_out; + for(int i=1;i<25;i++) + { + torch::Tensor c_mask_m = conf_preds[i] > mask_thresh; + torch::Tensor nonzero_index = torch::nonzero(c_mask_m); + torch::Tensor score_m = torch::index_select(conf_preds[i],0,nonzero_index.squeeze(1)); + torch::Tensor boxes_m = torch::index_select(decode_boxes_m,0,nonzero_index.squeeze(1)); + + torch::Tensor keep; + int count = 0; + float overlap = 0.45; + int top_k=1000; + nms(boxes_m, score_m, keep, count, overlap, top_k); + if(0 == count) { continue; } + + keep = keep.slice(0,0,count).clone(); + torch::Tensor score_my = score_m.index_select(0,keep); + torch::Tensor boxes_my = boxes_m.index_select(0,keep); + + if(score_my[0].item().toFloat() < conf_thresh) + { + continue; + } +// boxes_my.select(1,0).mul_(width); +// boxes_my.select(1,1).mul_(height); +// boxes_my.select(1,2).mul_(width); +// boxes_my.select(1,3).mul_(height); + torch::Tensor label_tensor = torch::full_like(score_my.unsqueeze(1),i); + torch::Tensor result_ = torch::cat({boxes_my.toType(torch::kFloat64),score_my.unsqueeze(1).toType(torch::kFloat64),label_tensor.toType(torch::kFloat64)},1); + if(0 == result_out.numel()) + { + result_out = result_.clone(); + }else + { + result_out = torch::cat({result_out,result_},0);//Splicing by line + } + } + if(0 == result_out.numel()) { std::cout<<"libtorch refinedet obj_small: nothing detect!"<(); + for(int i=0;i v_detections; + v_detections.push_back(0); //image_id + v_detections.push_back(id_label); //label + v_detections.push_back(score); //score + v_detections.push_back(x1); //xmin + v_detections.push_back(y1); //ymin + v_detections.push_back(x2); //xmax + v_detections.push_back(y2); //ymax + detections.push_back(v_detections); + } + cudaDeviceSynchronize(); + auto end_houchuli = std::chrono::system_clock::now(); + double during_time_houchuli = std::chrono::duration_cast(end_houchuli - start_houchuli).count(); + std::cout <<"time consume houchuli===" << during_time_houchuli << "ms" << std::endl; +} + +void base_transform(const cv::Mat &m_src,float *data) +{ + cv::Mat image; + cv::resize(m_src,image,cv::Size(INPUT_W,INPUT_H)); + if(1 == image.channels()) { cv::cvtColor(image,image,CV_GRAY2BGR); } + + for(int i=0;i(i); //Get the first address of the row pointer + for(int j=0;j(modelStream->data()), modelStream->size()); + modelStream->destroy(); + return 0; + +#elif defined INFER + std::ifstream file(path_engine, std::ios::binary); + if (file.good()) { + file.seekg(0, file.end); + size = file.tellg(); + file.seekg(0, file.beg); + trtModelStream = new char[size]; + assert(trtModelStream); + file.read(trtModelStream, size); + file.close(); + } + +#else + std::cerr << "arguments not right!" << std::endl; + std::cerr << "configure.h should difine SERIALIZE INFER" << std::endl; + std::cerr << "please check!" << std::endl; + return -1; +#endif + + std::vector file_names; + if (read_files_in_dir(p_dir_name, file_names) < 0) { + std::cout << "read_files_in_dir failed." << std::endl; + return -1; + } + + // prepare input data --------------------------- + float data[3 * INPUT_H * INPUT_W]; + + IRuntime* runtime = createInferRuntime(gLogger); //400M + assert(runtime != nullptr); + ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size); //777M + assert(engine != nullptr); + IExecutionContext* context = engine->createExecutionContext(); //971M + assert(context != nullptr); + delete[] trtModelStream; + + const int batchSize = 1; + const int inputIndex=0; + const int outputIndex_arm_loc=1; + const int outputIndex_arm_conf=3; + const int outputIndex_odm_loc=2; + const int outputIndex_odm_conf=4; + + //Initialize cuda memory: input and 4 output memory + void* buffers[5]; + // Create GPU buffers on device + CUDA_CHECK(cudaMalloc(&buffers[0], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float))); + + const int OUTPUT_SIZE_arm_loc = 25500; //40*40*12 + 20*20*12 + 10*10*12 + 5*5*12 = 25500 (Fixed value) + CUDA_CHECK(cudaMalloc(&buffers[outputIndex_arm_loc], batchSize * OUTPUT_SIZE_arm_loc * sizeof(float))); + + const int OUTPUT_SIZE_arm_conf = 12750; //40*40*6 + 20*20*6 + 10*10*6 + 5*5*6 = 12750 (Fixed value) + CUDA_CHECK(cudaMalloc(&buffers[outputIndex_arm_conf], batchSize * OUTPUT_SIZE_arm_conf * sizeof(float))); + + const int OUTPUT_SIZE_odm_loc = 25500; //40*40*12 + 20*20*12 + 10*10*12 + 5*5*12 = 25500 (Fixed value) + CUDA_CHECK(cudaMalloc(&buffers[outputIndex_odm_loc], batchSize * OUTPUT_SIZE_odm_loc * sizeof(float))); + + const int OUTPUT_SIZE_odm_conf = 159375; //40*40*(num_class*3) + 20*20**(num_class*3) + 10*10**(num_class*3) + 5*5**(num_class*3) //here num_class=25// =159375 + CUDA_CHECK(cudaMalloc(&buffers[outputIndex_odm_conf], batchSize * OUTPUT_SIZE_odm_conf * sizeof(float))); + + // Create stream + cudaStream_t stream; + CUDA_CHECK(cudaStreamCreate(&stream)); + + + int fcount = 0; + auto t_0 = std::chrono::steady_clock::now(); + for (auto f: file_names) { + fcount++; + std::cout << "\n" << fcount << " " << f << std::endl; + std::cout << std::string(p_dir_name) + "/" + f << std::endl; + + auto start_read = std::chrono::system_clock::now(); + cv::Mat img = cv::imread(std::string(p_dir_name) + "/" + f); + cudaDeviceSynchronize(); + auto end_read = std::chrono::system_clock::now(); + double during_time_read = std::chrono::duration_cast(end_read - start_read).count(); + std::cout <<"time consume during_time_read===" << during_time_read << "ms" << std::endl; + + if (img.empty()) continue; + + auto start_yuchuli = std::chrono::system_clock::now(); + base_transform(img,data); + cudaDeviceSynchronize(); + auto end_yuchuli = std::chrono::system_clock::now(); + double during_time_yuchuli = std::chrono::duration_cast(end_yuchuli - start_yuchuli).count(); + std::cout <<"time consume base_transform===" << during_time_yuchuli << "ms" << std::endl; + + auto start_doInfer = std::chrono::system_clock::now(); + std::vector> detections; + doInference(*context, buffers, stream, data, detections); + cudaDeviceSynchronize(); + auto end_doInfer = std::chrono::system_clock::now(); + double during_doinfer = std::chrono::duration_cast(end_doInfer - start_doInfer).count(); + std::cout <<"time consume doInference===" << during_doinfer << "ms" << std::endl; + + /* Print the detection results. */ + for (size_t i = 0; i < detections.size(); ++i) + { + const std::vector &d = detections[i]; + + CHECK_EQ(d.size(), 7); + const float score = d[2]; + + int label = int(d[1]); + if (label >= num_class || label < 0) + { + std::cout << "label_Error!" << std::endl; + continue; + } + if(score < TH) + { + continue; + } + cv::Rect r; + r.x = d[3] * img.cols; + r.y = d[4] * img.rows; + r.width = d[5] * img.cols - r.x; + r.height = d[6] * img.rows - r.y; + + RoiCorrect(img, r); + if(T_show) + { + cv::rectangle(img,r,cv::Scalar(255,0,0),2); + } + if (T_show == 0) + { + std::string name_1 = f.substr(0,f.size()-4); + std::string path_txt = save_path_txt + name_1 + ".txt"; + std::ofstream fout(path_txt); + fout << label_map[label] << " " << score << " " << r.x << " " << r.y << " " << r.x + r.width + << " " << r.y + r.height << std::endl; //使用自己的label + } + } + if(T_show) + { + cv::namedWindow("show",0); + cv::imshow("show",img); + cv::waitKey(0); + } + } + cudaStreamSynchronize(stream); + + // Release stream and buffers + cudaStreamDestroy(stream); + CUDA_CHECK(cudaFree(buffers[inputIndex])); + CUDA_CHECK(cudaFree(buffers[outputIndex_arm_loc])); + CUDA_CHECK(cudaFree(buffers[outputIndex_arm_conf])); + + CUDA_CHECK(cudaFree(buffers[outputIndex_odm_loc])); + CUDA_CHECK(cudaFree(buffers[outputIndex_odm_conf])); + + cudaDeviceSynchronize(); + auto ttt = std::chrono::duration_cast + (std::chrono::steady_clock::now() - t_0).count(); + std::cout << "all consume time="<destroy(); + engine->destroy(); + runtime->destroy(); + return 0; +} diff --git a/refinedet/utils.h b/refinedet/utils.h new file mode 100755 index 0000000..a074667 --- /dev/null +++ b/refinedet/utils.h @@ -0,0 +1,63 @@ +#ifndef __TRT_UTILS_H_ +#define __TRT_UTILS_H_ + +#include +#include +#include +#include +#include +#include + +#ifndef CUDA_CHECK + +#define CUDA_CHECK(callstr) \ + { \ + cudaError_t error_code = callstr; \ + if (error_code != cudaSuccess) { \ + std::cerr << "CUDA error " << error_code << " at " << __FILE__ << ":" << __LINE__; \ + assert(0); \ + } \ + } + +#endif + +namespace Tn +{ + template + void write(char*& buffer, const T& val) + { + *reinterpret_cast(buffer) = val; + buffer += sizeof(T); + } + + template + void read(const char*& buffer, T& val) + { + val = *reinterpret_cast(buffer); + buffer += sizeof(T); + } +} + +static inline int read_files_in_dir(const char *p_dir_name, std::vector &file_names) { + DIR *p_dir = opendir(p_dir_name); + if (p_dir == nullptr) { + return -1; + } + + struct dirent* p_file = nullptr; + while ((p_file = readdir(p_dir)) != nullptr) { + if (strcmp(p_file->d_name, ".") != 0 && + strcmp(p_file->d_name, "..") != 0) { + //std::string cur_file_name(p_dir_name); + //cur_file_name += "/"; + //cur_file_name += p_file->d_name; + std::string cur_file_name(p_file->d_name); + file_names.push_back(cur_file_name); + } + } + + closedir(p_dir); + return 0; +} + +#endif