verify yolov5 s/m/l/x and p6 backbone

yolov5/README.md

@@ -4,9 +4,10 @@ The Pytorch implementation is [ultralytics/yolov5](https://github.com/ultralytic
 
 ## Different versions of yolov5
 
-Currently, we support yolov5 v1.0(yolov5s only), v2.0, v3.0, v3.1 and v4.0.
+Currently, we support yolov5 v1.0(yolov5s only), v2.0, v3.0, v3.1, v4.0 and v5.0.
 
-- For yolov5 v4.0, download .pt from [yolov5 release v4.0](https://github.com/ultralytics/yolov5/releases/tag/v4.0), `git clone -b v4.0 https://github.com/ultralytics/yolov5.git` and `git clone https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in current page.
+- For yolov5 v5.0, download .pt from [yolov5 release v5.0](https://github.com/ultralytics/yolov5/releases/tag/v5.0), `git clone -b v5.0 https://github.com/ultralytics/yolov5.git` and `git clone https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in current page.
+- For yolov5 v4.0, download .pt from [yolov5 release v4.0](https://github.com/ultralytics/yolov5/releases/tag/v4.0), `git clone -b v4.0 https://github.com/ultralytics/yolov5.git` and `git clone -b yolov5-v4.0 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v4.0](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v4.0/yolov5).
 - For yolov5 v3.1, download .pt from [yolov5 release v3.1](https://github.com/ultralytics/yolov5/releases/tag/v3.1), `git clone -b v3.1 https://github.com/ultralytics/yolov5.git` and `git clone -b yolov5-v3.1 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v3.1](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v3.1/yolov5).
 - For yolov5 v3.0, download .pt from [yolov5 release v3.0](https://github.com/ultralytics/yolov5/releases/tag/v3.0), `git clone -b v3.0 https://github.com/ultralytics/yolov5.git` and `git clone -b yolov5-v3.0 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v3.0](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v3.0/yolov5).
 - For yolov5 v2.0, download .pt from [yolov5 release v2.0](https://github.com/ultralytics/yolov5/releases/tag/v2.0), `git clone -b v2.0 https://github.com/ultralytics/yolov5.git` and `git clone -b yolov5-v2.0 https://github.com/wang-xinyu/tensorrtx.git`, then follow how-to-run in [tensorrtx/yolov5-v2.0](https://github.com/wang-xinyu/tensorrtx/tree/yolov5-v2.0/yolov5).
@@ -28,20 +29,20 @@ Currently, we support yolov5 v1.0(yolov5s only), v2.0, v3.0, v3.1 and v4.0.
 1. generate .wts from pytorch with .pt, or download .wts from model zoo
 
 ```
-// git clone src code according to `Different versions of yolov5` above
+git clone -b v5.0 https://github.com/ultralytics/yolov5.git
-// download https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5s.pt
+git clone https://github.com/wang-xinyu/tensorrtx.git
-// copy tensorrtx/yolov5/gen_wts.py into ultralytics/yolov5
+// download https://github.com/ultralytics/yolov5/releases/download/v5.0/yolov5s.pt
-// ensure the file name is yolov5s.pt and yolov5s.wts in gen_wts.py
+cp {tensorrtx}/yolov5/gen_wts.py {ultralytics}/yolov5
-// go to ultralytics/yolov5
+cd {ultralytics}/yolov5
-python gen_wts.py
+python gen_wts.py yolov5s.pt
 // a file 'yolov5s.wts' will be generated.
 ```
 
 2. build tensorrtx/yolov5 and run
 
 ```
-// put yolov5s.wts into tensorrtx/yolov5
+cp {ultralytics}/yolov5/yolov5s.wts {tensorrtx}/yolov5/
-// go to tensorrtx/yolov5
+cd {tensorrtx}/yolov5/
 // update CLASS_NUM in yololayer.h if your model is trained on custom dataset
 mkdir build
 cd build

yolov5/gen_wts.py

@@ -1,14 +1,16 @@
 import torch
 import struct
+import sys
 from utils.torch_utils import select_device
 
 # Initialize
 device = select_device('cpu')
+pt_file = sys.argv[1]
 # Load model
-model = torch.load('weights/yolov5s.pt', map_location=device)['model'].float()  # load to FP32
+model = torch.load(pt_file, map_location=device)['model'].float()  # load to FP32
 model.to(device).eval()
 
-with open('yolov5s.wts', 'w') as f:
+with open(pt_file.split('.')[0] + '.wts', 'w') as f:
     f.write('{}\n'.format(len(model.state_dict().keys())))
     for k, v in model.state_dict().items():
         vr = v.reshape(-1).cpu().numpy()

yolov5/yolov5.cpp

@@ -128,6 +128,111 @@ ICudaEngine* build_engine(unsigned int maxBatchSize, IBuilder* builder, IBuilder
     return engine;
 }
 
+ICudaEngine* build_engine_p6(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt, float& gd, float& gw, std::string& wts_name) {
+    INetworkDefinition* network = builder->createNetworkV2(0U);
+
+    // Create input tensor of shape {3, INPUT_H, INPUT_W} 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(wts_name);
+
+    /* ------ yolov5 backbone------ */
+    auto focus0 = focus(network, weightMap, *data, 3, get_width(64, gw), 3, "model.0");
+    auto conv1 = convBlock(network, weightMap, *focus0->getOutput(0), get_width(128, gw), 3, 2, 1, "model.1");
+    auto c3_2 = C3(network, weightMap, *conv1->getOutput(0), get_width(128, gw), get_width(128, gw), get_depth(3, gd), true, 1, 0.5, "model.2");
+    auto conv3 = convBlock(network, weightMap, *c3_2->getOutput(0), get_width(256, gw), 3, 2, 1, "model.3");
+    auto c3_4 = C3(network, weightMap, *conv3->getOutput(0), get_width(256, gw), get_width(256, gw), get_depth(9, gd), true, 1, 0.5, "model.4");
+    auto conv5 = convBlock(network, weightMap, *c3_4->getOutput(0), get_width(512, gw), 3, 2, 1, "model.5");
+    auto c3_6 = C3(network, weightMap, *conv5->getOutput(0), get_width(512, gw), get_width(512, gw), get_depth(9, gd), true, 1, 0.5, "model.6");
+    auto conv7 = convBlock(network, weightMap, *c3_6->getOutput(0), get_width(768, gw), 3, 2, 1, "model.7");
+    auto c3_8 = C3(network, weightMap, *conv7->getOutput(0), get_width(768, gw), get_width(768, gw), get_depth(3, gd), true, 1, 0.5, "model.8");
+    auto conv9 = convBlock(network, weightMap, *c3_8->getOutput(0), get_width(1024, gw), 3, 2, 1, "model.9");
+    auto spp10 = SPP(network, weightMap, *conv9->getOutput(0), get_width(1024, gw), get_width(1024, gw), 3, 5, 7, "model.10");
+    auto c3_11 = C3(network, weightMap, *spp10->getOutput(0), get_width(1024, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.11");
+
+    /* ------ yolov5 head ------ */
+    auto conv12 = convBlock(network, weightMap, *c3_11->getOutput(0), get_width(768, gw), 1, 1, 1, "model.12");
+    auto upsample13 = network->addResize(*conv12->getOutput(0));
+    assert(upsample13);
+    upsample13->setResizeMode(ResizeMode::kNEAREST);
+    upsample13->setOutputDimensions(c3_8->getOutput(0)->getDimensions());
+    ITensor* inputTensors14[] = { upsample13->getOutput(0), c3_8->getOutput(0) };
+    auto cat14 = network->addConcatenation(inputTensors14, 2);
+    auto c3_15 = C3(network, weightMap, *cat14->getOutput(0), get_width(1536, gw), get_width(768, gw), get_depth(3, gd), false, 1, 0.5, "model.15");
+
+    auto conv16 = convBlock(network, weightMap, *c3_15->getOutput(0), get_width(512, gw), 1, 1, 1, "model.16");
+    auto upsample17 = network->addResize(*conv16->getOutput(0));
+    assert(upsample17);
+    upsample17->setResizeMode(ResizeMode::kNEAREST);
+    upsample17->setOutputDimensions(c3_6->getOutput(0)->getDimensions());
+    ITensor* inputTensors18[] = { upsample17->getOutput(0), c3_6->getOutput(0) };
+    auto cat18 = network->addConcatenation(inputTensors18, 2);
+    auto c3_19 = C3(network, weightMap, *cat18->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.19");
+
+    auto conv20 = convBlock(network, weightMap, *c3_19->getOutput(0), get_width(256, gw), 1, 1, 1, "model.20");
+    auto upsample21 = network->addResize(*conv20->getOutput(0));
+    assert(upsample21);
+    upsample21->setResizeMode(ResizeMode::kNEAREST);
+    upsample21->setOutputDimensions(c3_4->getOutput(0)->getDimensions());
+    ITensor* inputTensors21[] = { upsample21->getOutput(0), c3_4->getOutput(0) };
+    auto cat22 = network->addConcatenation(inputTensors21, 2);
+    auto c3_23 = C3(network, weightMap, *cat22->getOutput(0), get_width(512, gw), get_width(256, gw), get_depth(3, gd), false, 1, 0.5, "model.23");
+
+    auto conv24 = convBlock(network, weightMap, *c3_23->getOutput(0), get_width(256, gw), 3, 2, 1, "model.24");
+    ITensor* inputTensors25[] = { conv24->getOutput(0), conv20->getOutput(0) };
+    auto cat25 = network->addConcatenation(inputTensors25, 2);
+    auto c3_26 = C3(network, weightMap, *cat25->getOutput(0), get_width(1024, gw), get_width(512, gw), get_depth(3, gd), false, 1, 0.5, "model.26");
+
+    auto conv27 = convBlock(network, weightMap, *c3_26->getOutput(0), get_width(512, gw), 3, 2, 1, "model.27");
+    ITensor* inputTensors28[] = { conv27->getOutput(0), conv16->getOutput(0) };
+    auto cat28 = network->addConcatenation(inputTensors28, 2);
+    auto c3_29 = C3(network, weightMap, *cat28->getOutput(0), get_width(1536, gw), get_width(768, gw), get_depth(3, gd), false, 1, 0.5, "model.29");
+
+    auto conv30 = convBlock(network, weightMap, *c3_29->getOutput(0), get_width(768, gw), 3, 2, 1, "model.30");
+    ITensor* inputTensors31[] = { conv30->getOutput(0), conv12->getOutput(0) };
+    auto cat31 = network->addConcatenation(inputTensors31, 2);
+    auto c3_32 = C3(network, weightMap, *cat31->getOutput(0), get_width(2048, gw), get_width(1024, gw), get_depth(3, gd), false, 1, 0.5, "model.32");
+
+    /* ------ detect ------ */
+    IConvolutionLayer* det0 = network->addConvolutionNd(*c3_23->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.0.weight"], weightMap["model.33.m.0.bias"]);
+    IConvolutionLayer* det1 = network->addConvolutionNd(*c3_26->getOutput(0), 3 * (Yolo::CLASS_NUM + 5), DimsHW{ 1, 1 }, weightMap["model.33.m.1.weight"], weightMap["model.33.m.1.bias"]);
+    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);
+    yolo->getOutput(0)->setName(OUTPUT_BLOB_NAME);
+    network->markOutput(*yolo->getOutput(0));
+
+    // Build engine
+    builder->setMaxBatchSize(maxBatchSize);
+    config->setMaxWorkspaceSize(16 * (1 << 20));  // 16MB
+#if defined(USE_FP16)
+    config->setFlag(BuilderFlag::kFP16);
+#elif defined(USE_INT8)
+    std::cout << "Your platform support int8: " << (builder->platformHasFastInt8() ? "true" : "false") << std::endl;
+    assert(builder->platformHasFastInt8());
+    config->setFlag(BuilderFlag::kINT8);
+    Int8EntropyCalibrator2* calibrator = new Int8EntropyCalibrator2(1, INPUT_W, INPUT_H, "./coco_calib/", "int8calib.table", INPUT_BLOB_NAME);
+    config->setInt8Calibrator(calibrator);
+#endif
+
+    std::cout << "Building engine, please wait for a while..." << std::endl;
+    ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
+    std::cout << "Build engine successfully!" << 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, float& gd, float& gw, std::string& wts_name) {
     // Create builder
     IBuilder* builder = createInferBuilder(gLogger);