diff --git a/yolov3-tiny/CMakeLists.txt b/yolov3-tiny/CMakeLists.txt
new file mode 100644
index 0000000..fff091e
--- /dev/null
+++ b/yolov3-tiny/CMakeLists.txt
@@ -0,0 +1,42 @@
+cmake_minimum_required(VERSION 2.6)
+
+project(yolov3-tiny)
+
+add_definitions(-std=c++11)
+
+option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
+set(CMAKE_CXX_STANDARD 11)
+set(CMAKE_BUILD_TYPE Debug)
+
+find_package(CUDA REQUIRED)
+
+set(CUDA_NVCC_PLAGS ${CUDA_NVCC_PLAGS};-std=c++11;-g;-G;-gencode;arch=compute_30;code=sm_30)
+
+include_directories(${PROJECT_SOURCE_DIR}/include)
+if (CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64")
+ message("embed_platform on")
+ include_directories(/usr/local/cuda/targets/aarch64-linux/include)
+ link_directories(/usr/local/cuda/targets/aarch64-linux/lib)
+else()
+ message("embed_platform off")
+ include_directories(/usr/local/cuda/include)
+ link_directories(/usr/local/cuda/lib64)
+endif()
+
+
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")
+
+#cuda_add_library(leaky ${PROJECT_SOURCE_DIR}/leaky.cu)
+cuda_add_library(yololayer SHARED ${PROJECT_SOURCE_DIR}/yololayer.cu)
+
+find_package(OpenCV)
+include_directories(OpenCV_INCLUDE_DIRS)
+
+add_executable(yolov3-tiny ${PROJECT_SOURCE_DIR}/yolov3-tiny.cpp)
+target_link_libraries(yolov3-tiny nvinfer)
+target_link_libraries(yolov3-tiny cudart)
+target_link_libraries(yolov3-tiny yololayer)
+target_link_libraries(yolov3-tiny ${OpenCV_LIBS})
+
+add_definitions(-O2 -pthread)
+
diff --git a/yolov3-tiny/README.md b/yolov3-tiny/README.md
new file mode 100644
index 0000000..48bce28
--- /dev/null
+++ b/yolov3-tiny/README.md
@@ -0,0 +1,49 @@
+# yolov3-tiny
+
+The Pytorch implementation is [ultralytics/yolov3](https://github.com/ultralytics/yolov3).
+
+## Excute:
+
+```
+1. generate yolov3-tiny.wts from pytorch implementation with yolov3-tiny.cfg and yolov3-tiny.weights
+
+git clone https://github.com/ultralytics/yolov3.git
+// download its weights 'yolov3-tiny.pt' or 'yolov3-tiny.weights'
+// put tensorrtx/yolov3-tiny/gen_wts.py into ultralytics/yolov3 and run
+python gen_wts.py yolov3-tiny.weights
+// a file 'yolov3-tiny.wts' will be generated.
+
+2. put yolov3-tiny.wts into tensorrtx/yolov3-tiny, build and run
+
+// go to tensorrtx/yolov3-tiny
+mkdir build
+cd build
+cmake ..
+make
+sudo ./yolov3-tiny -s // serialize model to plan file i.e. 'yolov3-tiny.engine'
+sudo ./yolov3-tiny -d ../../yolov3-spp/samples // deserialize plan file and run inference, the images in samples will be processed.
+
+3. check the images generated, as follows. _zidane.jpg and _bus.jpg
+```
+
+
+
+
+
+
+
+
+
+## Config
+
+- Input shape defined in yololayer.h
+- Number of classes defined in yololayer.h
+- FP16/FP32 can be selected by the macro in yolov3-tiny.cpp
+- GPU id can be selected by the macro in yolov3-tiny.cpp
+- NMS thresh in yolov3-tiny.cpp
+- BBox confidence thresh in yolov3-tiny.cpp
+
+## More Information
+
+See the readme in [home page.](https://github.com/wang-xinyu/tensorrtx)
+
diff --git a/yolov3-tiny/gen_wts.py b/yolov3-tiny/gen_wts.py
new file mode 100644
index 0000000..27ea15e
--- /dev/null
+++ b/yolov3-tiny/gen_wts.py
@@ -0,0 +1,24 @@
+import struct
+import sys
+from models import *
+from utils.utils import *
+
+model = Darknet('cfg/yolov3-tiny.cfg', (608, 608))
+weights = sys.argv[1]
+dev = '0'
+if weights.endswith('.pt'): # pytorch format
+ model.load_state_dict(torch.load(weights, map_location=device)['model'])
+else: # darknet format
+ load_darknet_weights(model, weights)
+model = model.eval()
+
+f = open('yolov3-tiny.wts', 'w')
+f.write('{}\n'.format(len(model.state_dict().keys())))
+for k, v in model.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')
+
diff --git a/yolov3-tiny/logging.h b/yolov3-tiny/logging.h
new file mode 100644
index 0000000..602b69f
--- /dev/null
+++ b/yolov3-tiny/logging.h
@@ -0,0 +1,503 @@
+/*
+ * 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/yolov3-tiny/utils.h b/yolov3-tiny/utils.h
new file mode 100644
index 0000000..0de663c
--- /dev/null
+++ b/yolov3-tiny/utils.h
@@ -0,0 +1,94 @@
+#ifndef __TRT_UTILS_H_
+#define __TRT_UTILS_H_
+
+#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
+{
+ class Profiler : public nvinfer1::IProfiler
+ {
+ public:
+ void printLayerTimes(int itrationsTimes)
+ {
+ float totalTime = 0;
+ for (size_t i = 0; i < mProfile.size(); i++)
+ {
+ printf("%-40.40s %4.3fms\n", mProfile[i].first.c_str(), mProfile[i].second / itrationsTimes);
+ totalTime += mProfile[i].second;
+ }
+ printf("Time over all layers: %4.3f\n", totalTime / itrationsTimes);
+ }
+ private:
+ typedef std::pair Record;
+ std::vector mProfile;
+
+ virtual void reportLayerTime(const char* layerName, float ms)
+ {
+ auto record = std::find_if(mProfile.begin(), mProfile.end(), [&](const Record& r){ return r.first == layerName; });
+ if (record == mProfile.end())
+ mProfile.push_back(std::make_pair(layerName, ms));
+ else
+ record->second += ms;
+ }
+ };
+
+ //Logger for TensorRT info/warning/errors
+ class Logger : public nvinfer1::ILogger
+ {
+ public:
+
+ Logger(): Logger(Severity::kWARNING) {}
+
+ Logger(Severity severity): reportableSeverity(severity) {}
+
+ void log(Severity severity, const char* msg) override
+ {
+ // suppress messages with severity enum value greater than the reportable
+ if (severity > reportableSeverity) return;
+
+ switch (severity)
+ {
+ case Severity::kINTERNAL_ERROR: std::cerr << "INTERNAL_ERROR: "; break;
+ case Severity::kERROR: std::cerr << "ERROR: "; break;
+ case Severity::kWARNING: std::cerr << "WARNING: "; break;
+ case Severity::kINFO: std::cerr << "INFO: "; break;
+ default: std::cerr << "UNKNOWN: "; break;
+ }
+ std::cerr << msg << std::endl;
+ }
+
+ Severity reportableSeverity{Severity::kWARNING};
+ };
+
+ 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);
+ }
+}
+
+#endif
\ No newline at end of file
diff --git a/yolov3-tiny/yololayer.cu b/yolov3-tiny/yololayer.cu
new file mode 100644
index 0000000..5c648fb
--- /dev/null
+++ b/yolov3-tiny/yololayer.cu
@@ -0,0 +1,260 @@
+#include
+#include "yololayer.h"
+#include "utils.h"
+
+using namespace Yolo;
+
+namespace nvinfer1
+{
+ YoloLayerPlugin::YoloLayerPlugin()
+ {
+ mClassCount = CLASS_NUM;
+ mYoloKernel.clear();
+ mYoloKernel.push_back(yolo1);
+ mYoloKernel.push_back(yolo2);
+
+ mKernelCount = mYoloKernel.size();
+ }
+
+ YoloLayerPlugin::~YoloLayerPlugin()
+ {
+ }
+
+ // create the plugin at runtime from a byte stream
+ YoloLayerPlugin::YoloLayerPlugin(const void* data, size_t length)
+ {
+ using namespace Tn;
+ const char *d = reinterpret_cast(data), *a = d;
+ read(d, mClassCount);
+ read(d, mThreadCount);
+ read(d, mKernelCount);
+ mYoloKernel.resize(mKernelCount);
+ auto kernelSize = mKernelCount*sizeof(YoloKernel);
+ memcpy(mYoloKernel.data(),d,kernelSize);
+ d += kernelSize;
+
+ assert(d == a + length);
+ }
+
+ void YoloLayerPlugin::serialize(void* buffer) const
+ {
+ using namespace Tn;
+ char* d = static_cast(buffer), *a = d;
+ write(d, mClassCount);
+ write(d, mThreadCount);
+ write(d, mKernelCount);
+ auto kernelSize = mKernelCount*sizeof(YoloKernel);
+ memcpy(d,mYoloKernel.data(),kernelSize);
+ d += kernelSize;
+
+ assert(d == a + getSerializationSize());
+ }
+
+ size_t YoloLayerPlugin::getSerializationSize() const
+ {
+ return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(Yolo::YoloKernel) * mYoloKernel.size();
+ }
+
+ int YoloLayerPlugin::initialize()
+ {
+ return 0;
+ }
+
+ Dims YoloLayerPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims)
+ {
+ //output the result to channel
+ int totalsize = MAX_OUTPUT_BBOX_COUNT * sizeof(Detection) / sizeof(float);
+
+ return Dims3(totalsize + 1, 1, 1);
+ }
+
+ // Set plugin namespace
+ void YoloLayerPlugin::setPluginNamespace(const char* pluginNamespace)
+ {
+ mPluginNamespace = pluginNamespace;
+ }
+
+ const char* YoloLayerPlugin::getPluginNamespace() const
+ {
+ return mPluginNamespace;
+ }
+
+ // Return the DataType of the plugin output at the requested index
+ DataType YoloLayerPlugin::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const
+ {
+ return DataType::kFLOAT;
+ }
+
+ // Return true if output tensor is broadcast across a batch.
+ bool YoloLayerPlugin::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const
+ {
+ return false;
+ }
+
+ // Return true if plugin can use input that is broadcast across batch without replication.
+ bool YoloLayerPlugin::canBroadcastInputAcrossBatch(int inputIndex) const
+ {
+ return false;
+ }
+
+ void YoloLayerPlugin::configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput)
+ {
+ }
+
+ // Attach the plugin object to an execution context and grant the plugin the access to some context resource.
+ void YoloLayerPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator)
+ {
+ }
+
+ // Detach the plugin object from its execution context.
+ void YoloLayerPlugin::detachFromContext() {}
+
+ const char* YoloLayerPlugin::getPluginType() const
+ {
+ return "YoloLayer_TRT";
+ }
+
+ const char* YoloLayerPlugin::getPluginVersion() const
+ {
+ return "1";
+ }
+
+ void YoloLayerPlugin::destroy()
+ {
+ delete this;
+ }
+
+ // Clone the plugin
+ IPluginV2IOExt* YoloLayerPlugin::clone() const
+ {
+ YoloLayerPlugin *p = new YoloLayerPlugin();
+ p->setPluginNamespace(mPluginNamespace);
+ return p;
+ }
+
+ __device__ float Logist(float data){ return 1.0f / (1.0f + expf(-data)); };
+
+ __global__ void CalDetection(const float *input, float *output,int noElements,
+ int yoloWidth,int yoloHeight,const float anchors[CHECK_COUNT*2],int classes,int outputElem) {
+
+ int idx = threadIdx.x + blockDim.x * blockIdx.x;
+ if (idx >= noElements) return;
+
+ int total_grid = yoloWidth * yoloHeight;
+ int bnIdx = idx / total_grid;
+ idx = idx - total_grid*bnIdx;
+ int info_len_i = 5 + classes;
+ const float* curInput = input + bnIdx * (info_len_i * total_grid * CHECK_COUNT);
+
+ for (int k = 0; k < 3; ++k) {
+ int class_id = 0;
+ float max_cls_prob = 0.0;
+ for (int i = 5; i < info_len_i; ++i) {
+ float p = Logist(curInput[idx + k * info_len_i * total_grid + i * total_grid]);
+ if (p > max_cls_prob) {
+ max_cls_prob = p;
+ class_id = i - 5;
+ }
+ }
+ float box_prob = Logist(curInput[idx + k * info_len_i * total_grid + 4 * total_grid]);
+ if (max_cls_prob < IGNORE_THRESH || box_prob < IGNORE_THRESH) continue;
+
+ float *res_count = output + bnIdx*outputElem;
+ int count = (int)atomicAdd(res_count, 1);
+ if (count >= MAX_OUTPUT_BBOX_COUNT) return;
+ char* data = (char * )res_count + sizeof(float) + count*sizeof(Detection);
+ Detection* det = (Detection*)(data);
+
+ int row = idx / yoloWidth;
+ int col = idx % yoloWidth;
+
+ //Location
+ det->bbox[0] = (col + Logist(curInput[idx + k * info_len_i * total_grid + 0 * total_grid])) * INPUT_W / yoloWidth;
+ det->bbox[1] = (row + Logist(curInput[idx + k * info_len_i * total_grid + 1 * total_grid])) * INPUT_H / yoloHeight;
+ det->bbox[2] = expf(curInput[idx + k * info_len_i * total_grid + 2 * total_grid]) * anchors[2*k];
+ det->bbox[3] = expf(curInput[idx + k * info_len_i * total_grid + 3 * total_grid]) * anchors[2*k + 1];
+ det->det_confidence = box_prob;
+ det->class_id = class_id;
+ det->class_confidence = max_cls_prob;
+ }
+ }
+
+ void YoloLayerPlugin::forwardGpu(const float *const * inputs, float* output, cudaStream_t stream, int batchSize) {
+ void* devAnchor;
+ size_t AnchorLen = sizeof(float)* CHECK_COUNT*2;
+ CUDA_CHECK(cudaMalloc(&devAnchor,AnchorLen));
+
+ int outputElem = 1 + MAX_OUTPUT_BBOX_COUNT * sizeof(Detection) / sizeof(float);
+
+ for(int idx = 0 ; idx < batchSize; ++idx) {
+ CUDA_CHECK(cudaMemset(output + idx*outputElem, 0, sizeof(float)));
+ }
+ int numElem = 0;
+ for (unsigned int i = 0;i< mYoloKernel.size();++i)
+ {
+ const auto& yolo = mYoloKernel[i];
+ numElem = yolo.width*yolo.height*batchSize;
+ if (numElem < mThreadCount)
+ mThreadCount = numElem;
+ CUDA_CHECK(cudaMemcpy(devAnchor, yolo.anchors, AnchorLen, cudaMemcpyHostToDevice));
+ CalDetection<<< (yolo.width*yolo.height*batchSize + mThreadCount - 1) / mThreadCount, mThreadCount>>>
+ (inputs[i],output, numElem, yolo.width, yolo.height, (float *)devAnchor, mClassCount ,outputElem);
+ }
+
+ CUDA_CHECK(cudaFree(devAnchor));
+ }
+
+
+ int YoloLayerPlugin::enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream)
+ {
+ //assert(batchSize == 1);
+ //GPU
+ //CUDA_CHECK(cudaStreamSynchronize(stream));
+ forwardGpu((const float *const *)inputs, (float*)outputs[0], stream, batchSize);
+
+ return 0;
+ }
+
+ PluginFieldCollection YoloPluginCreator::mFC{};
+ std::vector YoloPluginCreator::mPluginAttributes;
+
+ YoloPluginCreator::YoloPluginCreator()
+ {
+ mPluginAttributes.clear();
+
+ mFC.nbFields = mPluginAttributes.size();
+ mFC.fields = mPluginAttributes.data();
+ }
+
+ const char* YoloPluginCreator::getPluginName() const
+ {
+ return "YoloLayer_TRT";
+ }
+
+ const char* YoloPluginCreator::getPluginVersion() const
+ {
+ return "1";
+ }
+
+ const PluginFieldCollection* YoloPluginCreator::getFieldNames()
+ {
+ return &mFC;
+ }
+
+ IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc)
+ {
+ YoloLayerPlugin* obj = new YoloLayerPlugin();
+ obj->setPluginNamespace(mNamespace.c_str());
+ return obj;
+ }
+
+ IPluginV2IOExt* YoloPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength)
+ {
+ // This object will be deleted when the network is destroyed, which will
+ // call MishPlugin::destroy()
+ YoloLayerPlugin* obj = new YoloLayerPlugin(serialData, serialLength);
+ obj->setPluginNamespace(mNamespace.c_str());
+ return obj;
+ }
+
+}
diff --git a/yolov3-tiny/yololayer.h b/yolov3-tiny/yololayer.h
new file mode 100644
index 0000000..b3282ff
--- /dev/null
+++ b/yolov3-tiny/yololayer.h
@@ -0,0 +1,150 @@
+#ifndef _YOLO_LAYER_H
+#define _YOLO_LAYER_H
+
+#include
+#include
+#include "NvInfer.h"
+
+namespace Yolo
+{
+ static constexpr int CHECK_COUNT = 3;
+ static constexpr float IGNORE_THRESH = 0.1f;
+ static constexpr int MAX_OUTPUT_BBOX_COUNT = 1000;
+ static constexpr int CLASS_NUM = 80;
+ static constexpr int INPUT_H = 608;
+ static constexpr int INPUT_W = 608;
+
+ struct YoloKernel
+ {
+ int width;
+ int height;
+ float anchors[CHECK_COUNT*2];
+ };
+
+ static constexpr YoloKernel yolo1 = {
+ INPUT_W / 32,
+ INPUT_H / 32,
+ {81,82, 135,169, 344,319}
+ };
+ static constexpr YoloKernel yolo2 = {
+ INPUT_W / 16,
+ INPUT_H / 16,
+ {23,27, 37,58, 81,82}
+ };
+
+ static constexpr int LOCATIONS = 4;
+ struct alignas(float) Detection{
+ //x y w h
+ float bbox[LOCATIONS];
+ float det_confidence;
+ float class_id;
+ float class_confidence;
+ };
+}
+
+
+namespace nvinfer1
+{
+ class YoloLayerPlugin: public IPluginV2IOExt
+ {
+ public:
+ explicit YoloLayerPlugin();
+ YoloLayerPlugin(const void* data, size_t length);
+
+ ~YoloLayerPlugin();
+
+ int getNbOutputs() const override
+ {
+ return 1;
+ }
+
+ Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override;
+
+ int initialize() override;
+
+ virtual void terminate() override {};
+
+ virtual size_t getWorkspaceSize(int maxBatchSize) const override { return 0;}
+
+ virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override;
+
+ virtual size_t getSerializationSize() const override;
+
+ virtual void serialize(void* buffer) const override;
+
+ bool supportsFormatCombination(int pos, const PluginTensorDesc* inOut, int nbInputs, int nbOutputs) const override {
+ return inOut[pos].format == TensorFormat::kLINEAR && inOut[pos].type == DataType::kFLOAT;
+ }
+
+ const char* getPluginType() const override;
+
+ const char* getPluginVersion() const override;
+
+ void destroy() override;
+
+ IPluginV2IOExt* clone() const override;
+
+ void setPluginNamespace(const char* pluginNamespace) override;
+
+ const char* getPluginNamespace() const override;
+
+ DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override;
+
+ bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const override;
+
+ bool canBroadcastInputAcrossBatch(int inputIndex) const override;
+
+ void attachToContext(
+ cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) override;
+
+ void configurePlugin(const PluginTensorDesc* in, int nbInput, const PluginTensorDesc* out, int nbOutput) override;
+
+ void detachFromContext() override;
+
+ private:
+ void forwardGpu(const float *const * inputs,float * output, cudaStream_t stream,int batchSize = 1);
+ int mClassCount;
+ int mKernelCount;
+ std::vector mYoloKernel;
+ int mThreadCount = 256;
+ const char* mPluginNamespace;
+ };
+
+ class YoloPluginCreator : public IPluginCreator
+ {
+ public:
+ YoloPluginCreator();
+
+ ~YoloPluginCreator() override = default;
+
+ const char* getPluginName() const override;
+
+ const char* getPluginVersion() const override;
+
+ const PluginFieldCollection* getFieldNames() override;
+
+ IPluginV2IOExt* createPlugin(const char* name, const PluginFieldCollection* fc) override;
+
+ IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override;
+
+ void setPluginNamespace(const char* libNamespace) override
+ {
+ mNamespace = libNamespace;
+ }
+
+ const char* getPluginNamespace() const override
+ {
+ return mNamespace.c_str();
+ }
+
+ private:
+ std::string mNamespace;
+ static PluginFieldCollection mFC;
+ static std::vector mPluginAttributes;
+ };
+
+
+
+};
+
+#endif
diff --git a/yolov3-tiny/yolov3-tiny.cpp b/yolov3-tiny/yolov3-tiny.cpp
new file mode 100644
index 0000000..adf7e6d
--- /dev/null
+++ b/yolov3-tiny/yolov3-tiny.cpp
@@ -0,0 +1,482 @@
+#include
+#include
+#include