diff --git a/yolov3-spp/.yolov3-spp-1cls.cpp.swp b/yolov3-spp/.yolov3-spp-1cls.cpp.swp new file mode 100644 index 0000000..6c08b1e Binary files /dev/null and b/yolov3-spp/.yolov3-spp-1cls.cpp.swp differ diff --git a/yolov3-spp/CMakeLists.txt b/yolov3-spp/CMakeLists.txt new file mode 100644 index 0000000..18762eb --- /dev/null +++ b/yolov3-spp/CMakeLists.txt @@ -0,0 +1,34 @@ +cmake_minimum_required(VERSION 2.6) + +project(yolov3-spp) + +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) +include_directories(/usr/local/cuda/targets/aarch64-linux/include) +link_directories(/usr/local/cuda/targets/aarch64-linux/lib) + +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-spp ${PROJECT_SOURCE_DIR}/plugin_factory.cpp ${PROJECT_SOURCE_DIR}/yolov3-spp.cpp) +target_link_libraries(yolov3-spp nvinfer nvinfer_plugin) +target_link_libraries(yolov3-spp cudart) +target_link_libraries(yolov3-spp yololayer) +target_link_libraries(yolov3-spp ${OpenCV_LIBS}) + +add_definitions(-O2 -pthread) + diff --git a/yolov3-spp/README.md b/yolov3-spp/README.md new file mode 100644 index 0000000..bf7eb36 --- /dev/null +++ b/yolov3-spp/README.md @@ -0,0 +1,36 @@ +# yolov3-spp + +For the Pytorch implementation, you can refer to [ultralytics/yolov3](https://github.com/ultralytics/yolov3) + +Following tricks are used in this yolov3-spp: + +- Yololayer plugin is different from the plugin used in [yolov3](https://github.com/wang-xinyu/tensorrtx/tree/master/yolov3). In this version, I reimplement the calculation of three yololayer plugins into one to improve speed. And the yololayer detect outputs are limited to maxmium 1000. The first number of output is number of targets in current image. +- Batchnorm layer, implemented by scale layer. + + + +## Excute: + +``` +// 1. generate yolov3-spp_ultralytics68.wts from pytorch implementation with yolov3-spp.cfg and ultralytics68.pt + +// 2. put yolov3-spp_ultralytics68.wts into yolov3-spp + +// 3. build and run + +cd yolov3-spp + +mkdir build + +cd build + +cmake .. + +make + +sudo ./yolov3-spp -s ../samples // serialize model to plan file i.e. 'yolov3-spp.engine' + +sudo ./yolov3-spp -d ../samples // deserialize plan file and run inference + +// 4. see if the output is same as pytorch implementation, and see the detect result in build +``` diff --git a/yolov3-spp/Utils.h b/yolov3-spp/Utils.h new file mode 100644 index 0000000..0de663c --- /dev/null +++ b/yolov3-spp/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-spp/YoloConfigs.h b/yolov3-spp/YoloConfigs.h new file mode 100644 index 0000000..2d2547b --- /dev/null +++ b/yolov3-spp/YoloConfigs.h @@ -0,0 +1,55 @@ +#ifndef _YOLO_CONFIGS_H_ +#define _YOLO_CONFIGS_H_ + + +namespace Yolo +{ + static constexpr int CHECK_COUNT = 3; + static constexpr float IGNORE_THRESH = 0.5f; + static constexpr int CLASS_NUM = 80; + static constexpr int INPUT_H = 256; + static constexpr int INPUT_W = 416; + + struct YoloKernel + { + int width; + int height; + float anchors[CHECK_COUNT*2]; + }; + + //YOLO 608 + YoloKernel yolo1 = { + 13, + 8, + {116,90, 156,198, 373,326} + }; + YoloKernel yolo2 = { + 26, + 16, + {30,61, 62,45, 59,119} + }; + YoloKernel yolo3 = { + 52, + 32, + {10,13, 16,30, 33,23} + }; + + //YOLO 416 + // YoloKernel yolo1 = { + // 13, + // 13, + // {116,90, 156,198, 373,326} + // }; + // YoloKernel yolo2 = { + // 26, + // 26, + // {30,61, 62,45, 59,119} + // }; + // YoloKernel yolo3 = { + // 52, + // 52, + // {10,13, 16,30, 33,23} + // }; +} + +#endif diff --git a/yolov3-spp/common.h b/yolov3-spp/common.h new file mode 100644 index 0000000..3b9c30c --- /dev/null +++ b/yolov3-spp/common.h @@ -0,0 +1,356 @@ +#ifndef _TRT_COMMON_H_ +#define _TRT_COMMON_H_ +#include "NvInfer.h" +#include "NvOnnxConfig.h" +#include "NvOnnxParser.h" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace std; + +#define CHECK(status) \ + do \ + { \ + auto ret = (status); \ + if (ret != 0) \ + { \ + std::cout << "Cuda failure: " << ret; \ + abort(); \ + } \ + } while (0) + +constexpr long double operator"" _GB(long double val) { return val * (1 << 30); } +constexpr long double operator"" _MB(long double val) { return val * (1 << 20); } +constexpr long double operator"" _KB(long double val) { return val * (1 << 10); } + +// These is necessary if we want to be able to write 1_GB instead of 1.0_GB. +// Since the return type is signed, -1_GB will work as expected. +constexpr long long int operator"" _GB(long long unsigned int val) { return val * (1 << 30); } +constexpr long long int operator"" _MB(long long unsigned int val) { return val * (1 << 20); } +constexpr long long int operator"" _KB(long long unsigned int val) { return val * (1 << 10); } + +// 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}; +}; + +// Locate path to file, given its filename or filepath suffix and possible dirs it might lie in +// Function will also walk back MAX_DEPTH dirs from CWD to check for such a file path +inline std::string locateFile(const std::string& filepathSuffix, const std::vector& directories) +{ + const int MAX_DEPTH{10}; + bool found{false}; + std::string filepath; + + for (auto& dir : directories) + { + filepath = dir + filepathSuffix; + + for (int i = 0; i < MAX_DEPTH && !found; i++) + { + std::ifstream checkFile(filepath); + found = checkFile.is_open(); + if (found) break; + filepath = "../" + filepath; // Try again in parent dir + } + + if (found) + { + break; + } + + filepath.clear(); + } + + if (filepath.empty()) { + std::string directoryList = std::accumulate(directories.begin() + 1, directories.end(), directories.front(), + [](const std::string& a, const std::string& b) { return a + "\n\t" + b; }); + throw std::runtime_error("Could not find " + filepathSuffix + " in data directories:\n\t" + directoryList); + } + return filepath; +} + +inline void readPGMFile(const std::string& fileName, uint8_t* buffer, int inH, int inW) +{ + std::ifstream infile(fileName, std::ifstream::binary); + assert(infile.is_open() && "Attempting to read from a file that is not open."); + std::string magic, h, w, max; + infile >> magic >> h >> w >> max; + infile.seekg(1, infile.cur); + infile.read(reinterpret_cast(buffer), inH * inW); +} + +namespace samples_common +{ + +inline void* safeCudaMalloc(size_t memSize) +{ + void* deviceMem; + CHECK(cudaMalloc(&deviceMem, memSize)); + if (deviceMem == nullptr) + { + std::cerr << "Out of memory" << std::endl; + exit(1); + } + return deviceMem; +} + +inline bool isDebug() +{ + return (std::getenv("TENSORRT_DEBUG") ? true : false); +} + +struct InferDeleter +{ + template + void operator()(T* obj) const + { + if (obj) { + obj->destroy(); + } + } +}; + +template +inline std::shared_ptr infer_object(T* obj) +{ + if (!obj) { + throw std::runtime_error("Failed to create object"); + } + return std::shared_ptr(obj, InferDeleter()); +} + +template +inline std::vector argsort(Iter begin, Iter end, bool reverse = false) +{ + std::vector inds(end - begin); + std::iota(inds.begin(), inds.end(), 0); + if (reverse) { + std::sort(inds.begin(), inds.end(), [&begin](size_t i1, size_t i2) { + return begin[i2] < begin[i1]; + }); + } + else + { + std::sort(inds.begin(), inds.end(), [&begin](size_t i1, size_t i2) { + return begin[i1] < begin[i2]; + }); + } + return inds; +} + +inline bool readReferenceFile(const std::string& fileName, std::vector& refVector) +{ + std::ifstream infile(fileName); + if (!infile.is_open()) { + cout << "ERROR: readReferenceFile: Attempting to read from a file that is not open." << endl; + return false; + } + std::string line; + while (std::getline(infile, line)) { + if (line.empty()) continue; + refVector.push_back(line); + } + infile.close(); + return true; +} + +template +inline std::vector classify(const vector& refVector, const result_vector_t& output, const size_t topK) +{ + auto inds = samples_common::argsort(output.cbegin(), output.cend(), true); + std::vector result; + for (size_t k = 0; k < topK; ++k) { + result.push_back(refVector[inds[k]]); + } + return result; +} + +//...LG returns top K indices, not values. +template +inline vector topK(const vector inp, const size_t k) +{ + vector result; + std::vector inds = samples_common::argsort(inp.cbegin(), inp.cend(), true); + result.assign(inds.begin(), inds.begin()+k); + return result; +} + +template +inline bool readASCIIFile(const string& fileName, const size_t size, vector& out) +{ + std::ifstream infile(fileName); + if (!infile.is_open()) { + cout << "ERROR readASCIIFile: Attempting to read from a file that is not open." << endl; + return false; + } + out.clear(); + out.reserve(size); + out.assign(std::istream_iterator(infile), std::istream_iterator()); + infile.close(); + return true; +} + +template +inline bool writeASCIIFile(const string& fileName, const vector& in) +{ + std::ofstream outfile(fileName); + if (!outfile.is_open()) { + cout << "ERROR: writeASCIIFile: Attempting to write to a file that is not open." << endl; + return false; + } + for (auto fn : in) { + outfile << fn << " "; + } + outfile.close(); + return true; +} + +inline void print_version() +{ +//... This can be only done after statically linking this support into parserONNX.library +#if 0 + std::cout << "Parser built against:" << std::endl; + std::cout << " ONNX IR version: " << nvonnxparser::onnx_ir_version_string(onnx::IR_VERSION) << std::endl; +#endif + std::cout << " TensorRT version: " + << NV_TENSORRT_MAJOR << "." + << NV_TENSORRT_MINOR << "." + << NV_TENSORRT_PATCH << "." + << NV_TENSORRT_BUILD << std::endl; +} + +inline string getFileType(const string& filepath) +{ + return filepath.substr(filepath.find_last_of(".") + 1); +} + +inline string toLower(const string& inp) +{ + string out = inp; + std::transform(out.begin(), out.end(), out.begin(), ::tolower); + return out; +} + +inline unsigned int getElementSize(nvinfer1::DataType t) +{ + switch (t) + { + case nvinfer1::DataType::kINT32: return 4; + case nvinfer1::DataType::kFLOAT: return 4; + case nvinfer1::DataType::kHALF: return 2; + case nvinfer1::DataType::kINT8: return 1; + } + throw std::runtime_error("Invalid DataType."); + return 0; +} + +inline int64_t volume(const nvinfer1::Dims& d) +{ + return std::accumulate(d.d, d.d + d.nbDims, 1, std::multiplies()); +} + +// Struct to maintain command-line arguments. +struct Args +{ + bool runInInt8 = false; +}; + +// Populates the Args struct with the provided command-line parameters. +inline void parseArgs(Args& args, int argc, char* argv[]) +{ + if (argc >= 1) + { + for (int i = 1; i < argc; ++i) + { + if (!strcmp(argv[i], "--int8")) args.runInInt8 = true; + } + } +} + +template +struct PPM +{ + std::string magic, fileName; + int h, w, max; + uint8_t buffer[C * H * W]; +}; + +struct BBox +{ + float x1, y1, x2, y2; +}; + +template +inline void writePPMFileWithBBox(const std::string& filename, PPM& ppm, const BBox& bbox) +{ + std::ofstream outfile("./" + filename, std::ofstream::binary); + assert(!outfile.fail()); + outfile << "P6" << "\n" << ppm.w << " " << ppm.h << "\n" << ppm.max << "\n"; + auto round = [](float x) -> int { return int(std::floor(x + 0.5f)); }; + const int x1 = std::min(std::max(0, round(int(bbox.x1))), W - 1); + const int x2 = std::min(std::max(0, round(int(bbox.x2))), W - 1); + const int y1 = std::min(std::max(0, round(int(bbox.y1))), H - 1); + const int y2 = std::min(std::max(0, round(int(bbox.y2))), H - 1); + for (int x = x1; x <= x2; ++x) + { + // bbox top border + ppm.buffer[(y1 * ppm.w + x) * 3] = 255; + ppm.buffer[(y1 * ppm.w + x) * 3 + 1] = 0; + ppm.buffer[(y1 * ppm.w + x) * 3 + 2] = 0; + // bbox bottom border + ppm.buffer[(y2 * ppm.w + x) * 3] = 255; + ppm.buffer[(y2 * ppm.w + x) * 3 + 1] = 0; + ppm.buffer[(y2 * ppm.w + x) * 3 + 2] = 0; + } + for (int y = y1; y <= y2; ++y) + { + // bbox left border + ppm.buffer[(y * ppm.w + x1) * 3] = 255; + ppm.buffer[(y * ppm.w + x1) * 3 + 1] = 0; + ppm.buffer[(y * ppm.w + x1) * 3 + 2] = 0; + // bbox right border + ppm.buffer[(y * ppm.w + x2) * 3] = 255; + ppm.buffer[(y * ppm.w + x2) * 3 + 1] = 0; + ppm.buffer[(y * ppm.w + x2) * 3 + 2] = 0; + } + outfile.write(reinterpret_cast(ppm.buffer), ppm.w * ppm.h * 3); +} + +} // namespace samples_common + +#endif // _TRT_COMMON_H_ diff --git a/yolov3-spp/gen_wts.py b/yolov3-spp/gen_wts.py new file mode 100644 index 0000000..7604c77 --- /dev/null +++ b/yolov3-spp/gen_wts.py @@ -0,0 +1,22 @@ +import struct +import sys +from models import * +from utils.utils import * + +model = Darknet('cfg/yolov3-spp.cfg', (416, 416)) +weights = sys.argv[1] +dev = '1' +device = torch_utils.select_device(dev) +model.load_state_dict(torch.load(weights, map_location=device)['model']) + + +f = open('yolov3-spp_ultralytics68.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-spp/leaky.cu b/yolov3-spp/leaky.cu new file mode 100644 index 0000000..d91aad4 --- /dev/null +++ b/yolov3-spp/leaky.cu @@ -0,0 +1,25 @@ +#include +#include +#include "leaky.cuh" + + +__global__ void _leakyReluKer(float const *in, float *out, int size) { + int index = threadIdx.x + blockIdx.x * blockDim.x; + if (index >= size) + return ; + + if (in[index] < 0) + out[index] = in[index] * 0.1; + else + out[index] = in[index]; +} + +extern "C" void culeaky(float const *in, float *out, int size) { + int block_size = 256; + int grid_size = (size + block_size - 1) / block_size; + _leakyReluKer<<>>(in, out, size); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + fprintf(stderr, "Failed to launch _leakyReluKer kernel (error code %s)!\n", cudaGetErrorString(err)); + } +} diff --git a/yolov3-spp/leaky.cuh b/yolov3-spp/leaky.cuh new file mode 100644 index 0000000..6d1c5f5 --- /dev/null +++ b/yolov3-spp/leaky.cuh @@ -0,0 +1,8 @@ +#ifndef LEAKY_H +#define LEAKY_H + +extern "C" + +void culeaky(float const *in, float *out, int size); + +#endif diff --git a/yolov3-spp/leakyplugin.cpp b/yolov3-spp/leakyplugin.cpp new file mode 100644 index 0000000..c1cd4d2 --- /dev/null +++ b/yolov3-spp/leakyplugin.cpp @@ -0,0 +1,62 @@ +#include "common.h" +#include "leaky.cuh" +#include "leakyplugin.h" + +using namespace nvinfer1; +using nvinfer1::LeakyPlugin; +using nvinfer1::PluginFactory; + +LeakyPlugin::LeakyPlugin() { +} + +LeakyPlugin::LeakyPlugin(const void* buffer, size_t size) { + assert(size == sizeof(input_size_)); + input_size_ = *reinterpret_cast(buffer); +} + +int LeakyPlugin::getNbOutputs() const { + return 1; +} + +Dims LeakyPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) { + assert(nbInputDims == 1); + assert(index == 0); + // Output dimensions + return DimsCHW(inputs[0].d[0], inputs[0].d[1], inputs[0].d[2]); +} + +void LeakyPlugin::configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) { + input_size_ = inputDims[0].d[0] * inputDims[0].d[1] * inputDims[0].d[2]; +} + +int LeakyPlugin::initialize() { + return 0; +} + +void LeakyPlugin::terminate() {} + +size_t LeakyPlugin::getWorkspaceSize(int maxBatchSize) const { + return 0; +} + +int LeakyPlugin::enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) { + culeaky(reinterpret_cast(inputs[0]), reinterpret_cast(outputs[0]), input_size_); + return 0; +} + +size_t LeakyPlugin::getSerializationSize() { + return sizeof(input_size_); +} + +void LeakyPlugin::serialize(void* buffer) { + *reinterpret_cast(buffer) = input_size_; +} + +IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialData, size_t serialLength) { + IPlugin *plugin = nullptr; + if (strstr(layerName, "leaky") != NULL) { + plugin = new LeakyPlugin(serialData, serialLength); + } + return plugin; +} + diff --git a/yolov3-spp/leakyplugin.h b/yolov3-spp/leakyplugin.h new file mode 100644 index 0000000..e6999c0 --- /dev/null +++ b/yolov3-spp/leakyplugin.h @@ -0,0 +1,32 @@ +#ifndef LEAKY_PLUGIN_H +#define LEAKY_PLUGIN_H +#include + +namespace nvinfer1 { +class LeakyPlugin : public IPlugin { + public: + LeakyPlugin(); + LeakyPlugin(const void* buffer, size_t size); + ~LeakyPlugin() override = default; + int getNbOutputs() const override; + Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override; + void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override; + int initialize() override; + void terminate() override; + size_t getWorkspaceSize(int maxBatchSize) const override; + int enqueue( + int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) override; + size_t getSerializationSize() override; + void serialize(void* buffer) override; + + private: + int input_size_; +}; + +class PluginFactory : public IPluginFactory { + public: + IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength) override; +}; + +} +#endif diff --git a/yolov3-spp/plugin_factory.cpp b/yolov3-spp/plugin_factory.cpp new file mode 100644 index 0000000..e8ceb4d --- /dev/null +++ b/yolov3-spp/plugin_factory.cpp @@ -0,0 +1,17 @@ +#include "common.h" +#include "plugin_factory.h" +#include "NvInferPlugin.h" +#include "yololayer.h" + +using namespace nvinfer1; +using nvinfer1::PluginFactory; + +IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialData, size_t serialLength) { + IPlugin *plugin = nullptr; + if (strstr(layerName, "leaky") != NULL) { + plugin = plugin::createPReLUPlugin(serialData, serialLength); + } else if (strstr(layerName, "yolo") != NULL) { + plugin = new YoloLayerPlugin(serialData, serialLength); + } + return plugin; +} diff --git a/yolov3-spp/plugin_factory.h b/yolov3-spp/plugin_factory.h new file mode 100644 index 0000000..0be0225 --- /dev/null +++ b/yolov3-spp/plugin_factory.h @@ -0,0 +1,12 @@ +#ifndef MY_PLUGIN_FACTORY_H +#define MY_PLUGIN_FACTORY_H +#include + +namespace nvinfer1 { +class PluginFactory : public IPluginFactory { + public: + IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength) override; +}; + +} +#endif diff --git a/yolov3-spp/samples/bus.jpg b/yolov3-spp/samples/bus.jpg new file mode 100644 index 0000000..b43e311 Binary files /dev/null and b/yolov3-spp/samples/bus.jpg differ diff --git a/yolov3-spp/samples/zidane.jpg b/yolov3-spp/samples/zidane.jpg new file mode 100644 index 0000000..92d72ea Binary files /dev/null and b/yolov3-spp/samples/zidane.jpg differ diff --git a/yolov3-spp/yololayer.cu b/yolov3-spp/yololayer.cu new file mode 100644 index 0000000..1c7c843 --- /dev/null +++ b/yolov3-spp/yololayer.cu @@ -0,0 +1,271 @@ +#include "YoloConfigs.h" +#include "yololayer.h" + +using namespace Yolo; + +namespace nvinfer1 +{ + YoloLayerPlugin::YoloLayerPlugin(const int cudaThread /*= 512*/):mThreadCount(cudaThread) + { + mClassCount = CLASS_NUM; + mYoloKernel.clear(); + mYoloKernel.push_back(yolo1); + mYoloKernel.push_back(yolo2); + mYoloKernel.push_back(yolo3); + + mKernelCount = mYoloKernel.size(); + } + + YoloLayerPlugin::~YoloLayerPlugin() + { + if(mInputBuffer) + CUDA_CHECK(cudaFreeHost(mInputBuffer)); + + if(mOutputBuffer) + CUDA_CHECK(cudaFreeHost(mOutputBuffer)); + } + + // 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) + { + 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() + { + return sizeof(mClassCount) + sizeof(mThreadCount) + sizeof(mKernelCount) + sizeof(Yolo::YoloKernel) * mYoloKernel.size(); + } + + int YoloLayerPlugin::initialize() + { + int totalCount = 0; + for(const auto& yolo : mYoloKernel) + totalCount += (LOCATIONS + 1) * yolo.width*yolo.height * CHECK_COUNT; + CUDA_CHECK(cudaHostAlloc(&mInputBuffer, totalCount * sizeof(float), cudaHostAllocDefault)); + + totalCount = 0;//detection count + for(const auto& yolo : mYoloKernel) + totalCount += yolo.width*yolo.height * CHECK_COUNT; + CUDA_CHECK(cudaHostAlloc(&mOutputBuffer, sizeof(float) + totalCount * sizeof(Detection), cudaHostAllocDefault)); + return 0; + } + + Dims YoloLayerPlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) + { + //output the result to channel + int totalCount = 0; + for(const auto& yolo : mYoloKernel) + totalCount += yolo.width*yolo.height * CHECK_COUNT * sizeof(Detection) / sizeof(float); + + return Dims3(totalCount + 1, 1, 1); + } + + /*void YoloLayerPlugin::forwardCpu(const float*const * inputs, float* outputs, cudaStream_t stream,int batchSize) + { + auto Logist = [=](float data){ + return 1./(1. + exp(-data)); + }; + + int totalOutputCount = 0; + int i = 0; + int totalCount = 0; + for(const auto& yolo : mYoloKernel) + { + totalOutputCount += yolo.width*yolo.height * CHECK_COUNT * sizeof(Detection) / sizeof(float); + totalCount += (LOCATIONS + 1 + mClassCount) * yolo.width*yolo.height * CHECK_COUNT; + ++ i; + } + + for (int idx = 0; idx < batchSize;idx++) + { + i = 0; + float* inputData = (float *)mInputBuffer;// + idx *totalCount; //if create more batch size + for(const auto& yolo : mYoloKernel) + { + int size = (LOCATIONS + 1 + mClassCount) * yolo.width*yolo.height * CHECK_COUNT; + CUDA_CHECK(cudaMemcpyAsync(inputData, (float *)inputs[i] + idx * size, size * sizeof(float), cudaMemcpyDeviceToHost, stream)); + inputData += size; + ++ i; + } + + CUDA_CHECK(cudaStreamSynchronize(stream)); + + inputData = (float *)mInputBuffer ;//+ idx *totalCount; //if create more batch size + std::vector result; + for (const auto& yolo : mYoloKernel) + { + int stride = yolo.width*yolo.height; + for (int j = 0;j < stride ;++j) + { + for (int k = 0;k < CHECK_COUNT; ++k ) + { + int beginIdx = (LOCATIONS + 1 + mClassCount)* stride *k + j; + int objIndex = beginIdx + LOCATIONS*stride; + + //check obj + float objProb = Logist(inputData[objIndex]); + if(objProb <= IGNORE_THRESH) + continue; + + //classes + int classId = -1; + float maxProb = IGNORE_THRESH; + for (int c = 0;c< mClassCount;++c){ + float cProb = Logist(inputData[beginIdx + (5 + c) * stride]) * objProb; + if(cProb > maxProb){ + maxProb = cProb; + classId = c; + } + } + + if(classId >= 0) { + Detection det; + int row = j / yolo.width; + int cols = j % yolo.width; + + //Location + det.bbox[0] = (cols + Logist(inputData[beginIdx]))/ yolo.width; + det.bbox[1] = (row + Logist(inputData[beginIdx+stride]))/ yolo.height; + det.bbox[2] = exp(inputData[beginIdx+2*stride]) * yolo.anchors[2*k]; + det.bbox[3] = exp(inputData[beginIdx+3*stride]) * yolo.anchors[2*k + 1]; + //det.classId = classId; + det.prob = maxProb; + + result.emplace_back(det); + } + } + } + + inputData += (LOCATIONS + 1 + mClassCount) * stride * CHECK_COUNT; + } + + + int detCount =result.size(); + auto data = (float *)mOutputBuffer;// + idx*(totalOutputCount + 1); //if create more batch size + float * begin = data; + //copy count; + data[0] = (float)detCount; + data++; + //copy result + memcpy(data,result.data(),result.size()*sizeof(Detection)); + + //(count + det result) + CUDA_CHECK(cudaMemcpyAsync(outputs, begin,sizeof(float) + result.size()*sizeof(Detection), cudaMemcpyHostToDevice, stream)); + + outputs += totalOutputCount + 1; + } + };*/ + + __device__ float Logist(float data){ return 1./(1. + exp(-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 info_len_i = 5 + classes; + //int info_len_o = 7; + int input_col = idx; + //int out_row = input_col; + + for (int k = 0; k < 3; ++k) { + float *res_count = output; + if(*res_count > 1000) break; + int count = (int)atomicAdd(res_count, 1); + char* data = (char * )res_count + sizeof(float) + count*sizeof(Detection); + Detection* det = (Detection*)(data); + + int class_id = 0; + float max_prob = 0.0; + for (int i = 5; i < info_len_i; ++i) { + float p = Logist(input[input_col + k * info_len_i * total_grid + i * total_grid]); + if (p > max_prob) { + max_prob = p; + class_id = i - 5; + } + } + + int row = idx / yoloWidth; + int col = idx % yoloWidth; + + //Location + det->bbox[0] = (col + Logist(input[input_col + k * info_len_i * total_grid + 0 * total_grid])) * INPUT_W / yoloWidth; + det->bbox[1] = (row + Logist(input[input_col + k * info_len_i * total_grid + 1 * total_grid])) * INPUT_H / yoloHeight; + det->bbox[2] = exp(input[input_col + k * info_len_i * total_grid + 2 * total_grid]) * anchors[2*k]; + det->bbox[3] = exp(input[input_col + k * info_len_i * total_grid + 3 * total_grid]) * anchors[2*k + 1]; + det->det_confidence = Logist(input[input_col + k * info_len_i * total_grid + 4 * total_grid]); + det->class_id = class_id; + det->class_confidence = max_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; + for (unsigned int i = 0;i< mYoloKernel.size();++i) + { + const auto& yolo = mYoloKernel[i]; + outputElem += yolo.width*yolo.height * CHECK_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 < 256) + 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); + + //CPU + //forwardCpu((const float *const *)inputs,(float *)outputs[0],stream,batchSize); + return 0; + }; + +} diff --git a/yolov3-spp/yololayer.h b/yolov3-spp/yololayer.h new file mode 100644 index 0000000..2d250c6 --- /dev/null +++ b/yolov3-spp/yololayer.h @@ -0,0 +1,80 @@ +#ifndef _YOLO_LAYER_H +#define _YOLO_LAYER_H + +#include +#include +#include +#include +#include +#include "NvInfer.h" +#include "Utils.h" +#include + +namespace Yolo +{ + struct YoloKernel; + + 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 IPluginExt + { + public: + explicit YoloLayerPlugin(const int cudaThread = 256); + YoloLayerPlugin(const void* data, size_t length); + + ~YoloLayerPlugin(); + + int getNbOutputs() const override + { + return 1; + } + + Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override; + + bool supportsFormat(DataType type, PluginFormat format) const override { + return type == DataType::kFLOAT && format == PluginFormat::kNCHW; + } + + void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, DataType type, PluginFormat format, int maxBatchSize) 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() override; + + virtual void serialize(void* buffer) override; + + void forwardGpu(const float *const * inputs,float * output, cudaStream_t stream,int batchSize = 1); + + void forwardCpu(const float *const * inputs,float * output, cudaStream_t stream,int batchSize = 1); + + private: + int mClassCount; + int mKernelCount; + std::vector mYoloKernel; + int mThreadCount; + //int mDetNum; + + //cpu + void* mInputBuffer {nullptr}; + void* mOutputBuffer {nullptr}; + }; +}; + +#endif diff --git a/yolov3-spp/yolov3-spp.cpp b/yolov3-spp/yolov3-spp.cpp new file mode 100644 index 0000000..6b58a88 --- /dev/null +++ b/yolov3-spp/yolov3-spp.cpp @@ -0,0 +1,598 @@ +#include "NvInfer.h" +#include "NvInferPlugin.h" +#include "cuda_runtime_api.h" +#include "common.h" +#include +#include +#include +#include +#include +#include +#include "plugin_factory.h" +#include "yololayer.h" +#include +#include + +#define USE_FP16 // comment out this if want to use FP32 + +// stuff we know about the network and the input/output blobs +static const int INPUT_H = 256; +static const int INPUT_W = 416; +static const int OUTPUT_SIZE = 1000 * 7 + 1; + +const char* INPUT_BLOB_NAME = "data"; +const char* OUTPUT_BLOB_NAME = "prob"; + +using namespace nvinfer1; +using namespace Yolo; + +static Logger gLogger; + + +cv::Mat preprocess_img(cv::Mat& img) { + int w, h, x, y; + float r_w = INPUT_W / (img.cols*1.0); + float r_h = INPUT_H / (img.rows*1.0); + if (r_h > r_w) { + w = INPUT_W; + h = r_w * img.rows; + x = 0; + y = (INPUT_H - h) / 2; + } else { + w = r_h* img.cols; + h = INPUT_H; + x = (INPUT_W - w) / 2; + y = 0; + } + cv::Mat re(h, w, CV_8UC3); + cv::resize(img, re, re.size(), 0, 0, cv::INTER_CUBIC); + cv::Mat out(INPUT_H, INPUT_W, CV_8UC3, cv::Scalar(128, 128, 128)); + re.copyTo(out(cv::Rect(x, y, re.cols, re.rows))); + return out; +} + +cv::Rect get_rect(cv::Mat& img, float bbox[4]) { + int l, r, t, b; + float r_w = INPUT_W / (img.cols * 1.0); + float r_h = INPUT_H / (img.rows * 1.0); + if (r_h > r_w) { + l = bbox[0] - bbox[2]/2.f; + r = bbox[0] + bbox[2]/2.f; + t = bbox[1] - bbox[3]/2.f - (INPUT_H - r_w * img.rows) / 2; + b = bbox[1] + bbox[3]/2.f - (INPUT_H - r_w * img.rows) / 2; + l = l / r_w; + r = r / r_w; + t = t / r_w; + b = b / r_w; + } else { + l = bbox[0] - bbox[2]/2.f - (INPUT_W - r_h * img.cols) / 2; + r = bbox[0] + bbox[2]/2.f - (INPUT_W - r_h * img.cols) / 2; + t = bbox[1] - bbox[3]/2.f; + b = bbox[1] + bbox[3]/2.f; + l = l / r_h; + r = r / r_h; + t = t / r_h; + b = b / r_h; + } + return cv::Rect(l, t, r-l, b-t); +} + +float iou(float lbox[4], float rbox[4]) { + float interBox[] = { + max(lbox[0] - lbox[2]/2.f , rbox[0] - rbox[2]/2.f), //left + min(lbox[0] + lbox[2]/2.f , rbox[0] + rbox[2]/2.f), //right + max(lbox[1] - lbox[3]/2.f , rbox[1] - rbox[3]/2.f), //top + min(lbox[1] + lbox[3]/2.f , rbox[1] + rbox[3]/2.f), //bottom + }; + + if(interBox[2] > interBox[3] || interBox[0] > interBox[1]) + return 0.0f; + + float interBoxS =(interBox[1]-interBox[0])*(interBox[3]-interBox[2]); + return interBoxS/(lbox[2]*lbox[3] + rbox[2]*rbox[3] -interBoxS); +} + +bool cmp(Detection& a, Detection& b) { + return a.det_confidence > b.det_confidence; +} + + +void nms(std::vector& res, float *output, float nms_thresh = 0.4) { + std::map> m; + for (int i = 0; i < output[0]; i++) { + if (output[1 + 7 * i + 4] <= 0.5) continue; + Detection det; + memcpy(&det, &output[1 + 7 * i], 7 * sizeof(float)); + if (m.count(det.class_id) == 0) m.emplace(det.class_id, std::vector()); + m[det.class_id].push_back(det); + } + for (auto it = m.begin(); it != m.end(); it++) { + //std::cout << it->second[0].class_id << " --- " << std::endl; + auto& dets = it->second; + std::sort(dets.begin(), dets.end(), cmp); + for (size_t m = 0; m < dets.size(); ++m) { + auto& item = dets[m]; + res.push_back(item); + for (size_t n = m + 1; n < dets.size(); ++n) { + if (iou(item.bbox, dets[n].bbox) > nms_thresh) { + dets.erase(dets.begin()+n); + --n; + } + } + } + } +} + +// Load weights from files shared with TensorRT samples. +// 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; + std::cout << "len " << len << std::endl; + + 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; +} + +ILayer* convBnLeaky(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int ksize, int s, int p, int linx) { + std::cout << linx << std::endl; + Weights emptywts{DataType::kFLOAT, nullptr, 0}; + IConvolutionLayer* conv1 = network->addConvolution(input, outch, DimsHW{ksize, ksize}, weightMap["module_list." + std::to_string(linx) + ".Conv2d.weight"], emptywts); + assert(conv1); + conv1->setStride(DimsHW{s, s}); + conv1->setPadding(DimsHW{p, p}); + + IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), "module_list." + std::to_string(linx) + ".BatchNorm2d", 1e-5); + + ITensor* inputTensors[] = {bn1->getOutput(0)}; + //LeakyPlugin *lr = new LeakyPlugin(); + auto lr = plugin::createPReLUPlugin(0.1); + auto lr1 = network->addPlugin(inputTensors, 1, *lr); + assert(lr1); + lr1->setName(("leaky" + std::to_string(linx)).c_str()); + return lr1; +} + +// Creat the engine using only the API and not any parser. +ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, DataType dt) +{ + INetworkDefinition* network = builder->createNetwork(); + + // Create input tensor of shape { 1, 1, 32, 32 } with name INPUT_BLOB_NAME + ITensor* data = network->addInput(INPUT_BLOB_NAME, dt, Dims3{3, INPUT_H, INPUT_W}); + assert(data); + + std::map weightMap = loadWeights("../yolov3-spp_ultralytics68.wts"); + Weights emptywts{DataType::kFLOAT, nullptr, 0}; + + // Yeah I am stupid, I just want to expand the complete arch of darknet.. + auto lr0 = convBnLeaky(network, weightMap, *data, 32, 3, 1, 1, 0); + auto lr1 = convBnLeaky(network, weightMap, *lr0->getOutput(0), 64, 3, 2, 1, 1); + auto lr2 = convBnLeaky(network, weightMap, *lr1->getOutput(0), 32, 1, 1, 0, 2); + auto lr3 = convBnLeaky(network, weightMap, *lr2->getOutput(0), 64, 3, 1, 1, 3); + auto ew4 = network->addElementWise(*lr3->getOutput(0), *lr1->getOutput(0), ElementWiseOperation::kSUM); + auto lr5 = convBnLeaky(network, weightMap, *ew4->getOutput(0), 128, 3, 2, 1, 5); + auto lr6 = convBnLeaky(network, weightMap, *lr5->getOutput(0), 64, 1, 1, 0, 6); + auto lr7 = convBnLeaky(network, weightMap, *lr6->getOutput(0), 128, 3, 1, 1, 7); + auto ew8 = network->addElementWise(*lr7->getOutput(0), *lr5->getOutput(0), ElementWiseOperation::kSUM); + auto lr9 = convBnLeaky(network, weightMap, *ew8->getOutput(0), 64, 1, 1, 0, 9); + auto lr10 = convBnLeaky(network, weightMap, *lr9->getOutput(0), 128, 3, 1, 1, 10); + auto ew11 = network->addElementWise(*lr10->getOutput(0), *ew8->getOutput(0), ElementWiseOperation::kSUM); + auto lr12 = convBnLeaky(network, weightMap, *ew11->getOutput(0), 256, 3, 2, 1, 12); + auto lr13 = convBnLeaky(network, weightMap, *lr12->getOutput(0), 128, 1, 1, 0, 13); + auto lr14 = convBnLeaky(network, weightMap, *lr13->getOutput(0), 256, 3, 1, 1, 14); + auto ew15 = network->addElementWise(*lr14->getOutput(0), *lr12->getOutput(0), ElementWiseOperation::kSUM); + auto lr16 = convBnLeaky(network, weightMap, *ew15->getOutput(0), 128, 1, 1, 0, 16); + auto lr17 = convBnLeaky(network, weightMap, *lr16->getOutput(0), 256, 3, 1, 1, 17); + auto ew18 = network->addElementWise(*lr17->getOutput(0), *ew15->getOutput(0), ElementWiseOperation::kSUM); + auto lr19 = convBnLeaky(network, weightMap, *ew18->getOutput(0), 128, 1, 1, 0, 19); + auto lr20 = convBnLeaky(network, weightMap, *lr19->getOutput(0), 256, 3, 1, 1, 20); + auto ew21 = network->addElementWise(*lr20->getOutput(0), *ew18->getOutput(0), ElementWiseOperation::kSUM); + auto lr22 = convBnLeaky(network, weightMap, *ew21->getOutput(0), 128, 1, 1, 0, 22); + auto lr23 = convBnLeaky(network, weightMap, *lr22->getOutput(0), 256, 3, 1, 1, 23); + auto ew24 = network->addElementWise(*lr23->getOutput(0), *ew21->getOutput(0), ElementWiseOperation::kSUM); + auto lr25 = convBnLeaky(network, weightMap, *ew24->getOutput(0), 128, 1, 1, 0, 25); + auto lr26 = convBnLeaky(network, weightMap, *lr25->getOutput(0), 256, 3, 1, 1, 26); + auto ew27 = network->addElementWise(*lr26->getOutput(0), *ew24->getOutput(0), ElementWiseOperation::kSUM); + auto lr28 = convBnLeaky(network, weightMap, *ew27->getOutput(0), 128, 1, 1, 0, 28); + auto lr29 = convBnLeaky(network, weightMap, *lr28->getOutput(0), 256, 3, 1, 1, 29); + auto ew30 = network->addElementWise(*lr29->getOutput(0), *ew27->getOutput(0), ElementWiseOperation::kSUM); + auto lr31 = convBnLeaky(network, weightMap, *ew30->getOutput(0), 128, 1, 1, 0, 31); + auto lr32 = convBnLeaky(network, weightMap, *lr31->getOutput(0), 256, 3, 1, 1, 32); + auto ew33 = network->addElementWise(*lr32->getOutput(0), *ew30->getOutput(0), ElementWiseOperation::kSUM); + auto lr34 = convBnLeaky(network, weightMap, *ew33->getOutput(0), 128, 1, 1, 0, 34); + auto lr35 = convBnLeaky(network, weightMap, *lr34->getOutput(0), 256, 3, 1, 1, 35); + auto ew36 = network->addElementWise(*lr35->getOutput(0), *ew33->getOutput(0), ElementWiseOperation::kSUM); + auto lr37 = convBnLeaky(network, weightMap, *ew36->getOutput(0), 512, 3, 2, 1, 37); + auto lr38 = convBnLeaky(network, weightMap, *lr37->getOutput(0), 256, 1, 1, 0, 38); + auto lr39 = convBnLeaky(network, weightMap, *lr38->getOutput(0), 512, 3, 1, 1, 39); + auto ew40 = network->addElementWise(*lr39->getOutput(0), *lr37->getOutput(0), ElementWiseOperation::kSUM); + auto lr41 = convBnLeaky(network, weightMap, *ew40->getOutput(0), 256, 1, 1, 0, 41); + auto lr42 = convBnLeaky(network, weightMap, *lr41->getOutput(0), 512, 3, 1, 1, 42); + auto ew43 = network->addElementWise(*lr42->getOutput(0), *ew40->getOutput(0), ElementWiseOperation::kSUM); + auto lr44 = convBnLeaky(network, weightMap, *ew43->getOutput(0), 256, 1, 1, 0, 44); + auto lr45 = convBnLeaky(network, weightMap, *lr44->getOutput(0), 512, 3, 1, 1, 45); + auto ew46 = network->addElementWise(*lr45->getOutput(0), *ew43->getOutput(0), ElementWiseOperation::kSUM); + auto lr47 = convBnLeaky(network, weightMap, *ew46->getOutput(0), 256, 1, 1, 0, 47); + auto lr48 = convBnLeaky(network, weightMap, *lr47->getOutput(0), 512, 3, 1, 1, 48); + auto ew49 = network->addElementWise(*lr48->getOutput(0), *ew46->getOutput(0), ElementWiseOperation::kSUM); + auto lr50 = convBnLeaky(network, weightMap, *ew49->getOutput(0), 256, 1, 1, 0, 50); + auto lr51 = convBnLeaky(network, weightMap, *lr50->getOutput(0), 512, 3, 1, 1, 51); + auto ew52 = network->addElementWise(*lr51->getOutput(0), *ew49->getOutput(0), ElementWiseOperation::kSUM); + auto lr53 = convBnLeaky(network, weightMap, *ew52->getOutput(0), 256, 1, 1, 0, 53); + auto lr54 = convBnLeaky(network, weightMap, *lr53->getOutput(0), 512, 3, 1, 1, 54); + auto ew55 = network->addElementWise(*lr54->getOutput(0), *ew52->getOutput(0), ElementWiseOperation::kSUM); + auto lr56 = convBnLeaky(network, weightMap, *ew55->getOutput(0), 256, 1, 1, 0, 56); + auto lr57 = convBnLeaky(network, weightMap, *lr56->getOutput(0), 512, 3, 1, 1, 57); + auto ew58 = network->addElementWise(*lr57->getOutput(0), *ew55->getOutput(0), ElementWiseOperation::kSUM); + auto lr59 = convBnLeaky(network, weightMap, *ew58->getOutput(0), 256, 1, 1, 0, 59); + auto lr60 = convBnLeaky(network, weightMap, *lr59->getOutput(0), 512, 3, 1, 1, 60); + auto ew61 = network->addElementWise(*lr60->getOutput(0), *ew58->getOutput(0), ElementWiseOperation::kSUM); + auto lr62 = convBnLeaky(network, weightMap, *ew61->getOutput(0), 1024, 3, 2, 1, 62); + auto lr63 = convBnLeaky(network, weightMap, *lr62->getOutput(0), 512, 1, 1, 0, 63); + auto lr64 = convBnLeaky(network, weightMap, *lr63->getOutput(0), 1024, 3, 1, 1, 64); + auto ew65 = network->addElementWise(*lr64->getOutput(0), *lr62->getOutput(0), ElementWiseOperation::kSUM); + auto lr66 = convBnLeaky(network, weightMap, *ew65->getOutput(0), 512, 1, 1, 0, 66); + auto lr67 = convBnLeaky(network, weightMap, *lr66->getOutput(0), 1024, 3, 1, 1, 67); + auto ew68 = network->addElementWise(*lr67->getOutput(0), *ew65->getOutput(0), ElementWiseOperation::kSUM); + auto lr69 = convBnLeaky(network, weightMap, *ew68->getOutput(0), 512, 1, 1, 0, 69); + auto lr70 = convBnLeaky(network, weightMap, *lr69->getOutput(0), 1024, 3, 1, 1, 70); + auto ew71 = network->addElementWise(*lr70->getOutput(0), *ew68->getOutput(0), ElementWiseOperation::kSUM); + auto lr72 = convBnLeaky(network, weightMap, *ew71->getOutput(0), 512, 1, 1, 0, 72); + auto lr73 = convBnLeaky(network, weightMap, *lr72->getOutput(0), 1024, 3, 1, 1, 73); + auto ew74 = network->addElementWise(*lr73->getOutput(0), *ew71->getOutput(0), ElementWiseOperation::kSUM); + auto lr75 = convBnLeaky(network, weightMap, *ew74->getOutput(0), 512, 1, 1, 0, 75); + auto lr76 = convBnLeaky(network, weightMap, *lr75->getOutput(0), 1024, 3, 1, 1, 76); + auto lr77 = convBnLeaky(network, weightMap, *lr76->getOutput(0), 512, 1, 1, 0, 77); + + auto pool78 = network->addPooling(*lr77->getOutput(0), PoolingType::kMAX, DimsHW{5,5}); + pool78->setPadding(DimsHW{2, 2}); + pool78->setStride(DimsHW{1, 1}); + auto pool80 = network->addPooling(*lr77->getOutput(0), PoolingType::kMAX, DimsHW{9,9}); + pool80->setPadding(DimsHW{4, 4}); + pool80->setStride(DimsHW{1, 1}); + auto pool82 = network->addPooling(*lr77->getOutput(0), PoolingType::kMAX, DimsHW{13,13}); + pool82->setPadding(DimsHW{6, 6}); + pool82->setStride(DimsHW{1, 1}); + + ITensor* inputTensors83[] = {pool82->getOutput(0), pool80->getOutput(0), pool78->getOutput(0), lr77->getOutput(0)}; + auto cat83 = network->addConcatenation(inputTensors83, 4); + + auto lr84 = convBnLeaky(network, weightMap, *cat83->getOutput(0), 512, 1, 1, 0, 84); + auto lr85 = convBnLeaky(network, weightMap, *lr84->getOutput(0), 1024, 3, 1, 1, 85); + auto lr86 = convBnLeaky(network, weightMap, *lr85->getOutput(0), 512, 1, 1, 0, 86); + auto lr87 = convBnLeaky(network, weightMap, *lr86->getOutput(0), 1024, 3, 1, 1, 87); + IConvolutionLayer* conv88 = network->addConvolution(*lr87->getOutput(0), 255, DimsHW{1, 1}, weightMap["module_list.88.Conv2d.weight"], weightMap["module_list.88.Conv2d.bias"]); + assert(conv88); + auto lr91 = convBnLeaky(network, weightMap, *lr86->getOutput(0), 256, 1, 1, 0, 91); + + float *deval = reinterpret_cast(malloc(sizeof(float) * 256 * 2 * 2)); + for (int i = 0; i < 256 * 2 * 2; i++) { + deval[i] = 1.0; + } + Weights deconvwts92{DataType::kFLOAT, deval, 256 * 2 * 2}; + IDeconvolutionLayer* deconv92 = network->addDeconvolution(*lr91->getOutput(0), 256, DimsHW{2, 2}, deconvwts92, emptywts); + assert(deconv92); + deconv92->setStride(DimsHW{2, 2}); + deconv92->setNbGroups(256); + weightMap["deconv92"] = deconvwts92; + + ITensor* inputTensors[] = {deconv92->getOutput(0), ew61->getOutput(0)}; + auto cat93 = network->addConcatenation(inputTensors, 2); + auto lr94 = convBnLeaky(network, weightMap, *cat93->getOutput(0), 256, 1, 1, 0, 94); + auto lr95 = convBnLeaky(network, weightMap, *lr94->getOutput(0), 512, 3, 1, 1, 95); + auto lr96 = convBnLeaky(network, weightMap, *lr95->getOutput(0), 256, 1, 1, 0, 96); + auto lr97 = convBnLeaky(network, weightMap, *lr96->getOutput(0), 512, 3, 1, 1, 97); + auto lr98 = convBnLeaky(network, weightMap, *lr97->getOutput(0), 256, 1, 1, 0, 98); + auto lr99 = convBnLeaky(network, weightMap, *lr98->getOutput(0), 512, 3, 1, 1, 99); + IConvolutionLayer* conv100 = network->addConvolution(*lr99->getOutput(0), 255, DimsHW{1, 1}, weightMap["module_list.100.Conv2d.weight"], weightMap["module_list.100.Conv2d.bias"]); + assert(conv100); + auto lr103 = convBnLeaky(network, weightMap, *lr98->getOutput(0), 128, 1, 1, 0, 103); + Weights deconvwts104{DataType::kFLOAT, deval, 128 * 2 * 2}; + IDeconvolutionLayer* deconv104 = network->addDeconvolution(*lr103->getOutput(0), 128, DimsHW{2, 2}, deconvwts104, emptywts); + assert(deconv104); + deconv104->setStride(DimsHW{2, 2}); + deconv104->setNbGroups(128); + ITensor* inputTensors1[] = {deconv104->getOutput(0), ew36->getOutput(0)}; + auto cat105 = network->addConcatenation(inputTensors1, 2); + auto lr106 = convBnLeaky(network, weightMap, *cat105->getOutput(0), 128, 1, 1, 0, 106); + auto lr107 = convBnLeaky(network, weightMap, *lr106->getOutput(0), 256, 3, 1, 1, 107); + auto lr108 = convBnLeaky(network, weightMap, *lr107->getOutput(0), 128, 1, 1, 0, 108); + auto lr109 = convBnLeaky(network, weightMap, *lr108->getOutput(0), 256, 3, 1, 1, 109); + auto lr110 = convBnLeaky(network, weightMap, *lr109->getOutput(0), 128, 1, 1, 0, 110); + auto lr111 = convBnLeaky(network, weightMap, *lr110->getOutput(0), 256, 3, 1, 1, 111); + IConvolutionLayer* conv112 = network->addConvolution(*lr111->getOutput(0), 255, DimsHW{1, 1}, weightMap["module_list.112.Conv2d.weight"], weightMap["module_list.112.Conv2d.bias"]); + assert(conv112); + auto yolo = new YoloLayerPlugin(); + ITensor* inputTensors_yolo[] = {conv88->getOutput(0), conv100->getOutput(0), conv112->getOutput(0)}; + auto yolo113 = network->addPlugin(inputTensors_yolo, 3, *yolo); + assert(yolo113); + yolo113->setName("yolo113"); + + yolo113->getOutput(0)->setName(OUTPUT_BLOB_NAME); + std::cout << "set name out" << std::endl; + network->markOutput(*yolo113->getOutput(0)); + + // Build engine + builder->setMaxBatchSize(maxBatchSize); + builder->setMaxWorkspaceSize(1 << 20); +#ifdef USE_FP16 + builder->setFp16Mode(true); +#endif + ICudaEngine* engine = builder->buildCudaEngine(*network); + std::cout << "build out" << 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); + + // Create model to populate the network, then set the outputs and create an engine + ICudaEngine* engine = createEngine(maxBatchSize, builder, DataType::kFLOAT); + assert(engine != nullptr); + + // Serialize the engine + (*modelStream) = engine->serialize(); + + // Close everything down + engine->destroy(); + builder->destroy(); +} + +void doInference(IExecutionContext& context, float* input, float* output, int batchSize) +{ + const ICudaEngine& engine = context.getEngine(); + + // Pointers to input and output device buffers to pass to engine. + // Engine requires exactly IEngine::getNbBindings() number of buffers. + assert(engine.getNbBindings() == 2); + void* buffers[2]; + + // In order to bind the buffers, we need to know the names of the input and output tensors. + // Note that indices are guaranteed to be less than IEngine::getNbBindings() + const int inputIndex = engine.getBindingIndex(INPUT_BLOB_NAME); + const int outputIndex = engine.getBindingIndex(OUTPUT_BLOB_NAME); + + // Create GPU buffers on device + CHECK(cudaMalloc(&buffers[inputIndex], batchSize * 3 * INPUT_H * INPUT_W * sizeof(float))); + CHECK(cudaMalloc(&buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float))); + + // Create stream + cudaStream_t stream; + CHECK(cudaStreamCreate(&stream)); + + // DMA input batch data to device, infer on the batch asynchronously, and DMA output back to host + CHECK(cudaMemcpyAsync(buffers[inputIndex], input, batchSize * 3 * INPUT_H * INPUT_W * sizeof(float), cudaMemcpyHostToDevice, stream)); + context.enqueue(batchSize, buffers, stream, nullptr); + CHECK(cudaMemcpyAsync(output, buffers[outputIndex], batchSize * OUTPUT_SIZE * sizeof(float), cudaMemcpyDeviceToHost, stream)); + cudaStreamSynchronize(stream); + + // Release stream and buffers + cudaStreamDestroy(stream); + CHECK(cudaFree(buffers[inputIndex])); + CHECK(cudaFree(buffers[outputIndex])); +} + +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; +} + + +int main(int argc, char** argv) +{ + std::cout << "beginning" << std::endl; + if (argc != 3) { + std::cerr << "arguments not right!" << std::endl; + std::cerr << "./yolov3-spp -s ./img // serialize model to plan file" << std::endl; + std::cerr << "./yolov3-spp -d ./img // deserialize plan file and run inference" << std::endl; + return -1; + } + + // create a model using the API directly and serialize it to a stream + char *trtModelStream{nullptr}; + size_t size{0}; + + if (std::string(argv[1]) == "-s") { + IHostMemory* modelStream{nullptr}; + APIToModel(1, &modelStream); + assert(modelStream != nullptr); + + std::ofstream p("yolov3-spp.engine"); + if (!p) + { + std::cerr << "could not open plan output file" << std::endl; + return -1; + } + p.write(reinterpret_cast(modelStream->data()), modelStream->size()); + modelStream->destroy(); + return 1; + } else if (std::string(argv[1]) == "-d") { + std::ifstream file("yolov3-spp.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 { + return -1; + } + + std::vector file_names; + if (read_files_in_dir(argv[2], file_names) < 0) { + std::cout << "read_files_in_dir failed." << std::endl; + return -1; + } + + // prepare input data --------------------------- + float data[3 * INPUT_H * INPUT_W]; + //for (int i = 0; i < 3 * INPUT_H * INPUT_W; i++) + // data[i] = 1.0; + static float prob[OUTPUT_SIZE]; + PluginFactory pf; + IRuntime* runtime = createInferRuntime(gLogger); + assert(runtime != nullptr); + ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, &pf); + assert(engine != nullptr); + IExecutionContext* context = engine->createExecutionContext(); + assert(context != nullptr); + + int fcount = 0; + for (auto f: file_names) { + fcount++; + std::cout << fcount << " " << f << std::endl; + //if (fcount < 5700) continue; + cv::Mat img = cv::imread(std::string(argv[2]) + "/" + f); + if (img.empty()) continue; + cv::Mat pr_img = preprocess_img(img); + //cv::imwrite("123.jpg", pr_img); + for (int i = 0; i < INPUT_H * INPUT_W; i++) { + data[i] = pr_img.at(i)[2] / 255.0; + data[i + INPUT_H * INPUT_W] = pr_img.at(i)[1] / 255.0; + data[i + 2 * INPUT_H * INPUT_W] = pr_img.at(i)[0] / 255.0; + } + + // Run inference + auto start = std::chrono::system_clock::now(); + doInference(*context, data, prob, 1); + std::vector res; + nms(res, prob); + for (int i=0; i<20; i++) { + std::cout << prob[i] << ","; + } + std::cout << res.size() << std::endl; + for (size_t j = 0; j < res.size(); j++) { + float *p = (float*)&res[j]; + for (size_t k = 0; k < 7; k++) { + std::cout << p[k] << ", "; + } + std::cout << std::endl; + cv::Rect r = get_rect(img, res[j].bbox); + cv::rectangle(img, r, cv::Scalar(0x27, 0xC1, 0x36), 2); + cv::putText(img, std::to_string((int)res[j].class_id), cv::Point(r.x, r.y - 1), cv::FONT_HERSHEY_PLAIN, 1.2, cv::Scalar(0xFF, 0xFF, 0xFF), 2); + } + auto end = std::chrono::system_clock::now(); + std::cout << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; + cv::imwrite("_" + f, img); + } + + // Destroy the engine + context->destroy(); + engine->destroy(); + runtime->destroy(); + + // Print histogram of the output distribution + //std::cout << "\nOutput:\n\n"; + //for (unsigned int i = 0; i < OUTPUT_SIZE/7; i++) + //{ + // if (prob[7 * i + 4] <= 0.5) continue; + // for (int j = 0; j < 7; j++) { + // std::cout << prob[7 * i + j] << ", "; + // } + // std::cout << std::endl; + + // //std::cout << prob[i] << ", "; + // //if (i % 10 == 0) std::cout << i / 10 << std::endl; + //} + //std::cout << std::endl; + + return 0; +}