#include "ibnnet.h" //#define USE_FP16 namespace trt { IBNNet::IBNNet(trt::EngineConfig &enginecfg, const IBN ibn) : _engineCfg(enginecfg) { switch(ibn) { case IBN::A: _ibn = "a"; break; case IBN::B: _ibn = "b"; break; case IBN::NONE: default: _ibn = ""; break; } } // create the engine using only the API and not any parser. ICudaEngine *IBNNet::createEngine(IBuilder* builder, IBuilderConfig* config) { // resnet50-ibna, resnet50-ibnb, resnet50 assert(_ibn == "a" or _ibn == "b" or _ibn == ""); INetworkDefinition* network = builder->createNetworkV2(0U); // Create input tensor of shape { 3, INPUT_H, INPUT_W } with name INPUT_BLOB_NAME ITensor* data = network->addInput(_engineCfg.input_name, _dt, Dims3{3, _engineCfg.input_h, _engineCfg.input_w}); assert(data); std::string path; if(_ibn == "") { path = "../resnet50.wts"; } else { path = "../resnet50-ibn" + _ibn + ".wts"; } std::map weightMap = loadWeights(path); Weights emptywts{DataType::kFLOAT, nullptr, 0}; std::map> ibn_layers{ { "a", {"a", "a", "a", "a", "a", "a", "a", "a", "a", "a", "a", "a", "a", "", "", ""}}, { "b", {"", "", "b", "", "", "","b", "", "", "", "", "", "", "", "", "",}}, { "", {16, ""}}}; const float mean[3] = {0.485, 0.456, 0.406}; // rgb const float std[3] = {0.229, 0.224, 0.225}; ITensor* pre_input = MeanStd(network, weightMap, data, "", mean, std, false); IConvolutionLayer* conv1 = network->addConvolutionNd(*pre_input, 64, DimsHW{7, 7}, weightMap["conv1.weight"], emptywts); assert(conv1); conv1->setStrideNd(DimsHW{2, 2}); conv1->setPaddingNd(DimsHW{3, 3}); IActivationLayer* relu1{nullptr}; if (_ibn == "b") { IScaleLayer* bn1 = addInstanceNorm2d(network, weightMap, *conv1->getOutput(0), "bn1", 1e-5); relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU); } else { IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), "bn1", 1e-5); relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU); } assert(relu1); // Add max pooling layer with stride of 2x2 and kernel size of 2x2. IPoolingLayer* pool1 = network->addPoolingNd(*relu1->getOutput(0), PoolingType::kMAX, DimsHW{3, 3}); assert(pool1); pool1->setStrideNd(DimsHW{2, 2}); pool1->setPaddingNd(DimsHW{1, 1}); IActivationLayer* x = bottleneck_ibn(network, weightMap, *pool1->getOutput(0), 64, 64, 1, "layer1.0.", ibn_layers[_ibn][0]); x = bottleneck_ibn(network, weightMap, *x->getOutput(0), 256, 64, 1, "layer1.1.", ibn_layers[_ibn][1]); x = bottleneck_ibn(network, weightMap, *x->getOutput(0), 256, 64, 1, "layer1.2.", ibn_layers[_ibn][2]); x = bottleneck_ibn(network, weightMap, *x->getOutput(0), 256, 128, 2, "layer2.0.", ibn_layers[_ibn][3]); x = bottleneck_ibn(network, weightMap, *x->getOutput(0), 512, 128, 1, "layer2.1.", ibn_layers[_ibn][4]); x = bottleneck_ibn(network, weightMap, *x->getOutput(0), 512, 128, 1, "layer2.2.", ibn_layers[_ibn][5]); x = bottleneck_ibn(network, weightMap, *x->getOutput(0), 512, 128, 1, "layer2.3.", ibn_layers[_ibn][6]); x = bottleneck_ibn(network, weightMap, *x->getOutput(0), 512, 256, 2, "layer3.0.", ibn_layers[_ibn][7]); x = bottleneck_ibn(network, weightMap, *x->getOutput(0), 1024, 256, 1, "layer3.1.", ibn_layers[_ibn][8]); x = bottleneck_ibn(network, weightMap, *x->getOutput(0), 1024, 256, 1, "layer3.2.", ibn_layers[_ibn][9]); x = bottleneck_ibn(network, weightMap, *x->getOutput(0), 1024, 256, 1, "layer3.3.", ibn_layers[_ibn][10]); x = bottleneck_ibn(network, weightMap, *x->getOutput(0), 1024, 256, 1, "layer3.4.", ibn_layers[_ibn][11]); x = bottleneck_ibn(network, weightMap, *x->getOutput(0), 1024, 256, 1, "layer3.5.", ibn_layers[_ibn][12]); x = bottleneck_ibn(network, weightMap, *x->getOutput(0), 1024, 512, 2, "layer4.0.", ibn_layers[_ibn][13]); x = bottleneck_ibn(network, weightMap, *x->getOutput(0), 2048, 512, 1, "layer4.1.", ibn_layers[_ibn][14]); x = bottleneck_ibn(network, weightMap, *x->getOutput(0), 2048, 512, 1, "layer4.2.", ibn_layers[_ibn][15]); IPoolingLayer* pool2 = network->addPoolingNd(*x->getOutput(0), PoolingType::kAVERAGE, DimsHW{7, 7}); assert(pool2); pool2->setStrideNd(DimsHW{1, 1}); IFullyConnectedLayer* fc1 = network->addFullyConnected(*pool2->getOutput(0), 1000, weightMap["fc.weight"], weightMap["fc.bias"]); assert(fc1); fc1->getOutput(0)->setName(_engineCfg.output_name); std::cout << "set name out" << std::endl; network->markOutput(*fc1->getOutput(0)); // Build engine builder->setMaxBatchSize(_engineCfg.max_batch_size); config->setMaxWorkspaceSize(1 << 20); #ifdef USE_FP16 config->setFlag(BuilderFlag::kFP16); #endif ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config); 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; } bool IBNNet::serializeEngine() { // Create builder auto builder = make_holder(createInferBuilder(gLogger)); auto config = make_holder(builder->createBuilderConfig()); // Create model to populate the network, then set the outputs and create an engine ICudaEngine *engine = createEngine(builder.get(), config.get()); assert(engine); // Serialize the engine TensorRTHolder modelStream = make_holder(engine->serialize()); assert(modelStream); std::ofstream p("./ibnnet.engine", std::ios::binary | std::ios::out); if (!p) { std::cerr << "could not open plan output file" << std::endl; return false; } p.write(reinterpret_cast(modelStream->data()), modelStream->size()); return true; } bool IBNNet::deserializeEngine() { std::ifstream file("./ibnnet.engine", std::ios::binary | std::ios::in); if (file.good()) { file.seekg(0, file.end); _engineCfg.stream_size = file.tellg(); file.seekg(0, file.beg); _engineCfg.trtModelStream = std::shared_ptr( new char[_engineCfg.stream_size], []( char* ptr ){ delete [] ptr; } ); assert(_engineCfg.trtModelStream.get()); file.read(_engineCfg.trtModelStream.get(), _engineCfg.stream_size); file.close(); _inferEngine = make_unique(_engineCfg); return true; } return false; } void IBNNet::preprocessing(const cv::Mat& img, float* const data, const std::size_t stride) { for (std::size_t i = 0; i < stride; ++i) { data[i] = img.at(i)[2] / 255.0; data[i + stride] = img.at(i)[1] / 255.0; data[i + (stride<<1)] = img.at(i)[0] / 255.0; } } bool IBNNet::inference(std::vector &input) { if(_inferEngine != nullptr) { const std::size_t stride = _engineCfg.input_w * _engineCfg.input_h; return _inferEngine.get()->doInference(input.size(), [&](float* data) { for(const auto &img : input) { preprocessing(img, data, stride); data += 3 * stride; } } ); } else { return false; } } float* IBNNet::getOutput() { if(_inferEngine != nullptr) return _inferEngine.get()->getOutput(); return nullptr; } int IBNNet::getDeviceID() { return _engineCfg.device_id; } }