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

yolov5 support s6/m6/l6/x6 models

Browse Source
This commit is contained in:
wang-xinyu 2021-04-25 12:17:15 +00:00
parent 3e9a4fc19b
commit 046dd5b417
5 changed files with 91 additions and 121 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
yolov5/README.md
View File

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

96
yolov5/common.hpp
View File

@ -256,68 +256,50 @@ ILayer* SPP(INetworkDefinition *network, std::map<std::string, Weights>& weightM
return cv2;
}
std::vector<float> getAnchors(std::map<std::string, Weights>& weightMap)
{
std::vector<float> 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<float*>(tempAnchors)[i]);
}
if ((i >= each_yololayer_anchorsnum) && (i < (2 * each_yololayer_anchorsnum)))
{
anchors_yolo.push_back(const_cast<float*>(tempAnchors)[i]);
}
if (i >= (2 * each_yololayer_anchorsnum))
{
anchors_yolo.push_back(const_cast<float*>(tempAnchors)[i]);
}
std::vector<std::vector<float>> getAnchors(std::map<std::string, Weights>& weightMap, std::string lname) {
std::vector<std::vector<float>> 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<float> anchor(p, p + anchor_len);
anchors.push_back(anchor);
}
return anchors_yolo;
return anchors;
}
IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, IConvolutionLayer* det0, IConvolutionLayer* det1, IConvolutionLayer* det2)
{
IPluginV2Layer* addYoLoLayer(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, std::string lname, std::vector<IConvolutionLayer*> dets) {
auto creator = getPluginRegistry()->getPluginCreator("YoloLayer_TRT", "1");
std::vector<float> 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<Yolo::YoloKernel> 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<ITensor*> 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

70
yolov5/yololayer.cu
View File

@ -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::YoloKernel> 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<Yolo::YoloKernel> 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;
}

4
yolov5/yololayer.h
View File

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

40
yolov5/yolov5.cpp
View File

@ -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<IConvolutionLayer*>{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<IConvolutionLayer*>{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) {