haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add yolov3-spp tensorrt implement

Browse Source
This commit is contained in:
chufei1995 2020-04-02 18:55:38 +08:00
parent a68ec8233c
commit a1883a82ad
18 changed files with 1702 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

BIN
yolov3-spp/.yolov3-spp-1cls.cpp.swp Normal file
View File

Binary file not shown.

34
yolov3-spp/CMakeLists.txt Normal file
View File

@ -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)

36
yolov3-spp/README.md Normal file
View File

@ -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
```

94
yolov3-spp/Utils.h Normal file
View File

@ -0,0 +1,94 @@
#ifndef __TRT_UTILS_H_
#define __TRT_UTILS_H_
#include <iostream>
#include <vector>
#include <algorithm>
#include <cudnn.h>
#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<std::string, float> Record;
std::vector<Record> 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<typename T>
void write(char*& buffer, const T& val)
{
*reinterpret_cast<T*>(buffer) = val;
buffer += sizeof(T);
}
template<typename T>
void read(const char*& buffer, T& val)
{
val = *reinterpret_cast<const T*>(buffer);
buffer += sizeof(T);
}
}
#endif

55
yolov3-spp/YoloConfigs.h Normal file
View File

@ -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

356
yolov3-spp/common.h Normal file
View File

@ -0,0 +1,356 @@
#ifndef _TRT_COMMON_H_
#define _TRT_COMMON_H_
#include "NvInfer.h"
#include "NvOnnxConfig.h"
#include "NvOnnxParser.h"
#include <cuda_runtime_api.h>
#include <algorithm>
#include <cassert>
#include <fstream>
#include <iostream>
#include <iterator>
#include <map>
#include <memory>
#include <numeric>
#include <string>
#include <vector>
#include <cstring>
#include <cmath>
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<std::string>& 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<char*>(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 <typename T>
void operator()(T* obj) const
{
if (obj) {
obj->destroy();
}
}
};
template <typename T>
inline std::shared_ptr<T> infer_object(T* obj)
{
if (!obj) {
throw std::runtime_error("Failed to create object");
}
return std::shared_ptr<T>(obj, InferDeleter());
}
template <class Iter>
inline std::vector<size_t> argsort(Iter begin, Iter end, bool reverse = false)
{
std::vector<size_t> 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<std::string>& 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 <typename result_vector_t>
inline std::vector<std::string> classify(const vector<string>& refVector, const result_vector_t& output, const size_t topK)
{
auto inds = samples_common::argsort(output.cbegin(), output.cend(), true);
std::vector<std::string> 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 <typename T>
inline vector<size_t> topK(const vector<T> inp, const size_t k)
{
vector<size_t> result;
std::vector<size_t> inds = samples_common::argsort(inp.cbegin(), inp.cend(), true);
result.assign(inds.begin(), inds.begin()+k);
return result;
}
template <typename T>
inline bool readASCIIFile(const string& fileName, const size_t size, vector<T>& 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<T>(infile), std::istream_iterator<T>());
infile.close();
return true;
}
template <typename T>
inline bool writeASCIIFile(const string& fileName, const vector<T>& 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<int64_t>());
}
// 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 <int C, int H, int W>
struct PPM
{
std::string magic, fileName;
int h, w, max;
uint8_t buffer[C * H * W];
};
struct BBox
{
float x1, y1, x2, y2;
};
template <int C, int H, int W>
inline void writePPMFileWithBBox(const std::string& filename, PPM<C, H, W>& 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<char*>(ppm.buffer), ppm.w * ppm.h * 3);
}
} // namespace samples_common
#endif // _TRT_COMMON_H_

22
yolov3-spp/gen_wts.py Normal file
View File

@ -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')

25
yolov3-spp/leaky.cu Normal file
View File

@ -0,0 +1,25 @@
#include <cuda_runtime.h>
#include <stdio.h>
#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<<<grid_size, block_size>>>(in, out, size);
cudaError_t err = cudaGetLastError();
if (err != cudaSuccess) {
fprintf(stderr, "Failed to launch _leakyReluKer kernel (error code %s)!\n", cudaGetErrorString(err));
}
}

8
yolov3-spp/leaky.cuh Normal file
View File

@ -0,0 +1,8 @@
#ifndef LEAKY_H
#define LEAKY_H
extern "C"
void culeaky(float const *in, float *out, int size);
#endif

62
yolov3-spp/leakyplugin.cpp Normal file
View File

@ -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<const int*>(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<float const*>(inputs[0]), reinterpret_cast<float*>(outputs[0]), input_size_);
return 0;
}
size_t LeakyPlugin::getSerializationSize() {
return sizeof(input_size_);
}
void LeakyPlugin::serialize(void* buffer) {
*reinterpret_cast<int*>(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;
}

32
yolov3-spp/leakyplugin.h Normal file
View File

@ -0,0 +1,32 @@
#ifndef LEAKY_PLUGIN_H
#define LEAKY_PLUGIN_H
#include <NvInfer.h>
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

17
yolov3-spp/plugin_factory.cpp Normal file
View File

@ -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;
}

12
yolov3-spp/plugin_factory.h Normal file
View File

@ -0,0 +1,12 @@
#ifndef MY_PLUGIN_FACTORY_H
#define MY_PLUGIN_FACTORY_H
#include <NvInfer.h>
namespace nvinfer1 {
class PluginFactory : public IPluginFactory {
public:
IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength) override;
};
}
#endif

BIN
yolov3-spp/samples/bus.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 476 KiB

BIN
yolov3-spp/samples/zidane.jpg Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 165 KiB

271
yolov3-spp/yololayer.cu Normal file
View File

@ -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<const char *>(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<char*>(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 <Detection> 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;
};
}

80
yolov3-spp/yololayer.h Normal file
View File

@ -0,0 +1,80 @@
#ifndef _YOLO_LAYER_H
#define _YOLO_LAYER_H
#include <assert.h>
#include <cmath>
#include <string.h>
#include <cudnn.h>
#include <cublas_v2.h>
#include "NvInfer.h"
#include "Utils.h"
#include <iostream>
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<Yolo::YoloKernel> mYoloKernel;
int mThreadCount;
//int mDetNum;
//cpu
void* mInputBuffer {nullptr};
void* mOutputBuffer {nullptr};
};
};
#endif

598
yolov3-spp/yolov3-spp.cpp Normal file
View File

@ -0,0 +1,598 @@
#include "NvInfer.h"
#include "NvInferPlugin.h"
#include "cuda_runtime_api.h"
#include "common.h"
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
#include <chrono>
#include "plugin_factory.h"
#include "yololayer.h"
#include <opencv2/opencv.hpp>
#include <dirent.h>
#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<Detection>& res, float *output, float nms_thresh = 0.4) {
std::map<float, std::vector<Detection>> 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<Detection>());
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] <data x size in hex>
std::map<std::string, Weights> loadWeights(const std::string file)
{
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> 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<uint32_t*>(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<std::string, Weights>& 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<float*>(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<float*>(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<float*>(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<std::string, Weights>& 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<std::string, Weights> 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<float*>(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<std::string> &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<const char*>(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<std::string> 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<cv::Vec3b>(i)[2] / 255.0;
data[i + INPUT_H * INPUT_W] = pr_img.at<cv::Vec3b>(i)[1] / 255.0;
data[i + 2 * INPUT_H * INPUT_W] = pr_img.at<cv::Vec3b>(i)[0] / 255.0;
}
// Run inference
auto start = std::chrono::system_clock::now();
doInference(*context, data, prob, 1);
std::vector<Detection> 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<std::chrono::milliseconds>(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;
}