diff --git a/retinaface/CMakeLists.txt b/retinaface/CMakeLists.txt index 685c804..7696bad 100644 --- a/retinaface/CMakeLists.txt +++ b/retinaface/CMakeLists.txt @@ -18,16 +18,15 @@ link_directories(/usr/local/cuda-9.0/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 ${PROJECT_SOURCE_DIR}/yololayer.cu) +cuda_add_library(decodeplugin SHARED ${PROJECT_SOURCE_DIR}/decode.cu) find_package(OpenCV) include_directories(OpenCV_INCLUDE_DIRS) -add_executable(retina_50 ${PROJECT_SOURCE_DIR}/retina_r50.cpp) +add_executable(retina_50 ${PROJECT_SOURCE_DIR}/plugin_factory.cpp ${PROJECT_SOURCE_DIR}/retina_r50.cpp) target_link_libraries(retina_50 nvinfer nvinfer_plugin) target_link_libraries(retina_50 cudart) -#target_link_libraries(retina yololayer) +target_link_libraries(retina_50 decodeplugin) target_link_libraries(retina_50 ${OpenCV_LIBRARIES}) add_definitions(-O2 -pthread) diff --git a/retinaface/decode.cu b/retinaface/decode.cu new file mode 100644 index 0000000..65a6688 --- /dev/null +++ b/retinaface/decode.cu @@ -0,0 +1,125 @@ +#include "decode.h" +#include "stdio.h" + +namespace nvinfer1 +{ + DecodePlugin::DecodePlugin(const int cudaThread):thread_count_(cudaThread) + { + } + + DecodePlugin::~DecodePlugin() + { + } + + // create the plugin at runtime from a byte stream + DecodePlugin::DecodePlugin(const void* data, size_t length) + { + } + + void DecodePlugin::serialize(void* buffer) + { + } + + size_t DecodePlugin::getSerializationSize() + { + return 0; + } + + int DecodePlugin::initialize() + { + return 0; + } + + Dims DecodePlugin::getOutputDimensions(int index, const Dims* inputs, int nbInputDims) + { + //output the result to channel + int totalCount = 0; + totalCount += input_h_ / 8 * input_w_ / 8 * 2 * sizeof(decodeplugin::Detection) / sizeof(float); + totalCount += input_h_ / 16 * input_w_ / 16 * 2 * sizeof(decodeplugin::Detection) / sizeof(float); + totalCount += input_h_ / 32 * input_w_ / 32 * 2 * sizeof(decodeplugin::Detection) / sizeof(float); + + return Dims3(totalCount + 1, 1, 1); + } + + __device__ float Logist(float data){ return 1./(1. + exp(-data)); }; + + __global__ void CalDetection(const float *input, float *output, int num_elem, int input_h, int input_w, int step, int anchor) { + + int idx = threadIdx.x + blockDim.x * blockIdx.x; + if (idx >= num_elem) return; + + int h = input_h / step; + int w = input_w / step; + int y = idx / w; + int x = idx % w; + const float *bbox_reg = &input[0]; + const float *cls_reg = &input[2 * 4 * num_elem]; + const float *lmk_reg = &input[2 * 4 * num_elem + 2 * 2 * num_elem]; + + for (int k = 0; k < 2; ++k) { + float conf1 = cls_reg[idx + k * num_elem * 2]; + float conf2 = cls_reg[idx + k * num_elem * 2 + num_elem]; + conf2 = exp(conf2) / (exp(conf1) + exp(conf2)); + if (conf2 <= 0.002) continue; + + float *res_count = output; + int count = (int)atomicAdd(res_count, 1); + char* data = (char *)res_count + sizeof(float) + count * sizeof(decodeplugin::Detection); + decodeplugin::Detection* det = (decodeplugin::Detection*)(data); + + float prior[4]; + prior[0] = ((float)x + 0.5) / w; + prior[1] = ((float)y + 0.5) / h; + prior[2] = (float)anchor / input_w; + prior[3] = (float)anchor / input_h; + printf("prior0, %f\n", prior[0]); + printf("bbox0, %f\n", bbox_reg[idx + k * num_elem * 4]); + + //Location + det->bbox[0] = prior[0] + bbox_reg[idx + k * num_elem * 4] * 0.1 * prior[2]; + det->bbox[1] = prior[1] + bbox_reg[idx + k * num_elem * 4 + num_elem] * 0.1 * prior[3]; + det->bbox[2] = prior[2] * exp(bbox_reg[idx + k * num_elem * 4 + num_elem * 2] * 0.2); + det->bbox[3] = prior[3] * exp(bbox_reg[idx + k * num_elem * 4 + num_elem * 3] * 0.2); + det->bbox[0] -= det->bbox[2] / 2; + det->bbox[1] -= det->bbox[3] / 2; + det->bbox[2] += det->bbox[0]; + det->bbox[3] += det->bbox[1]; + det->bbox[0] *= input_w; + det->bbox[1] *= input_h; + det->bbox[2] *= input_w; + det->bbox[3] *= input_h; + det->class_confidence = conf2; + anchor *= 2; + } + } + + void DecodePlugin::forwardGpu(const float *const * inputs, float * output, cudaStream_t stream, int batchSize) + { + int num_elem = 0; + int base_step = 8; + int base_anchor = 16; + int thread_count; + for (unsigned int i = 0; i < 3; ++i) + { + num_elem = input_h_ / base_step * input_w_ / base_step; + thread_count = (num_elem < thread_count_) ? num_elem : thread_count_; + CalDetection<<< (num_elem + thread_count - 1) / thread_count, thread_count>>> + (inputs[i], output, num_elem, input_h_, input_w_, base_step, base_anchor); + base_step *= 2; + base_anchor *= 4; + } + + } + + + int DecodePlugin::enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) + { + //assert(batchSize == 1); + //GPU + //CUDA_CHECK(cudaStreamSynchronize(stream)); + forwardGpu((const float *const *)inputs,(float *)outputs[0],stream,batchSize); + + return 0; + }; + +} diff --git a/retinaface/decode.h b/retinaface/decode.h new file mode 100644 index 0000000..8be89d7 --- /dev/null +++ b/retinaface/decode.h @@ -0,0 +1,61 @@ +#ifndef _DECODE_CU_H +#define _DECODE_CU_H + +#include "NvInfer.h" + +namespace decodeplugin +{ + struct alignas(float) Detection{ + //x y w h + float bbox[4]; + float class_confidence; + float landmark[10]; + }; +} + + +namespace nvinfer1 +{ + class DecodePlugin: public IPluginExt + { + public: + explicit DecodePlugin(const int cudaThread = 256); + DecodePlugin(const void* data, size_t length); + + ~DecodePlugin(); + + 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); + + private: + const int input_h_ = 384; + const int input_w_ = 640; + int thread_count_ = 256; + }; +}; + +#endif diff --git a/retinaface/plugin_factory.cpp b/retinaface/plugin_factory.cpp new file mode 100644 index 0000000..0ee35a7 --- /dev/null +++ b/retinaface/plugin_factory.cpp @@ -0,0 +1,17 @@ +#include "plugin_factory.h" +#include "NvInferPlugin.h" +#include "decode.h" +#include "common.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, "decode") != NULL) { + plugin = new DecodePlugin(serialData, serialLength); + } + return plugin; +} diff --git a/retinaface/plugin_factory.h b/retinaface/plugin_factory.h new file mode 100644 index 0000000..0be0225 --- /dev/null +++ b/retinaface/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/retinaface/retina_r50.cpp b/retinaface/retina_r50.cpp index 3651133..715cbfb 100644 --- a/retinaface/retina_r50.cpp +++ b/retinaface/retina_r50.cpp @@ -8,8 +8,8 @@ #include #include #include -//#include "plugin_factory.h" -//#include "yololayer.h" +#include "plugin_factory.h" +#include "decode.h" #include //#define USE_FP16 // comment out this if want to use FP32 @@ -18,20 +18,13 @@ // stuff we know about the network and the input/output blobs static const int INPUT_H = 384; // static const int INPUT_W = 640; -static const int OUTPUT_SIZE = 256 * 48 * 80; +static const int OUTPUT_SIZE = 10080 * 15 + 1; const char* INPUT_BLOB_NAME = "data"; const char* OUTPUT_BLOB_NAME = "prob"; using namespace nvinfer1; static Logger gLogger; -typedef struct { - float bbox[4]; - float det_confidence; - float class_id; - float class_confidence; -} Detection; - cv::Mat preprocess_img(cv::Mat& img, int input_dim) { int w, h, x, y; if (img.cols > img.rows) { @@ -91,34 +84,34 @@ float iou(float lbox[4], float rbox[4]) { return interBoxS/(lbox[2]*lbox[3] + rbox[2]*rbox[3] -interBoxS); } -bool cmp(Detection& a, Detection& b) { - return a.det_confidence > b.det_confidence; +bool cmp(decodeplugin::Detection& a, decodeplugin::Detection& b) { + return a.class_confidence > b.class_confidence; } -void nms(std::vector& res, float *output, float nms_thresh = 0.4) { - std::map> m; - for (int i = 0; i < OUTPUT_SIZE / 7; i++) { - if (output[7 * i + 4] <= 0.5) continue; - Detection det; - memcpy(&det, &output[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; - } - } - } - } +void nms(std::vector& res, float *output, float nms_thresh = 0.4) { + //std::map> m; + //for (int i = 0; i < OUTPUT_SIZE / 7; i++) { + // if (output[7 * i + 4] <= 0.5) continue; + // decodeplugin::Detection det; + // memcpy(&det, &output[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 @@ -162,6 +155,14 @@ std::map loadWeights(const std::string file) { return weightMap; } +Weights getWeights(std::map& weightMap, std::string key) { + if (weightMap.count(key) != 1) { + std::cerr << key << " not existed in weight map, fatal error!!!" << std::endl; + exit(-1); + } + return weightMap[key]; +} + 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; @@ -241,7 +242,7 @@ IActivationLayer* bottleneck(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, int kernelsize, int stride, int padding, bool userelu, std::string lname) { Weights emptywts{DataType::kFLOAT, nullptr, 0}; - IConvolutionLayer* conv1 = network->addConvolution(input, outch, DimsHW{kernelsize, kernelsize}, weightMap[lname + ".0.weight"], emptywts); + IConvolutionLayer* conv1 = network->addConvolution(input, outch, DimsHW{kernelsize, kernelsize}, getWeights(weightMap, lname + ".0.weight"), emptywts); assert(conv1); conv1->setStride(DimsHW{stride, stride}); conv1->setPadding(DimsHW{padding, padding}); @@ -256,6 +257,19 @@ ILayer* conv_bn_relu(INetworkDefinition *network, std::map return relu1; } +IActivationLayer* ssh(INetworkDefinition *network, std::map& weightMap, ITensor& input, std::string lname) { + auto conv3x3 = conv_bn_relu(network, weightMap, input, 256 / 2, 3, 1, 1, false, lname + ".conv3X3"); + auto conv5x5_1 = conv_bn_relu(network, weightMap, input, 256 / 4, 3, 1, 1, true, lname + ".conv5X5_1"); + auto conv5x5 = conv_bn_relu(network, weightMap, *conv5x5_1->getOutput(0), 256 / 4, 3, 1, 1, false, lname + ".conv5X5_2"); + auto conv7x7 = conv_bn_relu(network, weightMap, *conv5x5_1->getOutput(0), 256 / 4, 3, 1, 1, true, lname + ".conv7X7_2"); + conv7x7 = conv_bn_relu(network, weightMap, *conv7x7->getOutput(0), 256 / 4, 3, 1, 1, false, lname + ".conv7x7_3"); + ITensor* inputTensors[] = {conv3x3->getOutput(0), conv5x5->getOutput(0), conv7x7->getOutput(0)}; + auto cat = network->addConcatenation(inputTensors, 3); + IActivationLayer* relu1 = network->addActivation(*cat->getOutput(0), ActivationType::kRELU); + assert(relu1); + return relu1; +} + // Creat the engine using only the API and not any parser. ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, DataType dt) { INetworkDefinition* network = builder->createNetwork(); @@ -335,9 +349,39 @@ ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, DataType output1 = conv_bn_relu(network, weightMap, *output1->getOutput(0), 256, 3, 1, 1, true, "fpn.merge1"); // ------------- SSH --------------- - output1->getOutput(0)->setName(OUTPUT_BLOB_NAME); + auto ssh1 = ssh(network, weightMap, *output1->getOutput(0), "ssh1"); + auto ssh2 = ssh(network, weightMap, *output2->getOutput(0), "ssh2"); + auto ssh3 = ssh(network, weightMap, *output3->getOutput(0), "ssh3"); + + // ------------- Head --------------- + auto bbox_head1 = network->addConvolution(*ssh1->getOutput(0), 2 * 4, DimsHW{1, 1}, weightMap["BboxHead.0.conv1x1.weight"], weightMap["BboxHead.0.conv1x1.bias"]); + auto bbox_head2 = network->addConvolution(*ssh2->getOutput(0), 2 * 4, DimsHW{1, 1}, weightMap["BboxHead.1.conv1x1.weight"], weightMap["BboxHead.1.conv1x1.bias"]); + auto bbox_head3 = network->addConvolution(*ssh3->getOutput(0), 2 * 4, DimsHW{1, 1}, weightMap["BboxHead.2.conv1x1.weight"], weightMap["BboxHead.2.conv1x1.bias"]); + + auto cls_head1 = network->addConvolution(*ssh1->getOutput(0), 2 * 2, DimsHW{1, 1}, weightMap["ClassHead.0.conv1x1.weight"], weightMap["ClassHead.0.conv1x1.bias"]); + auto cls_head2 = network->addConvolution(*ssh2->getOutput(0), 2 * 2, DimsHW{1, 1}, weightMap["ClassHead.1.conv1x1.weight"], weightMap["ClassHead.1.conv1x1.bias"]); + auto cls_head3 = network->addConvolution(*ssh3->getOutput(0), 2 * 2, DimsHW{1, 1}, weightMap["ClassHead.2.conv1x1.weight"], weightMap["ClassHead.2.conv1x1.bias"]); + + auto lmk_head1 = network->addConvolution(*ssh1->getOutput(0), 2 * 10, DimsHW{1, 1}, weightMap["LandmarkHead.0.conv1x1.weight"], weightMap["LandmarkHead.0.conv1x1.bias"]); + auto lmk_head2 = network->addConvolution(*ssh2->getOutput(0), 2 * 10, DimsHW{1, 1}, weightMap["LandmarkHead.1.conv1x1.weight"], weightMap["LandmarkHead.1.conv1x1.bias"]); + auto lmk_head3 = network->addConvolution(*ssh3->getOutput(0), 2 * 10, DimsHW{1, 1}, weightMap["LandmarkHead.2.conv1x1.weight"], weightMap["LandmarkHead.2.conv1x1.bias"]); + + // ------------- Decode bbox, conf, landmark --------------- + ITensor* inputTensors1[] = {bbox_head1->getOutput(0), cls_head1->getOutput(0), lmk_head1->getOutput(0)}; + auto cat1 = network->addConcatenation(inputTensors1, 3); + ITensor* inputTensors2[] = {bbox_head2->getOutput(0), cls_head2->getOutput(0), lmk_head2->getOutput(0)}; + auto cat2 = network->addConcatenation(inputTensors2, 3); + ITensor* inputTensors3[] = {bbox_head3->getOutput(0), cls_head3->getOutput(0), lmk_head3->getOutput(0)}; + auto cat3 = network->addConcatenation(inputTensors3, 3); + auto decode = new DecodePlugin(); + ITensor* inputTensors[] = {cat1->getOutput(0), cat2->getOutput(0), cat3->getOutput(0)}; + auto decodelayer = network->addPlugin(inputTensors, 3, *decode); + assert(decodelayer); + decodelayer->setName("decode"); + + decodelayer->getOutput(0)->setName(OUTPUT_BLOB_NAME); std::cout << "set name out" << std::endl; - network->markOutput(*output1->getOutput(0)); + network->markOutput(*decodelayer->getOutput(0)); // Build engine builder->setMaxBatchSize(maxBatchSize); @@ -467,21 +511,32 @@ int main(int argc, char** argv) { // data[i + 2 * INPUT_H * INPUT_W] = pr_img.at(i)[0] / 255.0; //} - //PluginFactory pf; + PluginFactory pf; IRuntime* runtime = createInferRuntime(gLogger); assert(runtime != nullptr); - //ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, &pf); - ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr); + ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, &pf); + //ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size, nullptr); assert(engine != nullptr); IExecutionContext* context = engine->createExecutionContext(); assert(context != nullptr); // Run inference static float prob[OUTPUT_SIZE]; - for (int i = 0; i < 10; i++) { + for (int i = 0; i < 1; i++) { auto start = std::chrono::system_clock::now(); doInference(*context, data, prob, 1); - //std::vector res; + std::vector res; + std::cout << "output 0 -> " << prob[0] << std::endl; + for (int j = 0; j < (int)prob[0]; j++) { + decodeplugin::Detection det; + memcpy(&det, &prob[1 + 15 * j], sizeof(decodeplugin::Detection)); + res.push_back(det); + } + sort(res.begin(), res.end(), cmp); + for (int j = 0; j < res.size(); j++) { + std::cout << res[j].class_confidence << std::endl; + std::cout << res[j].bbox[0] << ", " << res[j].bbox[1] << ", " << res[j].bbox[2] << ", " << res[j].bbox[3] << std::endl; + } //nms(res, prob); //for (size_t j = 0; j < res.size(); j++) { // float *p = (float*)&res[j]; @@ -504,13 +559,13 @@ int main(int argc, char** argv) { runtime->destroy(); // Print histogram of the output distribution - std::cout << "\nOutput:\n\n"; - for (unsigned int i = 0; i < OUTPUT_SIZE; i++) - { - std::cout << prob[i] << ", "; - if (i % 10 == 0) std::cout << i / 10 << std::endl; - } - std::cout << std::endl; + //std::cout << "\nOutput:\n\n"; + //for (unsigned int i = 0; i < OUTPUT_SIZE; i++) + //{ + // std::cout << prob[i] << ", "; + // if (i % 10 == 0) std::cout << i / 10 << std::endl; + //} + //std::cout << std::endl; return 0; }