diff --git a/yolov5/README.md b/yolov5/README.md index d1272b9..4b0cc43 100644 --- a/yolov5/README.md +++ b/yolov5/README.md @@ -48,7 +48,7 @@ mkdir build cd build cmake .. make -sudo ./yolov5 -s [.wts] [.engine] [s/m/l/x or c gd gw] // serialize model to plan file +sudo ./yolov5 -s [.wts] [.engine] [s/m/l/x/s6/m6/l6/x6 or c/c6 gd gw] // serialize model to plan file sudo ./yolov5 -d [.engine] [image folder] // deserialize and run inference, the images in [image folder] will be processed. // For example yolov5s sudo ./yolov5 -s yolov5s.wts yolov5s.engine s diff --git a/yolov5/common.hpp b/yolov5/common.hpp index 3ef0377..2e1a1e8 100644 --- a/yolov5/common.hpp +++ b/yolov5/common.hpp @@ -256,68 +256,50 @@ ILayer* SPP(INetworkDefinition *network, std::map& weightM return cv2; } -std::vector getAnchors(std::map& weightMap) -{ - std::vector anchors_yolo; - Weights Yolo_Anchors = weightMap["model.24.anchor_grid"]; - assert(Yolo_Anchors.count == 18); - int each_yololayer_anchorsnum = Yolo_Anchors.count / 3; - const float* tempAnchors = (const float*)(Yolo_Anchors.values); - for (int i = 0; i < Yolo_Anchors.count; i++) - { - if (i < each_yololayer_anchorsnum) - { - anchors_yolo.push_back(const_cast(tempAnchors)[i]); - } - if ((i >= each_yololayer_anchorsnum) && (i < (2 * each_yololayer_anchorsnum))) - { - anchors_yolo.push_back(const_cast(tempAnchors)[i]); - } - if (i >= (2 * each_yololayer_anchorsnum)) - { - anchors_yolo.push_back(const_cast(tempAnchors)[i]); - } +std::vector> getAnchors(std::map& weightMap, std::string lname) { + std::vector> anchors; + Weights wts = weightMap[lname + ".anchor_grid"]; + int anchor_len = Yolo::CHECK_COUNT * 2; + for (int i = 0; i < wts.count / anchor_len; i++) { + auto *p = (const float*)wts.values + i * anchor_len; + std::vector anchor(p, p + anchor_len); + anchors.push_back(anchor); } - return anchors_yolo; + return anchors; } -IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map& weightMap, IConvolutionLayer* det0, IConvolutionLayer* det1, IConvolutionLayer* det2) -{ +IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map& weightMap, std::string lname, std::vector dets) { auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1"); - std::vector anchors_yolo = getAnchors(weightMap); - PluginField pluginMultidata[4]; - int NetData[4]; - NetData[0] = Yolo::CLASS_NUM; - NetData[1] = Yolo::INPUT_W; - NetData[2] = Yolo::INPUT_H; - NetData[3] = Yolo::MAX_OUTPUT_BBOX_COUNT; - pluginMultidata[0].data = NetData; - pluginMultidata[0].length = 3; - pluginMultidata[0].name = "netdata"; - pluginMultidata[0].type = PluginFieldType::kFLOAT32; - int scale[3] = { 8, 16, 32 }; - int plugindata[3][8]; - std::string names[3]; - for (int k = 1; k < 4; k++) - { - plugindata[k - 1][0] = Yolo::INPUT_W / scale[k - 1]; - plugindata[k - 1][1] = Yolo::INPUT_H / scale[k - 1]; - for (int i = 2; i < 8; i++) - { - plugindata[k - 1][i] = int(anchors_yolo[(k - 1) * 6 + i - 2]); - } - pluginMultidata[k].data = plugindata[k - 1]; - pluginMultidata[k].length = 8; - names[k - 1] = "yolodata" + std::to_string(k); - pluginMultidata[k].name = names[k - 1].c_str(); - pluginMultidata[k].type = PluginFieldType::kFLOAT32; + auto anchors = getAnchors(weightMap, lname); + PluginField plugin_fields[2]; + int netinfo[4] = {Yolo::CLASS_NUM, Yolo::INPUT_W, Yolo::INPUT_H, Yolo::MAX_OUTPUT_BBOX_COUNT}; + plugin_fields[0].data = netinfo; + plugin_fields[0].length = 4; + plugin_fields[0].name = "netinfo"; + plugin_fields[0].type = PluginFieldType::kFLOAT32; + int scale = 8; + std::vector kernels; + for (size_t i = 0; i < anchors.size(); i++) { + Yolo::YoloKernel kernel; + kernel.width = Yolo::INPUT_W / scale; + kernel.height = Yolo::INPUT_H / scale; + memcpy(kernel.anchors, &anchors[i][0], anchors[i].size() * sizeof(float)); + kernels.push_back(kernel); + scale *= 2; } - PluginFieldCollection pluginData; - pluginData.nbFields = 4; - pluginData.fields = pluginMultidata; - IPluginV2 *pluginObj = creator->createPlugin("yololayer", &pluginData); - ITensor* inputTensors_yolo[] = { det2->getOutput(0), det1->getOutput(0), det0->getOutput(0) }; - auto yolo = network->addPluginV2(inputTensors_yolo, 3, *pluginObj); + plugin_fields[1].data = &kernels[0]; + plugin_fields[1].length = kernels.size(); + plugin_fields[1].name = "kernels"; + plugin_fields[1].type = PluginFieldType::kFLOAT32; + PluginFieldCollection plugin_data; + plugin_data.nbFields = 2; + plugin_data.fields = plugin_fields; + IPluginV2 *plugin_obj = creator->createPlugin("yololayer", &plugin_data); + std::vector input_tensors; + for (auto det: dets) { + input_tensors.push_back(det->getOutput(0)); + } + auto yolo = network->addPluginV2(&input_tensors[0], input_tensors.size(), *plugin_obj); return yolo; } #endif diff --git a/yolov5/yololayer.cu b/yolov5/yololayer.cu index 9d95e29..525bf8d 100644 --- a/yolov5/yololayer.cu +++ b/yolov5/yololayer.cu @@ -192,7 +192,7 @@ namespace nvinfer1 int info_len_i = 5 + classes; const float* curInput = input + bnIdx * (info_len_i * total_grid * CHECK_COUNT); - for (int k = 0; k < 3; ++k) { + for (int k = 0; k < CHECK_COUNT; ++k) { float box_prob = Logist(curInput[idx + k * info_len_i * total_grid + 4 * total_grid]); if (box_prob < IGNORE_THRESH) continue; int class_id = 0; @@ -207,8 +207,8 @@ namespace nvinfer1 float *res_count = output + bnIdx * outputElem; int count = (int)atomicAdd(res_count, 1); if (count >= maxoutobject) return; - char* data = (char *)res_count + sizeof(float) + count * sizeof(Detection); - Detection* det = (Detection*)(data); + char *data = (char*)res_count + sizeof(float) + count * sizeof(Detection); + Detection *det = (Detection*)(data); int row = idx / yoloWidth; int col = idx % yoloWidth; @@ -217,12 +217,12 @@ namespace nvinfer1 // pytorch: // y = x[i].sigmoid() // y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i] # xy - // y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh - // X: (sigmoid(tx) + cx)/FeaturemapW * netwidth + // y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i] # wh + // X: (sigmoid(tx) + cx)/FeaturemapW * netwidth det->bbox[0] = (col - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 0 * total_grid])) * netwidth / yoloWidth; det->bbox[1] = (row - 0.5f + 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 1 * total_grid])) * netheight / yoloHeight; - // W: (Pw * e^tw) / FeaturemapW * netwidth + // W: (Pw * e^tw) / FeaturemapW * netwidth // v5: https://github.com/ultralytics/yolov5/issues/471 det->bbox[2] = 2.0f * Logist(curInput[idx + k * info_len_i * total_grid + 2 * total_grid]); det->bbox[2] = det->bbox[2] * det->bbox[2] * anchors[2 * k]; @@ -233,30 +233,28 @@ namespace nvinfer1 } } - void YoloLayerPlugin::forwardGpu(const float *const * inputs, float* output, cudaStream_t stream, int batchSize) + void YoloLayerPlugin::forwardGpu(const float* const* inputs, float *output, cudaStream_t stream, int batchSize) { int outputElem = 1 + mMaxOutObject * 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) - { + for (unsigned int i = 0; i < mYoloKernel.size(); ++i) { const auto& yolo = mYoloKernel[i]; - numElem = yolo.width*yolo.height*batchSize; - if (numElem < mThreadCount) - mThreadCount = numElem; + numElem = yolo.width * yolo.height * batchSize; + if (numElem < mThreadCount) mThreadCount = numElem; //printf("Net: %d %d \n", mYoloV5NetWidth, mYoloV5NetHeight); - CalDetection << < (yolo.width*yolo.height*batchSize + mThreadCount - 1) / mThreadCount, mThreadCount, 0, stream >> > - (inputs[i], output, numElem, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, yolo.width, yolo.height, (float *)mAnchor[i], mClassCount, outputElem); + CalDetection << < (numElem + mThreadCount - 1) / mThreadCount, mThreadCount, 0, stream >> > + (inputs[i], output, numElem, mYoloV5NetWidth, mYoloV5NetHeight, mMaxOutObject, yolo.width, yolo.height, (float*)mAnchor[i], mClassCount, outputElem); } } - int YoloLayerPlugin::enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) + int YoloLayerPlugin::enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) { - forwardGpu((const float *const *)inputs, (float*)outputs[0], stream, batchSize); + forwardGpu((const float* const*)inputs, (float*)outputs[0], stream, batchSize); return 0; } @@ -288,35 +286,17 @@ namespace nvinfer1 IPluginV2IOExt* YoloPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) { - int class_count = -1; - int input_w = -1; - int input_h = -1; - int max_output_object_count = -1; - std::vector yolo_kernels(3); - - const PluginField* fields = fc->fields; - for (int i = 0; i < fc->nbFields; i++) { - if (strcmp(fields[i].name, "netdata") == 0) { - assert(fields[i].type == PluginFieldType::kFLOAT32); - int *tmp = (int*)(fields[i].data); - class_count = tmp[0]; - input_w = tmp[1]; - input_h = tmp[2]; - max_output_object_count = tmp[3]; - } else if (strstr(fields[i].name, "yolodata") != NULL) { - assert(fields[i].type == PluginFieldType::kFLOAT32); - int *tmp = (int*)(fields[i].data); - YoloKernel kernel; - kernel.width = tmp[0]; - kernel.height = tmp[1]; - for (int j = 0; j < fields[i].length - 2; j++) { - kernel.anchors[j] = tmp[j + 2]; - } - yolo_kernels[2 - (fields[i].name[8] - '1')] = kernel; - } - } - assert(class_count && input_w && input_h && max_output_object_count); - YoloLayerPlugin* obj = new YoloLayerPlugin(class_count, input_w, input_h, max_output_object_count, yolo_kernels); + assert(fc->nbFields == 2); + assert(strcmp(fc->fields[0].name, "netinfo") == 0); + assert(strcmp(fc->fields[1].name, "kernels") == 0); + int *p_netinfo = (int*)(fc->fields[0].data); + int class_count = p_netinfo[0]; + int input_w = p_netinfo[1]; + int input_h = p_netinfo[2]; + int max_output_object_count = p_netinfo[3]; + std::vector kernels(fc->fields[1].length); + memcpy(&kernels[0], fc->fields[1].data, kernels.size() * sizeof(Yolo::YoloKernel)); + YoloLayerPlugin* obj = new YoloLayerPlugin(class_count, input_w, input_h, max_output_object_count, kernels); obj->setPluginNamespace(mNamespace.c_str()); return obj; } diff --git a/yolov5/yololayer.h b/yolov5/yololayer.h index 7357259..49f6474 100644 --- a/yolov5/yololayer.h +++ b/yolov5/yololayer.h @@ -51,7 +51,7 @@ namespace nvinfer1 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 int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) override; virtual size_t getSerializationSize() const override; @@ -87,7 +87,7 @@ namespace nvinfer1 void detachFromContext() override; private: - void forwardGpu(const float *const * inputs, float * output, cudaStream_t stream, int batchSize = 1); + void forwardGpu(const float* const* inputs, float *output, cudaStream_t stream, int batchSize = 1); int mThreadCount = 256; const char* mPluginNamespace; int mKernelCount; diff --git a/yolov5/yolov5.cpp b/yolov5/yolov5.cpp index 2fcf50b..6601a63 100644 --- a/yolov5/yolov5.cpp +++ b/yolov5/yolov5.cpp @@ -75,19 +75,18 @@ ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilder assert(upsample15); upsample15->setResizeMode(ResizeMode::kNEAREST); upsample15->setOutputDimensions(bottleneck_csp4->getOutput(0)->getDimensions()); - + ITensor* inputTensors16[] = { upsample15->getOutput(0), bottleneck_csp4->getOutput(0) }; auto cat16 = network->addConcatenation(inputTensors16, 2); auto bottleneck_csp17 = C3(network, weightMap, *cat16->getOutput(0), get_width(512, gw), get_width(256, gw), get_depth(3, gd), false, 1, 0.5, "model.17"); - // yolo layer 0 + /* ------ detect ------ */ IConvolutionLayer* det0 = network->addConvolutionNd(*bottleneck_csp17->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.0.weight"], weightMap["model.24.m.0.bias"]); auto conv18 = convBlock(network, weightMap, *bottleneck_csp17->getOutput(0), get_width(256, gw), 3, 2, 1, "model.18"); ITensor* inputTensors19[] = { conv18->getOutput(0), conv14->getOutput(0) }; auto cat19 = network->addConcatenation(inputTensors19, 2); auto bottleneck_csp20 = C3(network, weightMap, *cat19->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.20"); - //yolo layer 1 IConvolutionLayer* det1 = network->addConvolutionNd(*bottleneck_csp20->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.1.weight"], weightMap["model.24.m.1.bias"]); auto conv21 = convBlock(network, weightMap, *bottleneck_csp20->getOutput(0), get_width(512, gw), 3, 2, 1, "model.21"); ITensor* inputTensors22[] = { conv21->getOutput(0), conv10->getOutput(0) }; @@ -95,7 +94,7 @@ ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilder auto bottleneck_csp23 = C3(network, weightMap, *cat22->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.23"); IConvolutionLayer* det2 = network->addConvolutionNd(*bottleneck_csp23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.24.m.2.weight"], weightMap["model.24.m.2.bias"]); - auto yolo = addYoLoLayer(network, weightMap, det0, det1, det2); + auto yolo = addYoLoLayer(network, weightMap, "model.24", std::vector{det0, det1, det2}); yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME); network->markOutput(*yolo->getOutput(0)); @@ -200,7 +199,7 @@ ICudaEngine* build_engine_p6(unsigned int maxBatchSize, IBuilder* builder, IBuil IConvolutionLayer* det2 = network->addConvolutionNd(*c3_29->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.2.weight"], weightMap["model.33.m.2.bias"]); IConvolutionLayer* det3 = network->addConvolutionNd(*c3_32->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.3.weight"], weightMap["model.33.m.3.bias"]); - auto yolo = addYoLoLayer(network, weightMap, det0, det1, det2); + auto yolo = addYoLoLayer(network, weightMap, "model.33", std::vector{det0, det1, det2, det3}); yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME); network->markOutput(*yolo->getOutput(0)); @@ -233,13 +232,18 @@ ICudaEngine* build_engine_p6(unsigned int maxBatchSize, IBuilder* builder, IBuil return engine; } -void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream, float& gd, float& gw, std::string& wts_name) { +void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream, bool& is_p6, float& gd, float& gw, std::string& wts_name) { // Create builder IBuilder* builder = createInferBuilder(gLogger); IBuilderConfig* config = builder->createBuilderConfig(); // Create model to populate the network, then set the outputs and create an engine - ICudaEngine* engine = build_engine(maxBatchSize, builder, config, DataType::kFLOAT, gd, gw, wts_name); + ICudaEngine *engine = nullptr; + if (is_p6) { + engine = build_engine_p6(maxBatchSize, builder, config, DataType::kFLOAT, gd, gw, wts_name); + } else { + engine = build_engine(maxBatchSize, builder, config, DataType::kFLOAT, gd, gw, wts_name); + } assert(engine != nullptr); // Serialize the engine @@ -259,30 +263,33 @@ void doInference(IExecutionContext& context, cudaStream_t& stream, void **buffer cudaStreamSynchronize(stream); } -bool parse_args(int argc, char** argv, std::string& wts, std::string& engine, float& gd, float& gw, std::string& img_dir) { +bool parse_args(int argc, char** argv, std::string& wts, std::string& engine, bool& is_p6, float& gd, float& gw, std::string& img_dir) { if (argc < 4) return false; if (std::string(argv[1]) == "-s" && (argc == 5 || argc == 7)) { wts = std::string(argv[2]); engine = std::string(argv[3]); auto net = std::string(argv[4]); - if (net == "s") { + if (net[0] == 's') { gd = 0.33; gw = 0.50; - } else if (net == "m") { + } else if (net[0] == 'm') { gd = 0.67; gw = 0.75; - } else if (net == "l") { + } else if (net[0] == 'l') { gd = 1.0; gw = 1.0; - } else if (net == "x") { + } else if (net[0] == 'x') { gd = 1.33; gw = 1.25; - } else if (net == "c" && argc == 7) { + } else if (net[0] == 'c' && argc == 7) { gd = atof(argv[5]); gw = atof(argv[6]); } else { return false; } + if (net.size() == 2 && net[1] == '6') { + is_p6 = true; + } } else if (std::string(argv[1]) == "-d" && argc == 4) { engine = std::string(argv[2]); img_dir = std::string(argv[3]); @@ -297,11 +304,12 @@ int main(int argc, char** argv) { std::string wts_name = ""; std::string engine_name = ""; + bool is_p6 = false; float gd = 0.0f, gw = 0.0f; std::string img_dir; - if (!parse_args(argc, argv, wts_name, engine_name, gd, gw, img_dir)) { + if (!parse_args(argc, argv, wts_name, engine_name, is_p6, gd, gw, img_dir)) { std::cerr << "arguments not right!" << std::endl; - std::cerr << "./yolov5 -s [.wts] [.engine] [s/m/l/x or c gd gw] // serialize model to plan file" << std::endl; + std::cerr << "./yolov5 -s [.wts] [.engine] [s/m/l/x/s6/m6/l6/x6 or c/c6 gd gw] // serialize model to plan file" << std::endl; std::cerr << "./yolov5 -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; return -1; } @@ -309,7 +317,7 @@ int main(int argc, char** argv) { // create a model using the API directly and serialize it to a stream if (!wts_name.empty()) { IHostMemory* modelStream{ nullptr }; - APIToModel(BATCH_SIZE, &modelStream, gd, gw, wts_name); + APIToModel(BATCH_SIZE, &modelStream, is_p6, gd, gw, wts_name); assert(modelStream != nullptr); std::ofstream p(engine_name, std::ios::binary); if (!p) {