diff --git a/hrnet/hrnet-semantic-segmentation/CMakeLists.txt b/hrnet/hrnet-semantic-segmentation/CMakeLists.txt index 40ce0fd..da49a83 100644 --- a/hrnet/hrnet-semantic-segmentation/CMakeLists.txt +++ b/hrnet/hrnet-semantic-segmentation/CMakeLists.txt @@ -24,10 +24,17 @@ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -Ofast -Wfatal-errors - find_package(OpenCV) include_directories(${OpenCV_INCLUDE_DIRS}) -add_executable(hrnetseg ${PROJECT_SOURCE_DIR}/hrnetseg.cpp) -target_link_libraries(hrnetseg nvinfer) -target_link_libraries(hrnetseg cudart) -target_link_libraries(hrnetseg ${OpenCV_LIBS}) +add_executable(hrnet ${PROJECT_SOURCE_DIR}/hrnet.cpp) +target_link_libraries(hrnet nvinfer) +target_link_libraries(hrnet cudart) +target_link_libraries(hrnet ${OpenCV_LIBS}) + + +add_executable(hrnet_ocr ${PROJECT_SOURCE_DIR}/hrnet_ocr.cpp) +target_link_libraries(hrnet_ocr nvinfer) +target_link_libraries(hrnet_ocr cudart) +target_link_libraries(hrnet_ocr ${OpenCV_LIBS}) + add_definitions(-O2 -pthread) diff --git a/hrnet/hrnet-semantic-segmentation/README.md b/hrnet/hrnet-semantic-segmentation/README.md index ccf6591..6c42448 100644 --- a/hrnet/hrnet-semantic-segmentation/README.md +++ b/hrnet/hrnet-semantic-segmentation/README.md @@ -1,27 +1,76 @@ # HRNet-Semantic-Segmentation -The Pytorch implementation is [HRNet-Semantic-Segmentation-v1.1](https://github.com/HRNet/HRNet-Semantic-Segmentation/tree/pytorch-v1.1). The implemented model is **HRNetV2-W18-Small-v2** +This repo implemtents [HRNet-Semantic-Segmentation-v1.1](https://github.com/HRNet/HRNet-Semantic-Segmentation/tree/pytorch-v1.1) and [HRNet-Semantic-Segmentation-OCR](https://github.com/HRNet/HRNet-Semantic-Segmentation/tree/HRNet-OCR). ## How to Run - -* 1. generate .wts - - Download code and model from [HRNet-Semantic-Segmentation-v1.1](https://github.com/HRNet/HRNet-Semantic-Segmentation/tree/pytorch-v1.1) and config your environments. - - Put `demo.py` in the `YOUR_ROOT_DIR\HRNet-Semantic-Segmentation\tools ` folder, set `savewts in main()` as `True`, and run, the .wts will be generated. - -* 2. cmake and make +### For HRNet-Semantic-Segmentation-v1.1 +1. generate .wts, use config `experiments/cityscapes/seg_hrnet_w48_train_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml` and pretrained weight `hrnet_w48_cityscapes_cls19_1024x2048_trainset.pth` as example. change `PRETRAINED` in `experiments/cityscapes/seg_hrnet_w48_train_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml` to `""`. +``` +cp gen_wts.py $HRNET--Semantic-Segmentation-PROJECT-ROOT/tools +cd $HRNET--Semantic-Segmentation-PROJECT-ROOT +python tools/gen_wts.py --cfg experiments/cityscapes/seg_hrnet_w48_train_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml --ckpt_path hrnet_w48_cityscapes_cls19_1024x2048_trainset.pth --save_path hrnet_w48.wts +cp hrnet_w48.wts $HRNET-TENSORRT-ROOT +cd $HRNET-TENSORRT-ROOT +``` +2. cmake and make ``` mkdir build cd build cmake .. make - sudo ./hrnetseg -s // serialize model to plan file i.e. 'hrnetseg.engine' - sudo ./hrnetseg -d ../samples // deserialize plan file and run inference, the images in samples will be processed. ``` + first serialize model to plan file + ``` + ./hrnet -s [.wts] [.engine] [small or 18 or 32 or 48] # small for W18-Small-v2, 18 for W18, etc. + ``` + such as + ``` + ./hrnet -s ../hrnet_w48.wts ./hrnet_w48.engine 48 + ``` + then deserialize plan file and run inference + ``` + ./hrnet -d [.engine] [image dir] + ``` + such as + ``` + ./hrnet -d ./hrnet_w48.engine ../samples + ``` +### For HRNet-Semantic-Segmentation-OCR +1. generate .wts, use config `experiments/cityscapes/seg_hrnet_ocr_w48_train_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml` and pretrained weight `hrnet_ocr_cs_8162_torch11.pth` as example. change `PRETRAINED` in `experiments/cityscapes/seg_hrnet_ocr_w48_train_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml` to `""`. +``` +cp gen_wts.py $HRNET-OCR-TRAIN-PROJECT-ROOT/tools +cd $HRNET-OCR-PROJECT-ROOT +python tools/gen_wts.py --cfg experiments/cityscapes/seg_hrnet_ocr_w48_train_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml --ckpt_path hrnet_ocr_cs_8162_torch11.pth --save_path hrnet_ocr_w48.wts +cp hrnet_ocr_w48.wts $HRNET-OCR-TENSORRT-ROOT +cd $HRNET-OCR-TENSORRT-ROOT +``` +2. cmake and make + + ``` + mkdir build + cd build + cmake .. + make + ``` + first serialize model to plan file + ``` + ./hrnet_ocr -s [.wts] [.engine] [18 or 32 or 48] + ``` + such as + ``` + ./hrnet_ocr -s ../hrnet_ocr_w48.wts ./hrnet_ocr_w48.engine 48 + ``` + then deserialize plan file and run inference + ``` + ./hrnet_ocr -d [.engine] [image dir] + ``` + such as + ``` + ./hrnet_ocr -d ./hrnet_ocr_w48.engine ../samples + ``` ## Result TRT Result: diff --git a/hrnet/hrnet-semantic-segmentation/common.hpp b/hrnet/hrnet-semantic-segmentation/common.hpp index 11c4f83..8219c5b 100644 --- a/hrnet/hrnet-semantic-segmentation/common.hpp +++ b/hrnet/hrnet-semantic-segmentation/common.hpp @@ -12,15 +12,15 @@ using namespace nvinfer1; -#define CHECK(status) \ - do\ - {\ - auto ret = (status);\ - if (ret != 0)\ - {\ - std::cerr << "Cuda failure: " << ret << std::endl;\ - abort();\ - }\ +#define CHECK(status) \ + do \ + { \ + auto ret = (status); \ + if (ret != 0) \ + { \ + std::cerr << "Cuda failure: " << ret << std::endl; \ + abort(); \ + } \ } while (0) int read_files_in_dir(const char *p_dir_name, std::vector &file_names) { @@ -45,9 +45,21 @@ int read_files_in_dir(const char *p_dir_name, std::vector &file_nam return 0; } + // TensorRT weight files have a simple space delimited format: // [type] [size] -std::map loadWeights(const std::string file) { +void debug_print(ITensor *input_tensor, std::string head) +{ + std::cout << head << " : "; + + for (int i = 0; i < input_tensor->getDimensions().nbDims; i++) + { + std::cout << input_tensor->getDimensions().d[i] << " "; + } + std::cout << std::endl; +} +std::map loadWeights(const std::string file) +{ std::cout << "Loading weights: " << file << std::endl; std::map weightMap; @@ -62,7 +74,7 @@ std::map loadWeights(const std::string file) { while (count--) { - Weights wt{ DataType::kFLOAT, nullptr, 0 }; + Weights wt{DataType::kFLOAT, nullptr, 0}; uint32_t size; // Read name and type of blob @@ -71,7 +83,7 @@ std::map loadWeights(const std::string file) { wt.type = DataType::kFLOAT; // Load blob - uint32_t* val = reinterpret_cast(malloc(sizeof(val) * size)); + uint32_t *val = reinterpret_cast(malloc(sizeof(val) * size)); for (uint32_t x = 0, y = size; x < y; ++x) { input >> std::hex >> val[x]; @@ -85,96 +97,110 @@ std::map loadWeights(const std::string file) { return weightMap; } -cv::Mat createLTU(int len) { +cv::Mat createLTU(int len) +{ cv::Mat lookUpTable(1, 256, CV_8U); - uchar* p = lookUpTable.data; - for (int j = 0; j < 256; ++j) { - p[j] = (j * (256 / len) > 255) ? uchar(255):(uchar)(j * (256 / len)); + uchar *p = lookUpTable.data; + for (int j = 0; j < 256; ++j) + { + p[j] = (j * (256 / len) > 255) ? uchar(255) : (uchar)(j * (256 / len)); } return lookUpTable; } -ITensor* MeanStd(INetworkDefinition *network, ITensor* input, float* mean, float* std, bool div255) { - if (div255) { - Weights Div_225{ DataType::kFLOAT, nullptr, 3 }; - float *wgt = reinterpret_cast(malloc(sizeof(float) * 3)); - for (int i = 0; i < 3; ++i) { +ITensor *MeanStd(INetworkDefinition *network, ITensor *input, float *mean, float *std, bool div255) +{ + if (div255) + { + Weights Div_225{DataType::kFLOAT, nullptr, 3}; + float *wgt = reinterpret_cast(malloc(sizeof(float) * 3)); + for (int i = 0; i < 3; ++i) + { wgt[i] = 255.0f; } Div_225.values = wgt; - IConstantLayer* d = network->addConstant(Dims3{ 3, 1, 1 }, Div_225); + IConstantLayer *d = network->addConstant(Dims3{3, 1, 1}, Div_225); input = network->addElementWise(*input, *d->getOutput(0), ElementWiseOperation::kDIV)->getOutput(0); } - Weights Mean{ DataType::kFLOAT, nullptr, 3 }; + Weights Mean{DataType::kFLOAT, nullptr, 3}; Mean.values = mean; - IConstantLayer* m = network->addConstant(Dims3{ 3, 1, 1 }, Mean); - IElementWiseLayer* sub_mean = network->addElementWise(*input, *m->getOutput(0), ElementWiseOperation::kSUB); - if (std != nullptr) { - Weights Std{ DataType::kFLOAT, nullptr, 3 }; + IConstantLayer *m = network->addConstant(Dims3{3, 1, 1}, Mean); + IElementWiseLayer *sub_mean = network->addElementWise(*input, *m->getOutput(0), ElementWiseOperation::kSUB); + if (std != nullptr) + { + Weights Std{DataType::kFLOAT, nullptr, 3}; Std.values = std; - IConstantLayer* s = network->addConstant(Dims3{ 3, 1, 1 }, Std); - IElementWiseLayer* std_mean = network->addElementWise(*sub_mean->getOutput(0), *s->getOutput(0), ElementWiseOperation::kDIV); + IConstantLayer *s = network->addConstant(Dims3{3, 1, 1}, Std); + IElementWiseLayer *std_mean = network->addElementWise(*sub_mean->getOutput(0), *s->getOutput(0), ElementWiseOperation::kDIV); return std_mean->getOutput(0); } - else { + else + { return sub_mean->getOutput(0); } } -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; - float *mean = (float*)weightMap[lname + ".running_mean"].values; - float *var = (float*)weightMap[lname + ".running_var"].values; +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; + float *mean = (float *)weightMap[lname + ".running_mean"].values; + float *var = (float *)weightMap[lname + ".running_var"].values; int len = weightMap[lname + ".running_var"].count; //std::cout << "len " << len << std::endl; - float *scval = reinterpret_cast(malloc(sizeof(float) * len)); - for (int i = 0; i < len; i++) { + float *scval = reinterpret_cast(malloc(sizeof(float) * len)); + for (int i = 0; i < len; i++) + { scval[i] = gamma[i] / sqrt(var[i] + eps); } - Weights scale{ DataType::kFLOAT, scval, len }; + Weights scale{DataType::kFLOAT, scval, len}; - float *shval = reinterpret_cast(malloc(sizeof(float) * len)); - for (int i = 0; i < len; i++) { + float *shval = reinterpret_cast(malloc(sizeof(float) * len)); + for (int i = 0; i < len; i++) + { shval[i] = beta[i] - mean[i] * gamma[i] / sqrt(var[i] + eps); } - Weights shift{ DataType::kFLOAT, shval, len }; + Weights shift{DataType::kFLOAT, shval, len}; - float *pval = reinterpret_cast(malloc(sizeof(float) * len)); - for (int i = 0; i < len; i++) { + float *pval = reinterpret_cast(malloc(sizeof(float) * len)); + for (int i = 0; i < len; i++) + { pval[i] = 1.0; } - Weights power{ DataType::kFLOAT, pval, len }; + Weights power{DataType::kFLOAT, pval, len}; weightMap[lname + ".scale"] = scale; weightMap[lname + ".shift"] = shift; weightMap[lname + ".power"] = power; - IScaleLayer* scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power); + IScaleLayer *scale_1 = network->addScale(input, ScaleMode::kCHANNEL, shift, scale, power); assert(scale_1); return scale_1; } -ILayer* convBnLeaky(INetworkDefinition *network, - std::map& weightMap, - ITensor& input, int outch, int ksize, int s, int p, - std::string convname, std::string bnname, - bool relu = true, - bool bias = false) { - Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; - IConvolutionLayer* conv1; +ILayer *convBnRelu(INetworkDefinition *network, + std::map &weightMap, + ITensor &input, int outch, int ksize, int s, int p, + std::string convname, std::string bnname, + bool relu = true, + bool bias = false) +{ + Weights emptywts{DataType::kFLOAT, nullptr, 0}; + IConvolutionLayer *conv1; //Dims dim; if (!bias) { - conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[convname + ".weight"], emptywts); + conv1 = network->addConvolutionNd(input, outch, DimsHW{ksize, ksize}, weightMap[convname + ".weight"], emptywts); } else { - conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[convname + ".weight"], weightMap[convname + ".bias"]); + conv1 = network->addConvolutionNd(input, outch, DimsHW{ksize, ksize}, weightMap[convname + ".weight"], weightMap[convname + ".bias"]); } assert(conv1); - conv1->setStrideNd(DimsHW{ s, s }); - conv1->setPaddingNd(DimsHW{ p, p }); - IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), bnname, 1e-4); + conv1->setStrideNd(DimsHW{s, s}); + conv1->setPaddingNd(DimsHW{p, p}); + debug_print(conv1->getOutput(0), convname); + IScaleLayer *bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), bnname, 1e-5); + debug_print(bn1->getOutput(0), bnname); if (relu) { auto lr = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU); @@ -183,140 +209,146 @@ ILayer* convBnLeaky(INetworkDefinition *network, return bn1; } -IActivationLayer* ResBlock2Conv(INetworkDefinition *network, std::map& weightMap, ITensor& input, int inch, int outch, int stride, std::string lname) { - Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; - IConvolutionLayer* conv1 = network->addConvolutionNd(input, inch, DimsHW{ 1, 1 }, weightMap[lname + ".conv1.weight"], emptywts); +IActivationLayer *ResBlock2Conv(INetworkDefinition *network, std::map &weightMap, ITensor &input, int inch, int outch, int stride, std::string lname) +{ + Weights emptywts{DataType::kFLOAT, nullptr, 0}; + IConvolutionLayer *conv1 = network->addConvolutionNd(input, inch, DimsHW{1, 1}, weightMap[lname + ".conv1.weight"], emptywts); assert(conv1); - conv1->setStrideNd(DimsHW{ stride, stride }); - conv1->setPaddingNd(DimsHW{ 0, 0 }); - - IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn1", 1e-5); - IActivationLayer* relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU); + conv1->setStrideNd(DimsHW{stride, stride}); + conv1->setPaddingNd(DimsHW{0, 0}); + debug_print(conv1->getOutput(0), lname + "_1"); + IScaleLayer *bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn1", 1e-5); + IActivationLayer *relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU); assert(relu1); /// - IConvolutionLayer* conv2 = network->addConvolutionNd(*relu1->getOutput(0), inch, DimsHW{ 3, 3 }, weightMap[lname + ".conv2.weight"], emptywts); + IConvolutionLayer *conv2 = network->addConvolutionNd(*relu1->getOutput(0), inch, DimsHW{3, 3}, weightMap[lname + ".conv2.weight"], emptywts); assert(conv2); - conv2->setStrideNd(DimsHW{ stride, stride }); - conv2->setPaddingNd(DimsHW{ 1, 1 }); + conv2->setStrideNd(DimsHW{stride, stride}); + conv2->setPaddingNd(DimsHW{1, 1}); + debug_print(conv2->getOutput(0), lname + "_2"); + IScaleLayer *bn2 = addBatchNorm2d(network, weightMap, *conv2->getOutput(0), lname + ".bn2", 1e-5); - IScaleLayer* bn2 = addBatchNorm2d(network, weightMap, *conv2->getOutput(0), lname + ".bn2", 1e-5); - - IActivationLayer* relu2 = network->addActivation(*bn2->getOutput(0), ActivationType::kRELU); + IActivationLayer *relu2 = network->addActivation(*bn2->getOutput(0), ActivationType::kRELU); assert(relu2); ////// - IConvolutionLayer* conv3 = network->addConvolutionNd(*relu2->getOutput(0), outch, DimsHW{ 1, 1 }, weightMap[lname + ".conv3.weight"], emptywts); + IConvolutionLayer *conv3 = network->addConvolutionNd(*relu2->getOutput(0), outch, DimsHW{1, 1}, weightMap[lname + ".conv3.weight"], emptywts); assert(conv3); - conv1->setStrideNd(DimsHW{ stride, stride }); - conv3->setPaddingNd(DimsHW{ 0, 0 }); + conv3->setStrideNd(DimsHW{stride, stride}); + conv3->setPaddingNd(DimsHW{0, 0}); + debug_print(conv3->getOutput(0), lname + "_3"); + IScaleLayer *bn3 = addBatchNorm2d(network, weightMap, *conv3->getOutput(0), lname + ".bn3", 1e-5); - IScaleLayer* bn3 = addBatchNorm2d(network, weightMap, *conv3->getOutput(0), lname + ".bn3", 1e-5); - - IElementWiseLayer* ew1; - if (inch != outch) { - IConvolutionLayer* conv4 = network->addConvolutionNd(input, outch, DimsHW{ 1, 1 }, weightMap[lname + ".downsample.0.weight"], emptywts); - assert(conv4); - conv4->setStrideNd(DimsHW{ stride, stride }); - conv4->setPaddingNd(DimsHW{ 0, 0 }); - IScaleLayer* bn4 = addBatchNorm2d(network, weightMap, *conv4->getOutput(0), lname + ".downsample.1", 1e-5); - ew1 = network->addElementWise(*bn4->getOutput(0), *bn3->getOutput(0), ElementWiseOperation::kSUM); - } - else { - ew1 = network->addElementWise(input, *bn3->getOutput(0), ElementWiseOperation::kSUM); - } - IActivationLayer* relu3 = network->addActivation(*ew1->getOutput(0), ActivationType::kRELU); - assert(relu3); - return relu3; -} - -IActivationLayer* ResBlock(INetworkDefinition *network, std::map& weightMap, ITensor& input, int inch, int outch, int stride, std::string lname) { - Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; - // in 256 out 64 - IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ 1, 1 }, weightMap[lname + ".conv1.weight"], emptywts); - assert(conv1); - conv1->setStrideNd(DimsHW{ stride, stride }); - conv1->setPaddingNd(DimsHW{ 0, 0 }); - - IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn1", 1e-5); - - IActivationLayer* relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU); - assert(relu1); - /// - IConvolutionLayer* conv2 = network->addConvolutionNd(*relu1->getOutput(0), outch, DimsHW{ 3, 3 }, weightMap[lname + ".conv2.weight"], emptywts); - assert(conv2); - conv2->setStrideNd(DimsHW{ stride, stride }); - conv2->setPaddingNd(DimsHW{ 1, 1 }); - - IScaleLayer* bn2 = addBatchNorm2d(network, weightMap, *conv2->getOutput(0), lname + ".bn2", 1e-5); - - IActivationLayer* relu2 = network->addActivation(*bn2->getOutput(0), ActivationType::kRELU); - assert(relu2); - ////// - IConvolutionLayer* conv3 = network->addConvolutionNd(*relu2->getOutput(0), inch, DimsHW{ 1, 1 }, weightMap[lname + ".conv3.weight"], emptywts); - assert(conv3); - conv1->setStrideNd(DimsHW{ stride, stride }); - conv1->setPaddingNd(DimsHW{ 0, 0 }); - - IScaleLayer* bn3 = addBatchNorm2d(network, weightMap, *conv3->getOutput(0), lname + ".bn3", 1e-5); - - IElementWiseLayer* ew1; - ew1 = network->addElementWise(input, *bn3->getOutput(0), ElementWiseOperation::kSUM); - IActivationLayer* relu3 = network->addActivation(*ew1->getOutput(0), ActivationType::kRELU); - assert(relu3); - return relu3; -} - -IActivationLayer* liteResBlock(INetworkDefinition *network, std::map& weightMap, ITensor& input, int outch, std::string lname) { - Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; - // in 256 out 64 - IConvolutionLayer* conv1 = network->addConvolutionNd(input, outch, DimsHW{ 3, 3 }, weightMap[lname + ".conv1.weight"], emptywts); - assert(conv1); - conv1->setStrideNd(DimsHW{ 1, 1 }); - conv1->setPaddingNd(DimsHW{ 1, 1 }); - - IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn1", 1e-5); - - IActivationLayer* relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU); - assert(relu1); - /// - IConvolutionLayer* conv2 = network->addConvolutionNd(*relu1->getOutput(0), outch, DimsHW{ 3, 3 }, weightMap[lname + ".conv2.weight"], emptywts); - assert(conv2); - conv2->setStrideNd(DimsHW{ 1, 1 }); - conv2->setPaddingNd(DimsHW{ 1, 1 }); - - IScaleLayer* bn2 = addBatchNorm2d(network, weightMap, *conv2->getOutput(0), lname + ".bn2", 1e-5); - - IElementWiseLayer* ew1; - ew1 = network->addElementWise(input, *bn2->getOutput(0), ElementWiseOperation::kSUM); - - IActivationLayer* relu3 = network->addActivation(*ew1->getOutput(0), ActivationType::kRELU); - assert(relu3); - return relu3; -} - - -ILayer* convBnAddLeaky(INetworkDefinition *network, std::map& weightMap, ITensor& input, ITensor& addinput, int outch, int ksize, int s, int p, std::string convname, std::string bnname, bool bias = false) { - Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; - IConvolutionLayer* conv1; - //Dims dim; - if (!bias) + IElementWiseLayer *ew1; + if (inch != outch) { - conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[convname + ".weight"], emptywts); + IConvolutionLayer *conv4 = network->addConvolutionNd(input, outch, DimsHW{1, 1}, weightMap[lname + ".downsample.0.weight"], emptywts); + assert(conv4); + conv4->setStrideNd(DimsHW{stride, stride}); + conv4->setPaddingNd(DimsHW{0, 0}); + debug_print(conv4->getOutput(0), lname + "_4"); + IScaleLayer *bn4 = addBatchNorm2d(network, weightMap, *conv4->getOutput(0), lname + ".downsample.1", 1e-5); + ew1 = network->addElementWise(*bn4->getOutput(0), *bn3->getOutput(0), ElementWiseOperation::kSUM); } else { - conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[convname + ".weight"], weightMap[convname + ".bias"]); + ew1 = network->addElementWise(input, *bn3->getOutput(0), ElementWiseOperation::kSUM); + } + IActivationLayer *relu3 = network->addActivation(*ew1->getOutput(0), ActivationType::kRELU); + assert(relu3); + return relu3; +} + +IActivationLayer *ResBlock(INetworkDefinition *network, std::map &weightMap, ITensor &input, int inch, int outch, int stride, std::string lname) +{ + Weights emptywts{DataType::kFLOAT, nullptr, 0}; + // in 256 out 64 + IConvolutionLayer *conv1 = network->addConvolutionNd(input, outch, DimsHW{1, 1}, weightMap[lname + ".conv1.weight"], emptywts); + assert(conv1); + conv1->setStrideNd(DimsHW{stride, stride}); + conv1->setPaddingNd(DimsHW{0, 0}); + debug_print(conv1->getOutput(0), lname + "_1"); + IScaleLayer *bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn1", 1e-5); + + IActivationLayer *relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU); + assert(relu1); + /// + IConvolutionLayer *conv2 = network->addConvolutionNd(*relu1->getOutput(0), outch, DimsHW{3, 3}, weightMap[lname + ".conv2.weight"], emptywts); + assert(conv2); + conv2->setStrideNd(DimsHW{stride, stride}); + conv2->setPaddingNd(DimsHW{1, 1}); + debug_print(conv2->getOutput(0), lname + "_2"); + IScaleLayer *bn2 = addBatchNorm2d(network, weightMap, *conv2->getOutput(0), lname + ".bn2", 1e-5); + + IActivationLayer *relu2 = network->addActivation(*bn2->getOutput(0), ActivationType::kRELU); + assert(relu2); + ////// + IConvolutionLayer *conv3 = network->addConvolutionNd(*relu2->getOutput(0), inch, DimsHW{1, 1}, weightMap[lname + ".conv3.weight"], emptywts); + assert(conv3); + conv3->setStrideNd(DimsHW{stride, stride}); + conv3->setPaddingNd(DimsHW{0, 0}); + debug_print(conv3->getOutput(0), lname + "_3"); + IScaleLayer *bn3 = addBatchNorm2d(network, weightMap, *conv3->getOutput(0), lname + ".bn3", 1e-5); + + IElementWiseLayer *ew1; + ew1 = network->addElementWise(input, *bn3->getOutput(0), ElementWiseOperation::kSUM); + IActivationLayer *relu3 = network->addActivation(*ew1->getOutput(0), ActivationType::kRELU); + assert(relu3); + return relu3; +} + +IActivationLayer *liteResBlock(INetworkDefinition *network, std::map &weightMap, ITensor &input, int outch, std::string lname) +{ + Weights emptywts{DataType::kFLOAT, nullptr, 0}; + // in 256 out 64 + IConvolutionLayer *conv1 = network->addConvolutionNd(input, outch, DimsHW{3, 3}, weightMap[lname + ".conv1.weight"], emptywts); + assert(conv1); + conv1->setStrideNd(DimsHW{1, 1}); + conv1->setPaddingNd(DimsHW{1, 1}); + debug_print(conv1->getOutput(0), lname + "_1"); + IScaleLayer *bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), lname + ".bn1", 1e-5); + + IActivationLayer *relu1 = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU); + assert(relu1); + /// + IConvolutionLayer *conv2 = network->addConvolutionNd(*relu1->getOutput(0), outch, DimsHW{3, 3}, weightMap[lname + ".conv2.weight"], emptywts); + assert(conv2); + conv2->setStrideNd(DimsHW{1, 1}); + conv2->setPaddingNd(DimsHW{1, 1}); + debug_print(conv2->getOutput(0), lname + "_2"); + IScaleLayer *bn2 = addBatchNorm2d(network, weightMap, *conv2->getOutput(0), lname + ".bn2", 1e-5); + + IElementWiseLayer *ew1; + ew1 = network->addElementWise(input, *bn2->getOutput(0), ElementWiseOperation::kSUM); + debug_print(ew1->getOutput(0), lname + "_add"); + IActivationLayer *relu3 = network->addActivation(*ew1->getOutput(0), ActivationType::kRELU); + assert(relu3); + return relu3; +} + +ILayer *convBnAddRelu(INetworkDefinition *network, std::map &weightMap, ITensor &input, ITensor &addinput, int outch, int ksize, int s, int p, std::string convname, std::string bnname, bool bias = false) +{ + Weights emptywts{DataType::kFLOAT, nullptr, 0}; + IConvolutionLayer *conv1; + //Dims dim; + if (!bias) + { + conv1 = network->addConvolutionNd(input, outch, DimsHW{ksize, ksize}, weightMap[convname + ".weight"], emptywts); + } + else + { + conv1 = network->addConvolutionNd(input, outch, DimsHW{ksize, ksize}, weightMap[convname + ".weight"], weightMap[convname + ".bias"]); } assert(conv1); - conv1->setStrideNd(DimsHW{ s, s }); - conv1->setPaddingNd(DimsHW{ p, p }); - IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), bnname, 1e-4); - IElementWiseLayer* add = network->addElementWise(*bn1->getOutput(0), addinput, ElementWiseOperation::kSUM); + conv1->setStrideNd(DimsHW{s, s}); + conv1->setPaddingNd(DimsHW{p, p}); + debug_print(conv1->getOutput(0), convname); + IScaleLayer *bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), bnname, 1e-5); auto lr = network->addActivation(*bn1->getOutput(0), ActivationType::kRELU); + debug_print(lr->getOutput(0), convname + "_add"); return lr; } - -ILayer* netAddUpsampleBi(INetworkDefinition *network, ITensor *input, Dims outdims) +ILayer *netAddUpsampleBi(INetworkDefinition *network, ITensor *input, Dims outdims) { // Bi + True IResizeLayer *upSample = network->addResize(*input); @@ -326,36 +358,41 @@ ILayer* netAddUpsampleBi(INetworkDefinition *network, ITensor *input, Dims outdi return upSample; } -IElementWiseLayer* convBnUpAdd(INetworkDefinition *network, - std::map& weightMap, - ITensor& input, ITensor& addinput, - int outch, int ksize, int s, int p, - std::string convname, - std::string bnname, bool upsample, bool bias = false) { - Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; - IConvolutionLayer* conv1; - nvinfer1::Dims inp = input.getDimensions(); - if (!bias) { - conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[convname + ".weight"], emptywts); +IElementWiseLayer *convBnUpAdd(INetworkDefinition *network, + std::map &weightMap, + ITensor &input, ITensor &addinput, + int outch, int ksize, int s, int p, + std::string convname, + std::string bnname, bool upsample, bool bias = false) +{ + Weights emptywts{DataType::kFLOAT, nullptr, 0}; + IConvolutionLayer *conv1; + if (!bias) + { + conv1 = network->addConvolutionNd(input, outch, DimsHW{ksize, ksize}, weightMap[convname + ".weight"], emptywts); } - else { - conv1 = network->addConvolutionNd(input, outch, DimsHW{ ksize, ksize }, weightMap[convname + ".weight"], weightMap[convname + ".bias"]); + else + { + conv1 = network->addConvolutionNd(input, outch, DimsHW{ksize, ksize}, weightMap[convname + ".weight"], weightMap[convname + ".bias"]); } assert(conv1); - conv1->setStrideNd(DimsHW{ s, s }); - conv1->setPaddingNd(DimsHW{ p, p }); - nvinfer1::Dims inpDims = conv1->getOutput(0)->getDimensions(); - IScaleLayer* bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), bnname, 1e-5); - if (!upsample) { - IElementWiseLayer* add = network->addElementWise(*bn1->getOutput(0), addinput, ElementWiseOperation::kSUM); + conv1->setStrideNd(DimsHW{s, s}); + conv1->setPaddingNd(DimsHW{p, p}); + debug_print(conv1->getOutput(0), convname + "_1"); + IScaleLayer *bn1 = addBatchNorm2d(network, weightMap, *conv1->getOutput(0), bnname, 1e-5); + if (!upsample) + { + IElementWiseLayer *add = network->addElementWise(*bn1->getOutput(0), addinput, ElementWiseOperation::kSUM); + debug_print(add->getOutput(0), convname + "_add"); return add; } - else { + else + { nvinfer1::Dims dim = addinput.getDimensions(); - ILayer* up = netAddUpsampleBi(network, bn1->getOutput(0), dim); - IElementWiseLayer* add = network->addElementWise(*up->getOutput(0), addinput, ElementWiseOperation::kSUM); + ILayer *up = netAddUpsampleBi(network, bn1->getOutput(0), dim); + IElementWiseLayer *add = network->addElementWise(*up->getOutput(0), addinput, ElementWiseOperation::kSUM); + debug_print(conv1->getOutput(0), convname + "_1"); //auto lr = network->addActivation(*add->getOutput(0), ActivationType::kRELU); return add; } } - diff --git a/hrnet/hrnet-semantic-segmentation/demo.py b/hrnet/hrnet-semantic-segmentation/demo.py deleted file mode 100644 index 20b13cf..0000000 --- a/hrnet/hrnet-semantic-segmentation/demo.py +++ /dev/null @@ -1,175 +0,0 @@ -from __future__ import absolute_import -from __future__ import division -from __future__ import print_function - -import argparse -import os -import sys -import shutil -import pprint - -import torch -import torch.nn as nn -from torch.nn import functional as F -import torch.nn.parallel -import torch.backends.cudnn as cudnn -import torch.optim -import torch.utils.data -import torch.utils.data.distributed -import torchvision.datasets as datasets -import torchvision.transforms as transforms - -import _init_paths -import models -from config import config -from config import update_config -from utils.modelsummary import get_model_summary -from utils.utils import create_logger - -import cv2 -import numpy as np -from PIL import Image -import struct -import time - -import torch.onnx as torch_onnx - -ignore_label = -1 -label_mapping = {-1: ignore_label, 0: ignore_label, - 1: ignore_label, 2: ignore_label, - 3: ignore_label, 4: ignore_label, - 5: ignore_label, 6: ignore_label, - 7: 0, 8: 1, 9: ignore_label, - 10: ignore_label, 11: 2, 12: 3, - 13: 4, 14: ignore_label, 15: ignore_label, - 16: ignore_label, 17: 5, 18: ignore_label, - 19: 6, 20: 7, 21: 8, 22: 9, 23: 10, 24: 11, - 25: 12, 26: 13, 27: 14, 28: 15, - 29: ignore_label, 30: ignore_label, - 31: 16, 32: 17, 33: 18} - -def convert_label(label, inverse=False): - temp = label.copy() - if inverse: - for v, k in label_mapping.items(): - label[temp == k] = v - else: - for k, v in label_mapping.items(): - label[temp == k] = v - return label - -def get_palette(n): - palette = [0] * (n * 3) - for j in range(0, n): - lab = j - palette[j * 3 + 0] = 0 - palette[j * 3 + 1] = 0 - palette[j * 3 + 2] = 0 - i = 0 - while lab: - palette[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i)) - palette[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i)) - palette[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i)) - i += 1 - lab >>= 3 - return palette - -def save_pred(preds, sv_path, name): - palette = get_palette(256) - #preds = preds.cpu().numpy().copy() - preds = preds.detach().cpu().numpy().copy() - # numpy.argmax(array, axis) 用于返回一个numpy数组中最大值的索引值 - # (1 19 512 1024) 在19维度上输出每个维度最大值的索引 - # np.argmax(preds, axis=1) 返回 （1， 512， 1024） - preds = np.asarray(np.argmax(preds, axis=1), dtype=np.uint8) # argmax - for i in range(preds.shape[0]): - pred = convert_label(preds[i], inverse=True) - save_img = Image.fromarray(pred) - save_img.putpalette(palette) - #save_img.save(os.path.join(sv_path, name[i]+'.png')) - save_img.show() -def parse_args(): - parser = argparse.ArgumentParser(description='Train keypoints network') - - parser.add_argument('--cfg', - help='experiment configure file name', - default=r"../experiments/cityscapes/seg_hrnet_w18_small_v2_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml", - type=str) - parser.add_argument('opts', - help="Modify config options using the command-line", - default=None, - nargs=argparse.REMAINDER) - - args = parser.parse_args() - update_config(config, args) - - return args - -def main(): - savewts = False - save_onnx = False - args = parse_args() - - logger, final_output_dir, _ = create_logger(config, args.cfg, 'demo') - - logger.info(pprint.pformat(args)) - logger.info(pprint.pformat(config)) - - # cudnn related setting - cudnn.benchmark = config.CUDNN.BENCHMARK - torch.backends.cudnn.deterministic = config.CUDNN.DETERMINISTIC - torch.backends.cudnn.enabled = config.CUDNN.ENABLED - - # eval() 函数用来执行一个字符串表达式，并返回表达式的值。 - model = eval('models.'+config.MODEL.NAME+'.get_seg_model')(config) - - model_state_file = "../pretrained_models/hrnet_w18_small_v2_cityscapes_cls19_1024x2048_trainset.pth" - logger.info('=> loading model from {}'.format(model_state_file)) - pretrained_dict = torch.load(model_state_file) - model_dict = model.state_dict() - pretrained_dict = {k[6:]: v for k, v in pretrained_dict.items() - if k[6:] in model_dict.keys()} - for k, _ in pretrained_dict.items(): - logger.info( - '=> loading {} from pretrained model'.format(k)) - model_dict.update(pretrained_dict) - model.load_state_dict(model_dict) - - if savewts: - logger.info('=> saving {} '.format('HRNetSeg.wts')) - f = open('HRNetSeg.wts', 'w') - f.write('{}\n'.format(len(model.state_dict().keys()))) - for k, v in model.state_dict().items(): - vr = v.reshape(-1).cpu().numpy() - f.write('{} {} '.format(k, len(vr))) - for vv in vr: - f.write(' ') - f.write(struct.pack('>f', float(vv)).hex()) - f.write('\n') - exit(0) - # load img - #testImg = "../data/cityscapes/images/bonn_000014_000019_leftImg8bit.png" - testImg = r"E:\Datasets\oneimg\munster_000000_000019_rightImg8bit.png" - image = cv2.imread(testImg) #BGR 0-255 hwc - resized_img = cv2.resize(image, (config.TRAIN.IMAGE_SIZE[0], config.TRAIN.IMAGE_SIZE[1])) - resized_img = cv2.cvtColor(resized_img, cv2.COLOR_BGR2RGB) #RGB - # normalize - mean = [0.485, 0.456, 0.406] - std = [0.229, 0.224, 0.225] - inp_image = ((resized_img/255. - mean) / std).astype(np.float32) - inp_image = inp_image.transpose(2, 0, 1) - inp_image = torch.from_numpy(inp_image).unsqueeze(0) # to_tensor - model.eval() - if save_onnx: - logger.info('=> saving {} '.format('hrnet.onnx')) - torch_onnx.export(model, inp_image, "hrnet.onnx", verbose=True, input_names=["input"], - output_names=["output"],opset_version=12) - start = time.time() - output = model(inp_image) - if output.size()[-2] != config.TRAIN.IMAGE_SIZE[0] or output.size()[-1] != config.TRAIN.IMAGE_SIZE[1]: - output = F.upsample(output, (config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]), mode='bilinear', align_corners=True) - end = time.time() - print(end-start) - save_pred(output, "./", "1") -if __name__ == "__main__": - main() diff --git a/hrnet/hrnet-semantic-segmentation/gen_wts.py b/hrnet/hrnet-semantic-segmentation/gen_wts.py new file mode 100644 index 0000000..f133a9c --- /dev/null +++ b/hrnet/hrnet-semantic-segmentation/gen_wts.py @@ -0,0 +1,60 @@ +import argparse +import struct + +import _init_paths +import models +import torch +from config import config, update_config + + +def parse_args(): + parser = argparse.ArgumentParser(description="Train keypoints network") + + parser.add_argument("--cfg", help="experiment configure file name", type=str) + parser.add_argument("--ckpt_path", help="checkpoint path", required=True, type=str) + parser.add_argument("--save_path", help=".wts path", required=True, type=str) + + parser.add_argument( + "opts", + help="Modify config options using the command-line", + default=None, + nargs=argparse.REMAINDER, + ) + + args = parser.parse_args() + update_config(config, args) + + return args + + +def main(): + args = parse_args() + + model = eval("models." + config.MODEL.NAME + ".get_seg_model")(config) + + print("=> loading model from {}".format(args.ckpt_path)) + pretrained_dict = torch.load(args.ckpt_path, map_location="cpu") + model_dict = model.state_dict() + pretrained_dict = { + k[6:]: v for k, v in pretrained_dict.items() if k[6:] in model_dict.keys() + } + for k, _ in pretrained_dict.items(): + print("=> loading {} from pretrained model".format(k)) + model_dict.update(pretrained_dict) + model.load_state_dict(model_dict) + + print("=> saving {} ".format(args.save_path)) + f = open(args.save_path, "w") + f.write("{}\n".format(len(model.state_dict().keys()))) + for k, v in model.state_dict().items(): + vr = v.reshape(-1).cpu().numpy() + f.write("{} {} ".format(k, len(vr))) + for vv in vr: + f.write(" ") + f.write(struct.pack(">f", float(vv)).hex()) + f.write("\n") + f.close() + + +if __name__ == "__main__": + main() diff --git a/hrnet/hrnet-semantic-segmentation/hrnet.cpp b/hrnet/hrnet-semantic-segmentation/hrnet.cpp new file mode 100644 index 0000000..35a94a4 --- /dev/null +++ b/hrnet/hrnet-semantic-segmentation/hrnet.cpp @@ -0,0 +1,569 @@ +#include +#include +#include +#include +#include +#include +#include "common.hpp" +#include "logging.h" + +static Logger gLogger; +#define USE_FP32 +#define DEVICE 0 // GPU id +#define BATCH_SIZE 1 + +const char *INPUT_BLOB_NAME = "data"; +const char *OUTPUT_BLOB_NAME = "output"; +static const int INPUT_H = 512; +static const int INPUT_W = 1024; +static const int NUM_CLASSES = 19; +static const int OUTPUT_SIZE = INPUT_H * INPUT_W; + +// Creat the engine using only the API and not any parser. +ICudaEngine *createEngine(unsigned int maxBatchSize, IBuilder *builder, IBuilderConfig *config, DataType dt, std::string wtsPath, int width) +{ + 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{INPUT_H, INPUT_W, 3}); + assert(data); + + // hwc to chw + auto ps = network->addShuffle(*data); + ps->setFirstTranspose(nvinfer1::Permutation{2, 0, 1}); + float mean[3] = {0.485, 0.456, 0.406}; + float std[3] = {0.229, 0.224, 0.225}; + ITensor *preinput = MeanStd(network, ps->getOutput(0), mean, std, true); + + std::map weightMap = loadWeights(wtsPath); + auto relu_2 = convBnRelu(network, weightMap, *preinput, 64, 3, 2, 1, "conv1", "bn1"); + auto relu_5 = convBnRelu(network, weightMap, *relu_2->getOutput(0), 64, 3, 2, 1, "conv2", "bn2"); + auto relu_17 = ResBlock2Conv(network, weightMap, *relu_5->getOutput(0), 64, 256, 1, "layer1.0"); + auto relu_27 = ResBlock(network, weightMap, *relu_17->getOutput(0), 256, 64, 1, "layer1.1"); + auto relu_37 = ResBlock(network, weightMap, *relu_27->getOutput(0), 256, 64, 1, "layer1.2"); + auto relu_47 = ResBlock(network, weightMap, *relu_37->getOutput(0), 256, 64, 1, "layer1.3"); + + auto relu_50 = convBnRelu(network, weightMap, *relu_47->getOutput(0), width, 3, 1, 1, "transition1.0.0", "transition1.0.1"); + auto relu_60 = liteResBlock(network, weightMap, *relu_50->getOutput(0), width, "stage2.0.branches.0.0"); + auto relu_67 = liteResBlock(network, weightMap, *relu_60->getOutput(0), width, "stage2.0.branches.0.1"); + auto relu_74 = liteResBlock(network, weightMap, *relu_67->getOutput(0), width, "stage2.0.branches.0.2"); + auto relu_81 = liteResBlock(network, weightMap, *relu_74->getOutput(0), width, "stage2.0.branches.0.3"); + + auto relu_53 = convBnRelu(network, weightMap, *relu_47->getOutput(0), width * 2, 3, 2, 1, "transition1.1.0.0", "transition1.1.0.1"); + auto relu_88 = liteResBlock(network, weightMap, *relu_53->getOutput(0), width * 2, "stage2.0.branches.1.0"); + auto relu_95 = liteResBlock(network, weightMap, *relu_88->getOutput(0), width * 2, "stage2.0.branches.1.1"); + auto relu_102 = liteResBlock(network, weightMap, *relu_95->getOutput(0), width * 2, "stage2.0.branches.1.2"); + auto relu_109 = liteResBlock(network, weightMap, *relu_102->getOutput(0), width * 2, "stage2.0.branches.1.3"); + + auto add_131 = convBnUpAdd(network, weightMap, *relu_109->getOutput(0), *relu_81->getOutput(0), width, 1, 1, 0, "stage2.0.fuse_layers.0.1.0", "stage2.0.fuse_layers.0.1.1", true); + auto relu_132 = network->addActivation(*add_131->getOutput(0), ActivationType::kRELU); + + auto add_135 = convBnUpAdd(network, weightMap, *relu_81->getOutput(0), *relu_109->getOutput(0), width * 2, 3, 2, 1, "stage2.0.fuse_layers.1.0.0.0", "stage2.0.fuse_layers.1.0.0.1", false); + auto relu_136 = network->addActivation(*add_135->getOutput(0), ActivationType::kRELU); + + auto relu_146 = liteResBlock(network, weightMap, *relu_132->getOutput(0), width, "stage3.0.branches.0.0"); + auto relu_153 = liteResBlock(network, weightMap, *relu_146->getOutput(0), width, "stage3.0.branches.0.1"); + auto relu_160 = liteResBlock(network, weightMap, *relu_153->getOutput(0), width, "stage3.0.branches.0.2"); + auto relu_167 = liteResBlock(network, weightMap, *relu_160->getOutput(0), width, "stage3.0.branches.0.3"); + + auto relu_174 = liteResBlock(network, weightMap, *relu_136->getOutput(0), width * 2, "stage3.0.branches.1.0"); + auto relu_181 = liteResBlock(network, weightMap, *relu_174->getOutput(0), width * 2, "stage3.0.branches.1.1"); + auto relu_188 = liteResBlock(network, weightMap, *relu_181->getOutput(0), width * 2, "stage3.0.branches.1.2"); + auto relu_195 = liteResBlock(network, weightMap, *relu_188->getOutput(0), width * 2, "stage3.0.branches.1.3"); + + auto relu_139 = convBnRelu(network, weightMap, *relu_136->getOutput(0), width * 4, 3, 2, 1, "transition2.2.0.0", "transition2.2.0.1"); + auto relu_202 = liteResBlock(network, weightMap, *relu_139->getOutput(0), width * 4, "stage3.0.branches.2.0"); + auto relu_209 = liteResBlock(network, weightMap, *relu_202->getOutput(0), width * 4, "stage3.0.branches.2.1"); + auto relu_216 = liteResBlock(network, weightMap, *relu_209->getOutput(0), width * 4, "stage3.0.branches.2.2"); + auto relu_223 = liteResBlock(network, weightMap, *relu_216->getOutput(0), width * 4, "stage3.0.branches.2.3"); + + auto add_245 = convBnUpAdd(network, weightMap, *relu_195->getOutput(0), *relu_167->getOutput(0), width, 1, 1, 0, "stage3.0.fuse_layers.0.1.0", "stage3.0.fuse_layers.0.1.1", true); + auto add_267 = convBnUpAdd(network, weightMap, *relu_223->getOutput(0), *add_245->getOutput(0), width, 1, 1, 0, "stage3.0.fuse_layers.0.2.0", "stage3.0.fuse_layers.0.2.1", true); + auto relu_268 = network->addActivation(*add_267->getOutput(0), ActivationType::kRELU); + + auto add_271 = convBnUpAdd(network, weightMap, *relu_167->getOutput(0), *relu_195->getOutput(0), width * 2, 3, 2, 1, "stage3.0.fuse_layers.1.0.0.0", "stage3.0.fuse_layers.1.0.0.1", false); + auto add_293 = convBnUpAdd(network, weightMap, *relu_223->getOutput(0), *add_271->getOutput(0), width * 2, 1, 1, 0, "stage3.0.fuse_layers.1.2.0", "stage3.0.fuse_layers.1.2.1", true); + auto relu_294 = network->addActivation(*add_293->getOutput(0), ActivationType::kRELU); + + auto relu_297 = convBnRelu(network, weightMap, *relu_167->getOutput(0), width, 3, 2, 1, "stage3.0.fuse_layers.2.0.0.0", "stage3.0.fuse_layers.2.0.0.1"); + auto bn_299 = convBnRelu(network, weightMap, *relu_297->getOutput(0), width * 4, 3, 2, 1, "stage3.0.fuse_layers.2.0.1.0", "stage3.0.fuse_layers.2.0.1.1", false); + auto add_302 = convBnUpAdd(network, weightMap, *relu_195->getOutput(0), *bn_299->getOutput(0), width * 4, 3, 2, 1, "stage3.0.fuse_layers.2.1.0.0", "stage3.0.fuse_layers.2.1.0.1", false); + auto add_303 = network->addElementWise(*add_302->getOutput(0), *relu_223->getOutput(0), ElementWiseOperation::kSUM); + auto relu_304 = network->addActivation(*add_303->getOutput(0), ActivationType::kRELU); + + auto relu_311 = liteResBlock(network, weightMap, *relu_268->getOutput(0), width, "stage3.1.branches.0.0"); + auto relu_318 = liteResBlock(network, weightMap, *relu_311->getOutput(0), width, "stage3.1.branches.0.1"); + auto relu_325 = liteResBlock(network, weightMap, *relu_318->getOutput(0), width, "stage3.1.branches.0.2"); + auto relu_332 = liteResBlock(network, weightMap, *relu_325->getOutput(0), width, "stage3.1.branches.0.3"); + + auto relu_339 = liteResBlock(network, weightMap, *relu_294->getOutput(0), width * 2, "stage3.1.branches.1.0"); + auto relu_346 = liteResBlock(network, weightMap, *relu_339->getOutput(0), width * 2, "stage3.1.branches.1.1"); + auto relu_353 = liteResBlock(network, weightMap, *relu_346->getOutput(0), width * 2, "stage3.1.branches.1.2"); + auto relu_360 = liteResBlock(network, weightMap, *relu_353->getOutput(0), width * 2, "stage3.1.branches.1.3"); + + auto relu_367 = liteResBlock(network, weightMap, *relu_304->getOutput(0), width * 4, "stage3.1.branches.2.0"); + auto relu_374 = liteResBlock(network, weightMap, *relu_367->getOutput(0), width * 4, "stage3.1.branches.2.1"); + auto relu_381 = liteResBlock(network, weightMap, *relu_374->getOutput(0), width * 4, "stage3.1.branches.2.2"); + auto relu_388 = liteResBlock(network, weightMap, *relu_381->getOutput(0), width * 4, "stage3.1.branches.2.3"); + + auto add_410 = convBnUpAdd(network, weightMap, *relu_360->getOutput(0), *relu_332->getOutput(0), width, 1, 1, 0, "stage3.1.fuse_layers.0.1.0", "stage3.1.fuse_layers.0.1.1", true); + auto add_432 = convBnUpAdd(network, weightMap, *relu_388->getOutput(0), *add_410->getOutput(0), width, 1, 1, 0, "stage3.1.fuse_layers.0.2.0", "stage3.1.fuse_layers.0.2.1", true); + auto relu_433 = network->addActivation(*add_432->getOutput(0), ActivationType::kRELU); + + auto add_436 = convBnUpAdd(network, weightMap, *relu_332->getOutput(0), *relu_360->getOutput(0), width * 2, 3, 2, 1, "stage3.1.fuse_layers.1.0.0.0", "stage3.1.fuse_layers.1.0.0.1", false); + auto add_458 = convBnUpAdd(network, weightMap, *relu_388->getOutput(0), *add_436->getOutput(0), width * 2, 1, 1, 0, "stage3.1.fuse_layers.1.2.0", "stage3.1.fuse_layers.1.2.1", true); + auto relu_459 = network->addActivation(*add_458->getOutput(0), ActivationType::kRELU); + + auto relu_462 = convBnRelu(network, weightMap, *relu_332->getOutput(0), width, 3, 2, 1, "stage3.1.fuse_layers.2.0.0.0", "stage3.1.fuse_layers.2.0.0.1"); + auto bn_464 = convBnRelu(network, weightMap, *relu_462->getOutput(0), width * 4, 3, 2, 1, "stage3.1.fuse_layers.2.0.1.0", "stage3.1.fuse_layers.2.0.1.1", false); + auto add_467 = convBnUpAdd(network, weightMap, *relu_360->getOutput(0), *bn_464->getOutput(0), width * 4, 3, 2, 1, "stage3.1.fuse_layers.2.1.0.0", "stage3.1.fuse_layers.2.1.0.1", false); + auto add_468 = network->addElementWise(*add_467->getOutput(0), *relu_388->getOutput(0), ElementWiseOperation::kSUM); + auto relu_469 = network->addActivation(*add_468->getOutput(0), ActivationType::kRELU); + + auto relu_476 = liteResBlock(network, weightMap, *relu_433->getOutput(0), width, "stage3.2.branches.0.0"); + auto relu_483 = liteResBlock(network, weightMap, *relu_476->getOutput(0), width, "stage3.2.branches.0.1"); + auto relu_490 = liteResBlock(network, weightMap, *relu_483->getOutput(0), width, "stage3.2.branches.0.2"); + auto relu_497 = liteResBlock(network, weightMap, *relu_490->getOutput(0), width, "stage3.2.branches.0.3"); + + auto relu_504 = liteResBlock(network, weightMap, *relu_459->getOutput(0), width * 2, "stage3.2.branches.1.0"); + auto relu_511 = liteResBlock(network, weightMap, *relu_504->getOutput(0), width * 2, "stage3.2.branches.1.1"); + auto relu_518 = liteResBlock(network, weightMap, *relu_511->getOutput(0), width * 2, "stage3.2.branches.1.2"); + auto relu_525 = liteResBlock(network, weightMap, *relu_518->getOutput(0), width * 2, "stage3.2.branches.1.3"); + + auto relu_532 = liteResBlock(network, weightMap, *relu_469->getOutput(0), width * 4, "stage3.2.branches.2.0"); + auto relu_539 = liteResBlock(network, weightMap, *relu_532->getOutput(0), width * 4, "stage3.2.branches.2.1"); + auto relu_546 = liteResBlock(network, weightMap, *relu_539->getOutput(0), width * 4, "stage3.2.branches.2.2"); + auto relu_553 = liteResBlock(network, weightMap, *relu_546->getOutput(0), width * 4, "stage3.2.branches.2.3"); + + auto add_575 = convBnUpAdd(network, weightMap, *relu_525->getOutput(0), *relu_497->getOutput(0), width, 1, 1, 0, "stage3.2.fuse_layers.0.1.0", "stage3.2.fuse_layers.0.1.1", true); + auto add_597 = convBnUpAdd(network, weightMap, *relu_553->getOutput(0), *add_575->getOutput(0), width, 1, 1, 0, "stage3.2.fuse_layers.0.2.0", "stage3.2.fuse_layers.0.2.1", true); + + auto relu_598 = network->addActivation(*add_597->getOutput(0), ActivationType::kRELU); + + auto add_601 = convBnUpAdd(network, weightMap, *relu_497->getOutput(0), *relu_525->getOutput(0), width * 2, 3, 2, 1, "stage3.2.fuse_layers.1.0.0.0", "stage3.2.fuse_layers.1.0.0.1", false); + auto add_623 = convBnUpAdd(network, weightMap, *relu_553->getOutput(0), *add_601->getOutput(0), width * 2, 1, 1, 0, "stage3.2.fuse_layers.1.2.0", "stage3.2.fuse_layers.1.2.1", true); + auto relu_624 = network->addActivation(*add_623->getOutput(0), ActivationType::kRELU); + + auto relu_627 = convBnRelu(network, weightMap, *relu_497->getOutput(0), width, 3, 2, 1, "stage3.2.fuse_layers.2.0.0.0", "stage3.2.fuse_layers.2.0.0.1"); + auto bn_629 = convBnRelu(network, weightMap, *relu_627->getOutput(0), width * 4, 3, 2, 1, "stage3.2.fuse_layers.2.0.1.0", "stage3.2.fuse_layers.2.0.1.1", false); + auto add_632 = convBnUpAdd(network, weightMap, *relu_525->getOutput(0), *bn_629->getOutput(0), width * 4, 3, 2, 1, "stage3.2.fuse_layers.2.1.0.0", "stage3.2.fuse_layers.2.1.0.1", false); + auto add_633 = network->addElementWise(*relu_553->getOutput(0), *add_632->getOutput(0), ElementWiseOperation::kSUM); + auto relu_634 = network->addActivation(*add_633->getOutput(0), ActivationType::kRELU); + + auto relu_641 = liteResBlock(network, weightMap, *relu_598->getOutput(0), width, "stage3.3.branches.0.0"); + auto relu_648 = liteResBlock(network, weightMap, *relu_641->getOutput(0), width, "stage3.3.branches.0.1"); + auto relu_655 = liteResBlock(network, weightMap, *relu_648->getOutput(0), width, "stage3.3.branches.0.2"); + auto relu_662 = liteResBlock(network, weightMap, *relu_655->getOutput(0), width, "stage3.3.branches.0.3"); + + auto relu_669 = liteResBlock(network, weightMap, *relu_624->getOutput(0), width * 2, "stage3.3.branches.1.0"); + auto relu_676 = liteResBlock(network, weightMap, *relu_669->getOutput(0), width * 2, "stage3.3.branches.1.1"); + auto relu_683 = liteResBlock(network, weightMap, *relu_676->getOutput(0), width * 2, "stage3.3.branches.1.2"); + auto relu_690 = liteResBlock(network, weightMap, *relu_683->getOutput(0), width * 2, "stage3.3.branches.1.3"); + + auto relu_697 = liteResBlock(network, weightMap, *relu_634->getOutput(0), width * 4, "stage3.3.branches.2.0"); + auto relu_704 = liteResBlock(network, weightMap, *relu_697->getOutput(0), width * 4, "stage3.3.branches.2.1"); + auto relu_711 = liteResBlock(network, weightMap, *relu_704->getOutput(0), width * 4, "stage3.3.branches.2.2"); + auto relu_718 = liteResBlock(network, weightMap, *relu_711->getOutput(0), width * 4, "stage3.3.branches.2.3"); + + auto add_740 = convBnUpAdd(network, weightMap, *relu_690->getOutput(0), *relu_662->getOutput(0), width, 1, 1, 0, "stage3.3.fuse_layers.0.1.0", "stage3.3.fuse_layers.0.1.1", true); + auto add_762 = convBnUpAdd(network, weightMap, *relu_718->getOutput(0), *add_740->getOutput(0), width, 1, 1, 0, "stage3.3.fuse_layers.0.2.0", "stage3.3.fuse_layers.0.2.1", true); + auto relu_763 = network->addActivation(*add_762->getOutput(0), ActivationType::kRELU); + + auto add_766 = convBnUpAdd(network, weightMap, *relu_662->getOutput(0), *relu_690->getOutput(0), width * 2, 3, 2, 1, "stage3.3.fuse_layers.1.0.0.0", "stage3.3.fuse_layers.1.0.0.1", false); + auto add_788 = convBnUpAdd(network, weightMap, *relu_718->getOutput(0), *add_766->getOutput(0), width * 2, 1, 1, 0, "stage3.3.fuse_layers.1.2.0", "stage3.3.fuse_layers.1.2.1", true); + auto relu_789 = network->addActivation(*add_788->getOutput(0), ActivationType::kRELU); + + auto relu_792 = convBnRelu(network, weightMap, *relu_662->getOutput(0), width, 3, 2, 1, "stage3.3.fuse_layers.2.0.0.0", "stage3.3.fuse_layers.2.0.0.1"); + auto bn_794 = convBnRelu(network, weightMap, *relu_792->getOutput(0), width * 4, 3, 2, 1, "stage3.3.fuse_layers.2.0.1.0", "stage3.3.fuse_layers.2.0.1.1", false); + auto add_797 = convBnUpAdd(network, weightMap, *relu_690->getOutput(0), *bn_794->getOutput(0), width * 4, 3, 2, 1, "stage3.3.fuse_layers.2.1.0.0", "stage3.3.fuse_layers.2.1.0.1", false); + auto add_798 = network->addElementWise(*relu_718->getOutput(0), *add_797->getOutput(0), ElementWiseOperation::kSUM); + auto relu_799 = network->addActivation(*add_798->getOutput(0), ActivationType::kRELU); + + auto relu_809 = liteResBlock(network, weightMap, *relu_763->getOutput(0), width, "stage4.0.branches.0.0"); + auto relu_816 = liteResBlock(network, weightMap, *relu_809->getOutput(0), width, "stage4.0.branches.0.1"); + auto relu_823 = liteResBlock(network, weightMap, *relu_816->getOutput(0), width, "stage4.0.branches.0.2"); + auto relu_830 = liteResBlock(network, weightMap, *relu_823->getOutput(0), width, "stage4.0.branches.0.3"); + + auto relu_837 = liteResBlock(network, weightMap, *relu_789->getOutput(0), width * 2, "stage4.0.branches.1.0"); + auto relu_844 = liteResBlock(network, weightMap, *relu_837->getOutput(0), width * 2, "stage4.0.branches.1.1"); + auto relu_851 = liteResBlock(network, weightMap, *relu_844->getOutput(0), width * 2, "stage4.0.branches.1.2"); + auto relu_858 = liteResBlock(network, weightMap, *relu_851->getOutput(0), width * 2, "stage4.0.branches.1.3"); + + auto relu_865 = liteResBlock(network, weightMap, *relu_799->getOutput(0), width * 4, "stage4.0.branches.2.0"); + auto relu_872 = liteResBlock(network, weightMap, *relu_865->getOutput(0), width * 4, "stage4.0.branches.2.1"); + auto relu_879 = liteResBlock(network, weightMap, *relu_872->getOutput(0), width * 4, "stage4.0.branches.2.2"); + auto relu_886 = liteResBlock(network, weightMap, *relu_879->getOutput(0), width * 4, "stage4.0.branches.2.3"); //======== + + auto relu_802 = convBnRelu(network, weightMap, *relu_799->getOutput(0), width * 8, 3, 2, 1, "transition3.3.0.0", "transition3.3.0.1"); + auto relu_893 = liteResBlock(network, weightMap, *relu_802->getOutput(0), width * 8, "stage4.0.branches.3.0"); + auto relu_900 = liteResBlock(network, weightMap, *relu_893->getOutput(0), width * 8, "stage4.0.branches.3.1"); + auto relu_907 = liteResBlock(network, weightMap, *relu_900->getOutput(0), width * 8, "stage4.0.branches.3.2"); + auto relu_914 = liteResBlock(network, weightMap, *relu_907->getOutput(0), width * 8, "stage4.0.branches.3.3"); + + auto add_936 = convBnUpAdd(network, weightMap, *relu_858->getOutput(0), *relu_830->getOutput(0), width, 1, 1, 0, "stage4.0.fuse_layers.0.1.0", "stage4.0.fuse_layers.0.1.1", true); + auto add_958 = convBnUpAdd(network, weightMap, *relu_886->getOutput(0), *add_936->getOutput(0), width, 1, 1, 0, "stage4.0.fuse_layers.0.2.0", "stage4.0.fuse_layers.0.2.1", true); + auto add_980 = convBnUpAdd(network, weightMap, *relu_914->getOutput(0), *add_958->getOutput(0), width, 1, 1, 0, "stage4.0.fuse_layers.0.3.0", "stage4.0.fuse_layers.0.3.1", true); + auto relu_981 = network->addActivation(*add_980->getOutput(0), ActivationType::kRELU); + + auto add_984 = convBnUpAdd(network, weightMap, *relu_830->getOutput(0), *relu_858->getOutput(0), width * 2, 3, 2, 1, "stage4.0.fuse_layers.1.0.0.0", "stage4.0.fuse_layers.1.0.0.1", false); + auto add_1006 = convBnUpAdd(network, weightMap, *relu_886->getOutput(0), *add_984->getOutput(0), width * 2, 1, 1, 0, "stage4.0.fuse_layers.1.2.0", "stage4.0.fuse_layers.1.2.1", true); + auto add_1028 = convBnUpAdd(network, weightMap, *relu_914->getOutput(0), *add_1006->getOutput(0), width * 2, 1, 1, 0, "stage4.0.fuse_layers.1.3.0", "stage4.0.fuse_layers.1.3.1", true); + auto relu_1029 = network->addActivation(*add_1028->getOutput(0), ActivationType::kRELU); + + auto relu_1032 = convBnRelu(network, weightMap, *relu_830->getOutput(0), width, 3, 2, 1, "stage4.0.fuse_layers.2.0.0.0", "stage4.0.fuse_layers.2.0.0.1"); + auto bn_1034 = convBnRelu(network, weightMap, *relu_1032->getOutput(0), width * 4, 3, 2, 1, "stage4.0.fuse_layers.2.0.1.0", "stage4.0.fuse_layers.2.0.1.1", false); + + auto add_1037 = convBnUpAdd(network, weightMap, *relu_858->getOutput(0), *bn_1034->getOutput(0), width * 4, 3, 2, 1, + "stage4.0.fuse_layers.2.1.0.0", "stage4.0.fuse_layers.2.1.0.1", false); + auto add_1038 = network->addElementWise(*relu_886->getOutput(0), *add_1037->getOutput(0), ElementWiseOperation::kSUM); + auto add_1060 = convBnUpAdd(network, weightMap, *relu_914->getOutput(0), *add_1038->getOutput(0), width * 4, 1, 1, 0, + "stage4.0.fuse_layers.2.3.0", "stage4.0.fuse_layers.2.3.1", true); + auto relu_1061 = network->addActivation(*add_1060->getOutput(0), ActivationType::kRELU); + + auto relu_1064 = convBnRelu(network, weightMap, *relu_830->getOutput(0), width, 3, 2, 1, "stage4.0.fuse_layers.3.0.0.0", "stage4.0.fuse_layers.3.0.0.1"); + auto relu_1067 = convBnRelu(network, weightMap, *relu_1064->getOutput(0), width, 3, 2, 1, "stage4.0.fuse_layers.3.0.1.0", "stage4.0.fuse_layers.3.0.1.1"); + auto bn_1069 = convBnRelu(network, weightMap, *relu_1067->getOutput(0), width * 8, 3, 2, 1, "stage4.0.fuse_layers.3.0.2.0", "stage4.0.fuse_layers.3.0.2.1", false); + auto relu_1072 = convBnRelu(network, weightMap, *relu_858->getOutput(0), width * 2, 3, 2, 1, "stage4.0.fuse_layers.3.1.0.0", "stage4.0.fuse_layers.3.1.0.1"); + auto add_1075 = convBnUpAdd(network, weightMap, *relu_1072->getOutput(0), *bn_1069->getOutput(0), width * 8, 3, 2, 1, + "stage4.0.fuse_layers.3.1.1.0", "stage4.0.fuse_layers.3.1.1.1", false); + auto add_1078 = convBnUpAdd(network, weightMap, *relu_886->getOutput(0), *add_1075->getOutput(0), width * 8, 3, 2, 1, + "stage4.0.fuse_layers.3.2.0.0", "stage4.0.fuse_layers.3.2.0.1", false); + auto add_1079 = network->addElementWise(*relu_914->getOutput(0), *add_1078->getOutput(0), ElementWiseOperation::kSUM); + auto relu_1080 = network->addActivation(*add_1079->getOutput(0), ActivationType::kRELU); + + auto relu_1087 = liteResBlock(network, weightMap, *relu_981->getOutput(0), width, "stage4.1.branches.0.0"); + auto relu_1094 = liteResBlock(network, weightMap, *relu_1087->getOutput(0), width, "stage4.1.branches.0.1"); + auto relu_1101 = liteResBlock(network, weightMap, *relu_1094->getOutput(0), width, "stage4.1.branches.0.2"); + auto relu_1108 = liteResBlock(network, weightMap, *relu_1101->getOutput(0), width, "stage4.1.branches.0.3"); + + auto relu_1115 = liteResBlock(network, weightMap, *relu_1029->getOutput(0), width * 2, "stage4.1.branches.1.0"); + auto relu_1122 = liteResBlock(network, weightMap, *relu_1115->getOutput(0), width * 2, "stage4.1.branches.1.1"); + auto relu_1129 = liteResBlock(network, weightMap, *relu_1122->getOutput(0), width * 2, "stage4.1.branches.1.2"); + auto relu_1136 = liteResBlock(network, weightMap, *relu_1129->getOutput(0), width * 2, "stage4.1.branches.1.3"); + + auto relu_1143 = liteResBlock(network, weightMap, *relu_1061->getOutput(0), width * 4, "stage4.1.branches.2.0"); + auto relu_1150 = liteResBlock(network, weightMap, *relu_1143->getOutput(0), width * 4, "stage4.1.branches.2.1"); + auto relu_1157 = liteResBlock(network, weightMap, *relu_1150->getOutput(0), width * 4, "stage4.1.branches.2.2"); + auto relu_1164 = liteResBlock(network, weightMap, *relu_1157->getOutput(0), width * 4, "stage4.1.branches.2.3"); + + auto relu_1171 = liteResBlock(network, weightMap, *relu_1080->getOutput(0), width * 8, "stage4.1.branches.3.0"); + auto relu_1178 = liteResBlock(network, weightMap, *relu_1171->getOutput(0), width * 8, "stage4.1.branches.3.1"); + auto relu_1185 = liteResBlock(network, weightMap, *relu_1178->getOutput(0), width * 8, "stage4.1.branches.3.2"); + auto relu_1192 = liteResBlock(network, weightMap, *relu_1185->getOutput(0), width * 8, "stage4.1.branches.3.3"); + + auto add_1214 = convBnUpAdd(network, weightMap, *relu_1136->getOutput(0), *relu_1108->getOutput(0), width, 1, 1, 0, + "stage4.1.fuse_layers.0.1.0", "stage4.1.fuse_layers.0.1.1", true); + auto add_1236 = convBnUpAdd(network, weightMap, *relu_1164->getOutput(0), *add_1214->getOutput(0), width, 1, 1, 0, + "stage4.1.fuse_layers.0.2.0", "stage4.1.fuse_layers.0.2.1", true); + auto add_1258 = convBnUpAdd(network, weightMap, *relu_1192->getOutput(0), *add_1236->getOutput(0), width, 1, 1, 0, + "stage4.1.fuse_layers.0.3.0", "stage4.1.fuse_layers.0.3.1", true); + auto relu_1259 = network->addActivation(*add_1258->getOutput(0), ActivationType::kRELU); + + auto add_1262 = convBnUpAdd(network, weightMap, *relu_1108->getOutput(0), *relu_1136->getOutput(0), width * 2, 3, 2, 1, + "stage4.1.fuse_layers.1.0.0.0", "stage4.1.fuse_layers.1.0.0.1", false); + auto add_1284 = convBnUpAdd(network, weightMap, *relu_1164->getOutput(0), *add_1262->getOutput(0), width * 2, 1, 1, 0, + "stage4.1.fuse_layers.1.2.0", "stage4.1.fuse_layers.1.2.1", true); + auto add_1306 = convBnUpAdd(network, weightMap, *relu_1192->getOutput(0), *add_1284->getOutput(0), width * 2, 1, 1, 0, + "stage4.1.fuse_layers.1.3.0", "stage4.1.fuse_layers.1.3.1", true); + auto relu_1307 = network->addActivation(*add_1306->getOutput(0), ActivationType::kRELU); + + auto relu_1310 = convBnRelu(network, weightMap, *relu_1108->getOutput(0), width, 3, 2, 1, "stage4.1.fuse_layers.2.0.0.0", "stage4.1.fuse_layers.2.0.0.1"); + auto bn_1312 = convBnRelu(network, weightMap, *relu_1310->getOutput(0), width * 4, 3, 2, 1, "stage4.1.fuse_layers.2.0.1.0", "stage4.1.fuse_layers.2.0.1.1", false); + auto add_1315 = convBnUpAdd(network, weightMap, *relu_1136->getOutput(0), *bn_1312->getOutput(0), width * 4, 3, 2, 1, + "stage4.1.fuse_layers.2.1.0.0", "stage4.1.fuse_layers.2.1.0.1", false); + auto add_1316 = network->addElementWise(*relu_1164->getOutput(0), *add_1315->getOutput(0), ElementWiseOperation::kSUM); + auto add_1338 = convBnUpAdd(network, weightMap, *relu_1192->getOutput(0), *add_1316->getOutput(0), width * 4, 1, 1, 0, + "stage4.1.fuse_layers.2.3.0", "stage4.1.fuse_layers.2.3.1", true); + auto relu_1339 = network->addActivation(*add_1338->getOutput(0), ActivationType::kRELU); + + auto relu_1342 = convBnRelu(network, weightMap, *relu_1108->getOutput(0), width, 3, 2, 1, "stage4.1.fuse_layers.3.0.0.0", "stage4.1.fuse_layers.3.0.0.1"); + auto relu_1345 = convBnRelu(network, weightMap, *relu_1342->getOutput(0), width, 3, 2, 1, "stage4.1.fuse_layers.3.0.1.0", "stage4.1.fuse_layers.3.0.1.1"); + auto bn_1347 = convBnRelu(network, weightMap, *relu_1345->getOutput(0), width * 8, 3, 2, 1, "stage4.1.fuse_layers.3.0.2.0", "stage4.1.fuse_layers.3.0.2.1", false); + auto relu_1350 = convBnRelu(network, weightMap, *relu_1136->getOutput(0), width * 2, 3, 2, 1, "stage4.1.fuse_layers.3.1.0.0", "stage4.1.fuse_layers.3.1.0.1"); + auto add_1353 = convBnUpAdd(network, weightMap, *relu_1350->getOutput(0), *bn_1347->getOutput(0), width * 8, 3, 2, 1, + "stage4.1.fuse_layers.3.1.1.0", "stage4.1.fuse_layers.3.1.1.1", false); + auto add_1356 = convBnUpAdd(network, weightMap, *relu_1164->getOutput(0), *add_1353->getOutput(0), width * 8, 3, 2, 1, + "stage4.1.fuse_layers.3.2.0.0", "stage4.1.fuse_layers.3.2.0.1", false); + auto add_1357 = network->addElementWise(*relu_1192->getOutput(0), *add_1356->getOutput(0), ElementWiseOperation::kSUM); + auto relu_1358 = network->addActivation(*add_1357->getOutput(0), ActivationType::kRELU); + + auto relu_1365 = liteResBlock(network, weightMap, *relu_1259->getOutput(0), width, "stage4.2.branches.0.0"); + auto relu_1372 = liteResBlock(network, weightMap, *relu_1365->getOutput(0), width, "stage4.2.branches.0.1"); + auto relu_1379 = liteResBlock(network, weightMap, *relu_1372->getOutput(0), width, "stage4.2.branches.0.2"); + auto relu_1386 = liteResBlock(network, weightMap, *relu_1379->getOutput(0), width, "stage4.2.branches.0.3"); + + auto relu_1393 = liteResBlock(network, weightMap, *relu_1307->getOutput(0), width * 2, "stage4.2.branches.1.0"); + auto relu_1400 = liteResBlock(network, weightMap, *relu_1393->getOutput(0), width * 2, "stage4.2.branches.1.1"); + auto relu_1407 = liteResBlock(network, weightMap, *relu_1400->getOutput(0), width * 2, "stage4.2.branches.1.2"); + auto relu_1414 = liteResBlock(network, weightMap, *relu_1407->getOutput(0), width * 2, "stage4.2.branches.1.3"); + + auto relu_1421 = liteResBlock(network, weightMap, *relu_1339->getOutput(0), width * 4, "stage4.2.branches.2.0"); + auto relu_1428 = liteResBlock(network, weightMap, *relu_1421->getOutput(0), width * 4, "stage4.2.branches.2.1"); + auto relu_1435 = liteResBlock(network, weightMap, *relu_1428->getOutput(0), width * 4, "stage4.2.branches.2.2"); + auto relu_1442 = liteResBlock(network, weightMap, *relu_1435->getOutput(0), width * 4, "stage4.2.branches.2.3"); + + auto relu_1449 = liteResBlock(network, weightMap, *relu_1358->getOutput(0), width * 8, "stage4.2.branches.3.0"); + auto relu_1456 = liteResBlock(network, weightMap, *relu_1449->getOutput(0), width * 8, "stage4.2.branches.3.1"); + auto relu_1463 = liteResBlock(network, weightMap, *relu_1456->getOutput(0), width * 8, "stage4.2.branches.3.2"); + auto relu_1470 = liteResBlock(network, weightMap, *relu_1463->getOutput(0), width * 8, "stage4.2.branches.3.3"); + + auto add_1492 = convBnUpAdd(network, weightMap, *relu_1414->getOutput(0), *relu_1386->getOutput(0), width, 1, 1, 0, + "stage4.2.fuse_layers.0.1.0", "stage4.2.fuse_layers.0.1.1", true); + auto add_1514 = convBnUpAdd(network, weightMap, *relu_1442->getOutput(0), *add_1492->getOutput(0), width, 1, 1, 0, + "stage4.2.fuse_layers.0.2.0", "stage4.2.fuse_layers.0.2.1", true); + + auto add_1536 = convBnUpAdd(network, weightMap, *relu_1470->getOutput(0), *add_1514->getOutput(0), width, 1, 1, 0, + "stage4.2.fuse_layers.0.3.0", "stage4.2.fuse_layers.0.3.1", true); + auto relu_1537 = network->addActivation(*add_1536->getOutput(0), ActivationType::kRELU); + + auto add_1540 = convBnUpAdd(network, weightMap, *relu_1386->getOutput(0), *relu_1414->getOutput(0), + width * 2, 3, 2, 1, "stage4.2.fuse_layers.1.0.0.0", "stage4.2.fuse_layers.1.0.0.1", false); + auto add_1562 = convBnUpAdd(network, weightMap, *relu_1442->getOutput(0), *add_1540->getOutput(0), + width * 2, 1, 1, 0, "stage4.2.fuse_layers.1.2.0", "stage4.2.fuse_layers.1.2.1", true); + auto add_1584 = convBnUpAdd(network, weightMap, *relu_1470->getOutput(0), *add_1562->getOutput(0), + width * 2, 1, 1, 0, "stage4.2.fuse_layers.1.3.0", "stage4.2.fuse_layers.1.3.1", true); + auto relu_1585 = network->addActivation(*add_1584->getOutput(0), ActivationType::kRELU); + + auto relu_1588 = convBnRelu(network, weightMap, *relu_1386->getOutput(0), width, 3, 2, 1, "stage4.2.fuse_layers.2.0.0.0", "stage4.2.fuse_layers.2.0.0.1"); + auto bn_1590 = convBnRelu(network, weightMap, *relu_1588->getOutput(0), width * 4, 3, 2, 1, "stage4.2.fuse_layers.2.0.1.0", "stage4.2.fuse_layers.2.0.1.1", false); + auto add_1593 = convBnUpAdd(network, weightMap, *relu_1414->getOutput(0), *bn_1590->getOutput(0), width * 4, 3, 2, 1, + "stage4.2.fuse_layers.2.1.0.0", "stage4.2.fuse_layers.2.1.0.1", false); + auto add_1594 = network->addElementWise(*relu_1442->getOutput(0), *add_1593->getOutput(0), ElementWiseOperation::kSUM); + auto add_1616 = convBnUpAdd(network, weightMap, *relu_1470->getOutput(0), *add_1594->getOutput(0), width * 4, 1, 1, 0, + "stage4.2.fuse_layers.2.3.0", "stage4.2.fuse_layers.2.3.1", true); + auto relu_1617 = network->addActivation(*add_1616->getOutput(0), ActivationType::kRELU); + + auto relu_1620 = convBnRelu(network, weightMap, *relu_1386->getOutput(0), width, 3, 2, 1, "stage4.2.fuse_layers.3.0.0.0", "stage4.2.fuse_layers.3.0.0.1"); + auto relu_1623 = convBnRelu(network, weightMap, *relu_1620->getOutput(0), width, 3, 2, 1, "stage4.2.fuse_layers.3.0.1.0", "stage4.2.fuse_layers.3.0.1.1"); + auto bn_1625 = convBnRelu(network, weightMap, *relu_1623->getOutput(0), width * 8, 3, 2, 1, "stage4.2.fuse_layers.3.0.2.0", "stage4.2.fuse_layers.3.0.2.1", false); + auto relu_1628 = convBnRelu(network, weightMap, *relu_1414->getOutput(0), width * 2, 3, 2, 1, "stage4.2.fuse_layers.3.1.0.0", "stage4.2.fuse_layers.3.1.0.1"); + auto add_1631 = convBnUpAdd(network, weightMap, *relu_1628->getOutput(0), *bn_1625->getOutput(0), width * 8, 3, 2, 1, + "stage4.2.fuse_layers.3.1.1.0", "stage4.2.fuse_layers.3.1.1.1", false); + auto add_1634 = convBnUpAdd(network, weightMap, *relu_1442->getOutput(0), *add_1631->getOutput(0), width * 8, 3, 2, 1, + "stage4.2.fuse_layers.3.2.0.0", "stage4.2.fuse_layers.3.2.0.1", false); + auto add_1635 = network->addElementWise(*relu_1470->getOutput(0), *add_1634->getOutput(0), ElementWiseOperation::kSUM); + auto relu_1636 = network->addActivation(*add_1635->getOutput(0), ActivationType::kRELU); + + nvinfer1::Dims dim = relu_1537->getOutput(0)->getDimensions(); + dim.d[0] = relu_1585->getOutput(0)->getDimensions().d[0]; + auto resize_1655 = netAddUpsampleBi(network, relu_1585->getOutput(0), dim); + dim.d[0] = relu_1617->getOutput(0)->getDimensions().d[0]; + auto resize_1668 = netAddUpsampleBi(network, relu_1617->getOutput(0), dim); + dim.d[0] = relu_1636->getOutput(0)->getDimensions().d[0]; + auto resize_1681 = netAddUpsampleBi(network, relu_1636->getOutput(0), dim); + + ITensor *concatTensors[] = {relu_1537->getOutput(0), resize_1655->getOutput(0), resize_1668->getOutput(0), resize_1681->getOutput(0)}; + auto concat_1682 = network->addConcatenation(concatTensors, 4); + concat_1682->setAxis(0); + auto relu_1685 = convBnRelu(network, weightMap, *concat_1682->getOutput(0), width * 15, 1, 1, 0, "last_layer.0", "last_layer.1", true, true); + auto conv_1686 = network->addConvolutionNd(*relu_1685->getOutput(0), NUM_CLASSES, DimsHW{1, 1}, weightMap["last_layer.3.weight"], weightMap["last_layer.3.bias"]); + conv_1686->setStrideNd(DimsHW{1, 1}); + conv_1686->setPaddingNd(DimsHW{0, 0}); + debug_print(conv_1686->getOutput(0), "conv_1686"); + dim.d[0] = NUM_CLASSES; + dim.d[1] = INPUT_H; + dim.d[2] = INPUT_W; + auto feature_map = netAddUpsampleBi(network, conv_1686->getOutput(0), dim); + debug_print(feature_map->getOutput(0), "feature_map"); + auto topk = network->addTopK(*feature_map->getOutput(0), TopKOperation::kMAX, 1, 0X01); + debug_print(topk->getOutput(0), "topk"); + std::cout << "set name out" << std::endl; + // topk->getOutput(1) 1 is index + topk->getOutput(1)->setName(OUTPUT_BLOB_NAME); + network->markOutput(*topk->getOutput(1)); + + builder->setMaxBatchSize(maxBatchSize); + config->setMaxWorkspaceSize((1 << 30)); // 1G +#ifdef USE_FP16 + std::cout << "use fp16" << std::endl; + config->setFlag(BuilderFlag::kFP16); +#endif + ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config); + std::cout << "build success!" << std::endl; + network->destroy(); + for (auto &mem : weightMap) + { + free((void *)(mem.second.values)); + } + return engine; +} +void APIToModel(unsigned int maxBatchSize, IHostMemory **modelStream, std::string wtsPath, int width) +{ + IBuilder *builder = createInferBuilder(gLogger); + IBuilderConfig *config = builder->createBuilderConfig(); + ICudaEngine *engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT, wtsPath, width); + assert(engine != nullptr); + (*modelStream) = engine->serialize(); + engine->destroy(); + builder->destroy(); +} + +bool parse_args(int argc, char **argv, std::string &wts, std::string &engine, int &width, std::string &img_dir) +{ + if (std::string(argv[1]) == "-s" && argc == 5) + { + wts = std::string(argv[2]); + engine = std::string(argv[3]); + width = std::stoi(argv[4]); + } + else if (std::string(argv[1]) == "-d" && argc == 4) + { + engine = std::string(argv[2]); + img_dir = std::string(argv[3]); + } + else + { + return false; + } + return true; +} +void doInference(IExecutionContext &context, cudaStream_t &stream, void **buffers, int batchSize) +{ + context.enqueue(batchSize, buffers, stream, nullptr); + cudaStreamSynchronize(stream); + cudaDeviceSynchronize(); +} + +int main(int argc, char **argv) +{ + cudaSetDevice(DEVICE); + std::string wtsPath = ""; + std::string engine_name = ""; + int width; + std::string img_dir; + // parse args + if (!parse_args(argc, argv, wtsPath, engine_name, width, img_dir)) + { + std::cerr << "arguments not right!" << std::endl; + std::cerr << "./hrnet -s [.wts] [.engine] [18 or 32 or 48] // serialize model to plan file" << std::endl; + std::cerr << "./hrnet -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; + return -1; + } + // create a model using the API directly and serialize it to a stream + if (!wtsPath.empty()) + { + IHostMemory *modelStream{nullptr}; + APIToModel(BATCH_SIZE, &modelStream, wtsPath, width); + assert(modelStream != nullptr); + std::ofstream p(engine_name, std::ios::binary); + if (!p) + { + std::cerr << "could not open plan output file" << std::endl; + return -1; + } + p.write(reinterpret_cast(modelStream->data()), modelStream->size()); + modelStream->destroy(); + return 0; + } + + // deserialize the .engine and run inference + char *trtModelStream{nullptr}; + size_t size{0}; + std::ifstream file(engine_name, std::ios::binary); + if (file.good()) + { + file.seekg(0, file.end); + size = file.tellg(); + file.seekg(0, file.beg); + trtModelStream = new char[size]; + assert(trtModelStream); + file.read(trtModelStream, size); + file.close(); + } + else + { + std::cerr << "could not open plan file" << std::endl; + } + + std::vector file_names; + if (read_files_in_dir(img_dir.c_str(), file_names) < 0) + { + std::cout << "read_files_in_dir failed." << std::endl; + return -1; + } + // prepare input data --------------------------- + cudaSetDeviceFlags(cudaDeviceMapHost); + float *data; + int *prob; // using int. output is index + CHECK(cudaHostAlloc((void **)&data, BATCH_SIZE * 3 * INPUT_H * INPUT_W * sizeof(float), cudaHostAllocMapped)); + CHECK(cudaHostAlloc((void **)&prob, BATCH_SIZE * OUTPUT_SIZE * sizeof(int), cudaHostAllocMapped)); + + IRuntime *runtime = createInferRuntime(gLogger); + assert(runtime != nullptr); + ICudaEngine *engine = runtime->deserializeCudaEngine(trtModelStream, size); + assert(engine != nullptr); + IExecutionContext *context = engine->createExecutionContext(); + assert(context != nullptr); + delete[] trtModelStream; + void *buffers[2]; + // In order to bind the buffers, we need to know the names of the input and output tensors. + // Note that indices are guaranteed to be less than IEngine::getNbBindings() + const int inputIndex = engine->getBindingIndex(INPUT_BLOB_NAME); + const int outputIndex = engine->getBindingIndex(OUTPUT_BLOB_NAME); + assert(inputIndex == 0); + assert(outputIndex == 1); + cudaStream_t stream; + CHECK(cudaStreamCreate(&stream)); + + for (int f = 0; f < (int)file_names.size(); f++) + { + std::cout << file_names[f] << std::endl; + cv::Mat pr_img; + cv::Mat img_BGR = cv::imread(img_dir + "/" + file_names[f], 1); // BGR + cv::Mat img; + cv::cvtColor(img_BGR, img, cv::COLOR_BGR2RGB); + if (img.empty()) + continue; + cv::resize(img, pr_img, cv::Size(INPUT_W, INPUT_H)); + img = pr_img.clone(); // for img show + pr_img.convertTo(pr_img, CV_32FC3); + if (!pr_img.isContinuous()) + { + pr_img = pr_img.clone(); + } + std::memcpy(data, pr_img.data, BATCH_SIZE * 3 * INPUT_W * INPUT_H * sizeof(float)); + + cudaHostGetDevicePointer((void **)&buffers[inputIndex], (void *)data, 0); // buffers[inputIndex]-->data + cudaHostGetDevicePointer((void **)&buffers[outputIndex], (void *)prob, 0); // buffers[outputIndex] --> prob + + // Run inference + auto start = std::chrono::high_resolution_clock::now(); + doInference(*context, stream, buffers, BATCH_SIZE); + auto end = std::chrono::high_resolution_clock::now(); + std::cout << "infer time: " << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; + + cv::Mat outimg(INPUT_H, INPUT_W, CV_8UC1); + for (int row = 0; row < INPUT_H; ++row) + { + uchar *uc_pixel = outimg.data + row * outimg.step; + for (int col = 0; col < INPUT_W; ++col) + { + uc_pixel[col] = (uchar)prob[row * INPUT_W + col]; + } + } + cv::Mat im_color; + cv::cvtColor(outimg, im_color, cv::COLOR_GRAY2RGB); + cv::Mat lut = createLTU(NUM_CLASSES); + cv::LUT(im_color, lut, im_color); + // false color + cv::cvtColor(im_color, im_color, cv::COLOR_RGB2GRAY); + cv::applyColorMap(im_color, im_color, cv::COLORMAP_HOT); + // cv::imshow("False Color Map", im_color); + cv::imwrite(std::to_string(f) + "_false_color_map.png", im_color); + //fusion + cv::Mat fusionImg; + cv::addWeighted(img, 1, im_color, 0.8, 1, fusionImg); + // cv::imshow("Fusion Img", fusionImg); + // cv::waitKey(0); + cv::imwrite(std::to_string(f) + "_fusion_img.png", fusionImg); + } + + // Release stream and buffers + cudaStreamDestroy(stream); + CHECK(cudaFreeHost(buffers[inputIndex])); + CHECK(cudaFreeHost(buffers[outputIndex])); + // Destroy the engine + context->destroy(); + engine->destroy(); + runtime->destroy(); + return 0; +} diff --git a/hrnet/hrnet-semantic-segmentation/hrnet_ocr.cpp b/hrnet/hrnet-semantic-segmentation/hrnet_ocr.cpp new file mode 100644 index 0000000..106d108 --- /dev/null +++ b/hrnet/hrnet-semantic-segmentation/hrnet_ocr.cpp @@ -0,0 +1,690 @@ +#include +#include +#include +#include +#include +#include +#include "common.hpp" +#include "logging.h" + +static Logger gLogger; +#define USE_FP32 +#define DEVICE 0 // GPU id +#define BATCH_SIZE 1 // + +const char *INPUT_BLOB_NAME = "data"; +const char *OUTPUT_BLOB_NAME = "output"; +static const int INPUT_H = 512; +static const int INPUT_W = 1024; +static const int NUM_CLASSES = 19; +static const int OUTPUT_SIZE = INPUT_H * INPUT_W; + +// Creat the engine using only the API and not any parser. +ICudaEngine *createEngine(unsigned int maxBatchSize, IBuilder *builder, IBuilderConfig *config, DataType dt, std::string wtsPath, int width) +{ + 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{INPUT_H, INPUT_W, 3}); + assert(data); + + // hwc to chw + auto ps = network->addShuffle(*data); + ps->setFirstTranspose(nvinfer1::Permutation{2, 0, 1}); + float mean[3] = {0.485, 0.456, 0.406}; + float std[3] = {0.229, 0.224, 0.225}; + ITensor *preinput = MeanStd(network, ps->getOutput(0), mean, std, true); + + std::map weightMap = loadWeights(wtsPath); + auto relu_2 = convBnRelu(network, weightMap, *preinput, 64, 3, 2, 1, "conv1", "bn1"); + auto relu_5 = convBnRelu(network, weightMap, *relu_2->getOutput(0), 64, 3, 2, 1, "conv2", "bn2"); + auto relu_17 = ResBlock2Conv(network, weightMap, *relu_5->getOutput(0), 64, 256, 1, "layer1.0"); + auto relu_27 = ResBlock(network, weightMap, *relu_17->getOutput(0), 256, 64, 1, "layer1.1"); + auto relu_37 = ResBlock(network, weightMap, *relu_27->getOutput(0), 256, 64, 1, "layer1.2"); + auto relu_47 = ResBlock(network, weightMap, *relu_37->getOutput(0), 256, 64, 1, "layer1.3"); + + auto relu_50 = convBnRelu(network, weightMap, *relu_47->getOutput(0), width, 3, 1, 1, "transition1.0.0", "transition1.0.1"); + auto relu_60 = liteResBlock(network, weightMap, *relu_50->getOutput(0), width, "stage2.0.branches.0.0"); + auto relu_67 = liteResBlock(network, weightMap, *relu_60->getOutput(0), width, "stage2.0.branches.0.1"); + auto relu_74 = liteResBlock(network, weightMap, *relu_67->getOutput(0), width, "stage2.0.branches.0.2"); + auto relu_81 = liteResBlock(network, weightMap, *relu_74->getOutput(0), width, "stage2.0.branches.0.3"); + + auto relu_53 = convBnRelu(network, weightMap, *relu_47->getOutput(0), width * 2, 3, 2, 1, "transition1.1.0.0", "transition1.1.0.1"); + auto relu_88 = liteResBlock(network, weightMap, *relu_53->getOutput(0), width * 2, "stage2.0.branches.1.0"); + auto relu_95 = liteResBlock(network, weightMap, *relu_88->getOutput(0), width * 2, "stage2.0.branches.1.1"); + auto relu_102 = liteResBlock(network, weightMap, *relu_95->getOutput(0), width * 2, "stage2.0.branches.1.2"); + auto relu_109 = liteResBlock(network, weightMap, *relu_102->getOutput(0), width * 2, "stage2.0.branches.1.3"); + + auto add_131 = convBnUpAdd(network, weightMap, *relu_109->getOutput(0), *relu_81->getOutput(0), width, 1, 1, 0, "stage2.0.fuse_layers.0.1.0", "stage2.0.fuse_layers.0.1.1", true); + auto relu_132 = network->addActivation(*add_131->getOutput(0), ActivationType::kRELU); + + auto add_135 = convBnUpAdd(network, weightMap, *relu_81->getOutput(0), *relu_109->getOutput(0), width * 2, 3, 2, 1, "stage2.0.fuse_layers.1.0.0.0", "stage2.0.fuse_layers.1.0.0.1", false); + auto relu_136 = network->addActivation(*add_135->getOutput(0), ActivationType::kRELU); + + auto relu_146 = liteResBlock(network, weightMap, *relu_132->getOutput(0), width, "stage3.0.branches.0.0"); + auto relu_153 = liteResBlock(network, weightMap, *relu_146->getOutput(0), width, "stage3.0.branches.0.1"); + auto relu_160 = liteResBlock(network, weightMap, *relu_153->getOutput(0), width, "stage3.0.branches.0.2"); + auto relu_167 = liteResBlock(network, weightMap, *relu_160->getOutput(0), width, "stage3.0.branches.0.3"); + + auto relu_174 = liteResBlock(network, weightMap, *relu_136->getOutput(0), width * 2, "stage3.0.branches.1.0"); + auto relu_181 = liteResBlock(network, weightMap, *relu_174->getOutput(0), width * 2, "stage3.0.branches.1.1"); + auto relu_188 = liteResBlock(network, weightMap, *relu_181->getOutput(0), width * 2, "stage3.0.branches.1.2"); + auto relu_195 = liteResBlock(network, weightMap, *relu_188->getOutput(0), width * 2, "stage3.0.branches.1.3"); + + auto relu_139 = convBnRelu(network, weightMap, *relu_136->getOutput(0), width * 4, 3, 2, 1, "transition2.2.0.0", "transition2.2.0.1"); + auto relu_202 = liteResBlock(network, weightMap, *relu_139->getOutput(0), width * 4, "stage3.0.branches.2.0"); + auto relu_209 = liteResBlock(network, weightMap, *relu_202->getOutput(0), width * 4, "stage3.0.branches.2.1"); + auto relu_216 = liteResBlock(network, weightMap, *relu_209->getOutput(0), width * 4, "stage3.0.branches.2.2"); + auto relu_223 = liteResBlock(network, weightMap, *relu_216->getOutput(0), width * 4, "stage3.0.branches.2.3"); + + auto add_245 = convBnUpAdd(network, weightMap, *relu_195->getOutput(0), *relu_167->getOutput(0), width, 1, 1, 0, "stage3.0.fuse_layers.0.1.0", "stage3.0.fuse_layers.0.1.1", true); + auto add_267 = convBnUpAdd(network, weightMap, *relu_223->getOutput(0), *add_245->getOutput(0), width, 1, 1, 0, "stage3.0.fuse_layers.0.2.0", "stage3.0.fuse_layers.0.2.1", true); + auto relu_268 = network->addActivation(*add_267->getOutput(0), ActivationType::kRELU); + + auto add_271 = convBnUpAdd(network, weightMap, *relu_167->getOutput(0), *relu_195->getOutput(0), width * 2, 3, 2, 1, "stage3.0.fuse_layers.1.0.0.0", "stage3.0.fuse_layers.1.0.0.1", false); + auto add_293 = convBnUpAdd(network, weightMap, *relu_223->getOutput(0), *add_271->getOutput(0), width * 2, 1, 1, 0, "stage3.0.fuse_layers.1.2.0", "stage3.0.fuse_layers.1.2.1", true); + auto relu_294 = network->addActivation(*add_293->getOutput(0), ActivationType::kRELU); + + auto relu_297 = convBnRelu(network, weightMap, *relu_167->getOutput(0), width, 3, 2, 1, "stage3.0.fuse_layers.2.0.0.0", "stage3.0.fuse_layers.2.0.0.1"); + auto bn_299 = convBnRelu(network, weightMap, *relu_297->getOutput(0), width * 4, 3, 2, 1, "stage3.0.fuse_layers.2.0.1.0", "stage3.0.fuse_layers.2.0.1.1", false); + auto add_302 = convBnUpAdd(network, weightMap, *relu_195->getOutput(0), *bn_299->getOutput(0), width * 4, 3, 2, 1, "stage3.0.fuse_layers.2.1.0.0", "stage3.0.fuse_layers.2.1.0.1", false); + auto add_303 = network->addElementWise(*add_302->getOutput(0), *relu_223->getOutput(0), ElementWiseOperation::kSUM); + auto relu_304 = network->addActivation(*add_303->getOutput(0), ActivationType::kRELU); + + auto relu_311 = liteResBlock(network, weightMap, *relu_268->getOutput(0), width, "stage3.1.branches.0.0"); + auto relu_318 = liteResBlock(network, weightMap, *relu_311->getOutput(0), width, "stage3.1.branches.0.1"); + auto relu_325 = liteResBlock(network, weightMap, *relu_318->getOutput(0), width, "stage3.1.branches.0.2"); + auto relu_332 = liteResBlock(network, weightMap, *relu_325->getOutput(0), width, "stage3.1.branches.0.3"); + + auto relu_339 = liteResBlock(network, weightMap, *relu_294->getOutput(0), width * 2, "stage3.1.branches.1.0"); + auto relu_346 = liteResBlock(network, weightMap, *relu_339->getOutput(0), width * 2, "stage3.1.branches.1.1"); + auto relu_353 = liteResBlock(network, weightMap, *relu_346->getOutput(0), width * 2, "stage3.1.branches.1.2"); + auto relu_360 = liteResBlock(network, weightMap, *relu_353->getOutput(0), width * 2, "stage3.1.branches.1.3"); + + auto relu_367 = liteResBlock(network, weightMap, *relu_304->getOutput(0), width * 4, "stage3.1.branches.2.0"); + auto relu_374 = liteResBlock(network, weightMap, *relu_367->getOutput(0), width * 4, "stage3.1.branches.2.1"); + auto relu_381 = liteResBlock(network, weightMap, *relu_374->getOutput(0), width * 4, "stage3.1.branches.2.2"); + auto relu_388 = liteResBlock(network, weightMap, *relu_381->getOutput(0), width * 4, "stage3.1.branches.2.3"); + + auto add_410 = convBnUpAdd(network, weightMap, *relu_360->getOutput(0), *relu_332->getOutput(0), width, 1, 1, 0, "stage3.1.fuse_layers.0.1.0", "stage3.1.fuse_layers.0.1.1", true); + auto add_432 = convBnUpAdd(network, weightMap, *relu_388->getOutput(0), *add_410->getOutput(0), width, 1, 1, 0, "stage3.1.fuse_layers.0.2.0", "stage3.1.fuse_layers.0.2.1", true); + auto relu_433 = network->addActivation(*add_432->getOutput(0), ActivationType::kRELU); + + auto add_436 = convBnUpAdd(network, weightMap, *relu_332->getOutput(0), *relu_360->getOutput(0), width * 2, 3, 2, 1, "stage3.1.fuse_layers.1.0.0.0", "stage3.1.fuse_layers.1.0.0.1", false); + auto add_458 = convBnUpAdd(network, weightMap, *relu_388->getOutput(0), *add_436->getOutput(0), width * 2, 1, 1, 0, "stage3.1.fuse_layers.1.2.0", "stage3.1.fuse_layers.1.2.1", true); + auto relu_459 = network->addActivation(*add_458->getOutput(0), ActivationType::kRELU); + + auto relu_462 = convBnRelu(network, weightMap, *relu_332->getOutput(0), width, 3, 2, 1, "stage3.1.fuse_layers.2.0.0.0", "stage3.1.fuse_layers.2.0.0.1"); + auto bn_464 = convBnRelu(network, weightMap, *relu_462->getOutput(0), width * 4, 3, 2, 1, "stage3.1.fuse_layers.2.0.1.0", "stage3.1.fuse_layers.2.0.1.1", false); + auto add_467 = convBnUpAdd(network, weightMap, *relu_360->getOutput(0), *bn_464->getOutput(0), width * 4, 3, 2, 1, "stage3.1.fuse_layers.2.1.0.0", "stage3.1.fuse_layers.2.1.0.1", false); + auto add_468 = network->addElementWise(*add_467->getOutput(0), *relu_388->getOutput(0), ElementWiseOperation::kSUM); + auto relu_469 = network->addActivation(*add_468->getOutput(0), ActivationType::kRELU); + + auto relu_476 = liteResBlock(network, weightMap, *relu_433->getOutput(0), width, "stage3.2.branches.0.0"); + auto relu_483 = liteResBlock(network, weightMap, *relu_476->getOutput(0), width, "stage3.2.branches.0.1"); + auto relu_490 = liteResBlock(network, weightMap, *relu_483->getOutput(0), width, "stage3.2.branches.0.2"); + auto relu_497 = liteResBlock(network, weightMap, *relu_490->getOutput(0), width, "stage3.2.branches.0.3"); + + auto relu_504 = liteResBlock(network, weightMap, *relu_459->getOutput(0), width * 2, "stage3.2.branches.1.0"); + auto relu_511 = liteResBlock(network, weightMap, *relu_504->getOutput(0), width * 2, "stage3.2.branches.1.1"); + auto relu_518 = liteResBlock(network, weightMap, *relu_511->getOutput(0), width * 2, "stage3.2.branches.1.2"); + auto relu_525 = liteResBlock(network, weightMap, *relu_518->getOutput(0), width * 2, "stage3.2.branches.1.3"); + + auto relu_532 = liteResBlock(network, weightMap, *relu_469->getOutput(0), width * 4, "stage3.2.branches.2.0"); + auto relu_539 = liteResBlock(network, weightMap, *relu_532->getOutput(0), width * 4, "stage3.2.branches.2.1"); + auto relu_546 = liteResBlock(network, weightMap, *relu_539->getOutput(0), width * 4, "stage3.2.branches.2.2"); + auto relu_553 = liteResBlock(network, weightMap, *relu_546->getOutput(0), width * 4, "stage3.2.branches.2.3"); + + auto add_575 = convBnUpAdd(network, weightMap, *relu_525->getOutput(0), *relu_497->getOutput(0), width, 1, 1, 0, "stage3.2.fuse_layers.0.1.0", "stage3.2.fuse_layers.0.1.1", true); + auto add_597 = convBnUpAdd(network, weightMap, *relu_553->getOutput(0), *add_575->getOutput(0), width, 1, 1, 0, "stage3.2.fuse_layers.0.2.0", "stage3.2.fuse_layers.0.2.1", true); + + auto relu_598 = network->addActivation(*add_597->getOutput(0), ActivationType::kRELU); + + auto add_601 = convBnUpAdd(network, weightMap, *relu_497->getOutput(0), *relu_525->getOutput(0), width * 2, 3, 2, 1, "stage3.2.fuse_layers.1.0.0.0", "stage3.2.fuse_layers.1.0.0.1", false); + auto add_623 = convBnUpAdd(network, weightMap, *relu_553->getOutput(0), *add_601->getOutput(0), width * 2, 1, 1, 0, "stage3.2.fuse_layers.1.2.0", "stage3.2.fuse_layers.1.2.1", true); + auto relu_624 = network->addActivation(*add_623->getOutput(0), ActivationType::kRELU); + + auto relu_627 = convBnRelu(network, weightMap, *relu_497->getOutput(0), width, 3, 2, 1, "stage3.2.fuse_layers.2.0.0.0", "stage3.2.fuse_layers.2.0.0.1"); + auto bn_629 = convBnRelu(network, weightMap, *relu_627->getOutput(0), width * 4, 3, 2, 1, "stage3.2.fuse_layers.2.0.1.0", "stage3.2.fuse_layers.2.0.1.1", false); + auto add_632 = convBnUpAdd(network, weightMap, *relu_525->getOutput(0), *bn_629->getOutput(0), width * 4, 3, 2, 1, "stage3.2.fuse_layers.2.1.0.0", "stage3.2.fuse_layers.2.1.0.1", false); + auto add_633 = network->addElementWise(*relu_553->getOutput(0), *add_632->getOutput(0), ElementWiseOperation::kSUM); + auto relu_634 = network->addActivation(*add_633->getOutput(0), ActivationType::kRELU); + + auto relu_641 = liteResBlock(network, weightMap, *relu_598->getOutput(0), width, "stage3.3.branches.0.0"); + auto relu_648 = liteResBlock(network, weightMap, *relu_641->getOutput(0), width, "stage3.3.branches.0.1"); + auto relu_655 = liteResBlock(network, weightMap, *relu_648->getOutput(0), width, "stage3.3.branches.0.2"); + auto relu_662 = liteResBlock(network, weightMap, *relu_655->getOutput(0), width, "stage3.3.branches.0.3"); + + auto relu_669 = liteResBlock(network, weightMap, *relu_624->getOutput(0), width * 2, "stage3.3.branches.1.0"); + auto relu_676 = liteResBlock(network, weightMap, *relu_669->getOutput(0), width * 2, "stage3.3.branches.1.1"); + auto relu_683 = liteResBlock(network, weightMap, *relu_676->getOutput(0), width * 2, "stage3.3.branches.1.2"); + auto relu_690 = liteResBlock(network, weightMap, *relu_683->getOutput(0), width * 2, "stage3.3.branches.1.3"); + + auto relu_697 = liteResBlock(network, weightMap, *relu_634->getOutput(0), width * 4, "stage3.3.branches.2.0"); + auto relu_704 = liteResBlock(network, weightMap, *relu_697->getOutput(0), width * 4, "stage3.3.branches.2.1"); + auto relu_711 = liteResBlock(network, weightMap, *relu_704->getOutput(0), width * 4, "stage3.3.branches.2.2"); + auto relu_718 = liteResBlock(network, weightMap, *relu_711->getOutput(0), width * 4, "stage3.3.branches.2.3"); + + auto add_740 = convBnUpAdd(network, weightMap, *relu_690->getOutput(0), *relu_662->getOutput(0), width, 1, 1, 0, "stage3.3.fuse_layers.0.1.0", "stage3.3.fuse_layers.0.1.1", true); + auto add_762 = convBnUpAdd(network, weightMap, *relu_718->getOutput(0), *add_740->getOutput(0), width, 1, 1, 0, "stage3.3.fuse_layers.0.2.0", "stage3.3.fuse_layers.0.2.1", true); + auto relu_763 = network->addActivation(*add_762->getOutput(0), ActivationType::kRELU); + + auto add_766 = convBnUpAdd(network, weightMap, *relu_662->getOutput(0), *relu_690->getOutput(0), width * 2, 3, 2, 1, "stage3.3.fuse_layers.1.0.0.0", "stage3.3.fuse_layers.1.0.0.1", false); + auto add_788 = convBnUpAdd(network, weightMap, *relu_718->getOutput(0), *add_766->getOutput(0), width * 2, 1, 1, 0, "stage3.3.fuse_layers.1.2.0", "stage3.3.fuse_layers.1.2.1", true); + auto relu_789 = network->addActivation(*add_788->getOutput(0), ActivationType::kRELU); + + auto relu_792 = convBnRelu(network, weightMap, *relu_662->getOutput(0), width, 3, 2, 1, "stage3.3.fuse_layers.2.0.0.0", "stage3.3.fuse_layers.2.0.0.1"); + auto bn_794 = convBnRelu(network, weightMap, *relu_792->getOutput(0), width * 4, 3, 2, 1, "stage3.3.fuse_layers.2.0.1.0", "stage3.3.fuse_layers.2.0.1.1", false); + auto add_797 = convBnUpAdd(network, weightMap, *relu_690->getOutput(0), *bn_794->getOutput(0), width * 4, 3, 2, 1, "stage3.3.fuse_layers.2.1.0.0", "stage3.3.fuse_layers.2.1.0.1", false); + auto add_798 = network->addElementWise(*relu_718->getOutput(0), *add_797->getOutput(0), ElementWiseOperation::kSUM); + auto relu_799 = network->addActivation(*add_798->getOutput(0), ActivationType::kRELU); + + auto relu_809 = liteResBlock(network, weightMap, *relu_763->getOutput(0), width, "stage4.0.branches.0.0"); + auto relu_816 = liteResBlock(network, weightMap, *relu_809->getOutput(0), width, "stage4.0.branches.0.1"); + auto relu_823 = liteResBlock(network, weightMap, *relu_816->getOutput(0), width, "stage4.0.branches.0.2"); + auto relu_830 = liteResBlock(network, weightMap, *relu_823->getOutput(0), width, "stage4.0.branches.0.3"); + + auto relu_837 = liteResBlock(network, weightMap, *relu_789->getOutput(0), width * 2, "stage4.0.branches.1.0"); + auto relu_844 = liteResBlock(network, weightMap, *relu_837->getOutput(0), width * 2, "stage4.0.branches.1.1"); + auto relu_851 = liteResBlock(network, weightMap, *relu_844->getOutput(0), width * 2, "stage4.0.branches.1.2"); + auto relu_858 = liteResBlock(network, weightMap, *relu_851->getOutput(0), width * 2, "stage4.0.branches.1.3"); + + auto relu_865 = liteResBlock(network, weightMap, *relu_799->getOutput(0), width * 4, "stage4.0.branches.2.0"); + auto relu_872 = liteResBlock(network, weightMap, *relu_865->getOutput(0), width * 4, "stage4.0.branches.2.1"); + auto relu_879 = liteResBlock(network, weightMap, *relu_872->getOutput(0), width * 4, "stage4.0.branches.2.2"); + auto relu_886 = liteResBlock(network, weightMap, *relu_879->getOutput(0), width * 4, "stage4.0.branches.2.3"); //======== + + auto relu_802 = convBnRelu(network, weightMap, *relu_799->getOutput(0), width * 8, 3, 2, 1, "transition3.3.0.0", "transition3.3.0.1"); + auto relu_893 = liteResBlock(network, weightMap, *relu_802->getOutput(0), width * 8, "stage4.0.branches.3.0"); + auto relu_900 = liteResBlock(network, weightMap, *relu_893->getOutput(0), width * 8, "stage4.0.branches.3.1"); + auto relu_907 = liteResBlock(network, weightMap, *relu_900->getOutput(0), width * 8, "stage4.0.branches.3.2"); + auto relu_914 = liteResBlock(network, weightMap, *relu_907->getOutput(0), width * 8, "stage4.0.branches.3.3"); + + auto add_936 = convBnUpAdd(network, weightMap, *relu_858->getOutput(0), *relu_830->getOutput(0), width, 1, 1, 0, "stage4.0.fuse_layers.0.1.0", "stage4.0.fuse_layers.0.1.1", true); + auto add_958 = convBnUpAdd(network, weightMap, *relu_886->getOutput(0), *add_936->getOutput(0), width, 1, 1, 0, "stage4.0.fuse_layers.0.2.0", "stage4.0.fuse_layers.0.2.1", true); + auto add_980 = convBnUpAdd(network, weightMap, *relu_914->getOutput(0), *add_958->getOutput(0), width, 1, 1, 0, "stage4.0.fuse_layers.0.3.0", "stage4.0.fuse_layers.0.3.1", true); + auto relu_981 = network->addActivation(*add_980->getOutput(0), ActivationType::kRELU); + + auto add_984 = convBnUpAdd(network, weightMap, *relu_830->getOutput(0), *relu_858->getOutput(0), width * 2, 3, 2, 1, "stage4.0.fuse_layers.1.0.0.0", "stage4.0.fuse_layers.1.0.0.1", false); + auto add_1006 = convBnUpAdd(network, weightMap, *relu_886->getOutput(0), *add_984->getOutput(0), width * 2, 1, 1, 0, "stage4.0.fuse_layers.1.2.0", "stage4.0.fuse_layers.1.2.1", true); + auto add_1028 = convBnUpAdd(network, weightMap, *relu_914->getOutput(0), *add_1006->getOutput(0), width * 2, 1, 1, 0, "stage4.0.fuse_layers.1.3.0", "stage4.0.fuse_layers.1.3.1", true); + auto relu_1029 = network->addActivation(*add_1028->getOutput(0), ActivationType::kRELU); + + auto relu_1032 = convBnRelu(network, weightMap, *relu_830->getOutput(0), width, 3, 2, 1, "stage4.0.fuse_layers.2.0.0.0", "stage4.0.fuse_layers.2.0.0.1"); + auto bn_1034 = convBnRelu(network, weightMap, *relu_1032->getOutput(0), width * 4, 3, 2, 1, "stage4.0.fuse_layers.2.0.1.0", "stage4.0.fuse_layers.2.0.1.1", false); + + auto add_1037 = convBnUpAdd(network, weightMap, *relu_858->getOutput(0), *bn_1034->getOutput(0), width * 4, 3, 2, 1, + "stage4.0.fuse_layers.2.1.0.0", "stage4.0.fuse_layers.2.1.0.1", false); + auto add_1038 = network->addElementWise(*relu_886->getOutput(0), *add_1037->getOutput(0), ElementWiseOperation::kSUM); + auto add_1060 = convBnUpAdd(network, weightMap, *relu_914->getOutput(0), *add_1038->getOutput(0), width * 4, 1, 1, 0, + "stage4.0.fuse_layers.2.3.0", "stage4.0.fuse_layers.2.3.1", true); + auto relu_1061 = network->addActivation(*add_1060->getOutput(0), ActivationType::kRELU); + + auto relu_1064 = convBnRelu(network, weightMap, *relu_830->getOutput(0), width, 3, 2, 1, "stage4.0.fuse_layers.3.0.0.0", "stage4.0.fuse_layers.3.0.0.1"); + auto relu_1067 = convBnRelu(network, weightMap, *relu_1064->getOutput(0), width, 3, 2, 1, "stage4.0.fuse_layers.3.0.1.0", "stage4.0.fuse_layers.3.0.1.1"); + auto bn_1069 = convBnRelu(network, weightMap, *relu_1067->getOutput(0), width * 8, 3, 2, 1, "stage4.0.fuse_layers.3.0.2.0", "stage4.0.fuse_layers.3.0.2.1", false); + auto relu_1072 = convBnRelu(network, weightMap, *relu_858->getOutput(0), width * 2, 3, 2, 1, "stage4.0.fuse_layers.3.1.0.0", "stage4.0.fuse_layers.3.1.0.1"); + auto add_1075 = convBnUpAdd(network, weightMap, *relu_1072->getOutput(0), *bn_1069->getOutput(0), width * 8, 3, 2, 1, + "stage4.0.fuse_layers.3.1.1.0", "stage4.0.fuse_layers.3.1.1.1", false); + auto add_1078 = convBnUpAdd(network, weightMap, *relu_886->getOutput(0), *add_1075->getOutput(0), width * 8, 3, 2, 1, + "stage4.0.fuse_layers.3.2.0.0", "stage4.0.fuse_layers.3.2.0.1", false); + auto add_1079 = network->addElementWise(*relu_914->getOutput(0), *add_1078->getOutput(0), ElementWiseOperation::kSUM); + auto relu_1080 = network->addActivation(*add_1079->getOutput(0), ActivationType::kRELU); + + auto relu_1087 = liteResBlock(network, weightMap, *relu_981->getOutput(0), width, "stage4.1.branches.0.0"); + auto relu_1094 = liteResBlock(network, weightMap, *relu_1087->getOutput(0), width, "stage4.1.branches.0.1"); + auto relu_1101 = liteResBlock(network, weightMap, *relu_1094->getOutput(0), width, "stage4.1.branches.0.2"); + auto relu_1108 = liteResBlock(network, weightMap, *relu_1101->getOutput(0), width, "stage4.1.branches.0.3"); + + auto relu_1115 = liteResBlock(network, weightMap, *relu_1029->getOutput(0), width * 2, "stage4.1.branches.1.0"); + auto relu_1122 = liteResBlock(network, weightMap, *relu_1115->getOutput(0), width * 2, "stage4.1.branches.1.1"); + auto relu_1129 = liteResBlock(network, weightMap, *relu_1122->getOutput(0), width * 2, "stage4.1.branches.1.2"); + auto relu_1136 = liteResBlock(network, weightMap, *relu_1129->getOutput(0), width * 2, "stage4.1.branches.1.3"); + + auto relu_1143 = liteResBlock(network, weightMap, *relu_1061->getOutput(0), width * 4, "stage4.1.branches.2.0"); + auto relu_1150 = liteResBlock(network, weightMap, *relu_1143->getOutput(0), width * 4, "stage4.1.branches.2.1"); + auto relu_1157 = liteResBlock(network, weightMap, *relu_1150->getOutput(0), width * 4, "stage4.1.branches.2.2"); + auto relu_1164 = liteResBlock(network, weightMap, *relu_1157->getOutput(0), width * 4, "stage4.1.branches.2.3"); + + auto relu_1171 = liteResBlock(network, weightMap, *relu_1080->getOutput(0), width * 8, "stage4.1.branches.3.0"); + auto relu_1178 = liteResBlock(network, weightMap, *relu_1171->getOutput(0), width * 8, "stage4.1.branches.3.1"); + auto relu_1185 = liteResBlock(network, weightMap, *relu_1178->getOutput(0), width * 8, "stage4.1.branches.3.2"); + auto relu_1192 = liteResBlock(network, weightMap, *relu_1185->getOutput(0), width * 8, "stage4.1.branches.3.3"); + + auto add_1214 = convBnUpAdd(network, weightMap, *relu_1136->getOutput(0), *relu_1108->getOutput(0), width, 1, 1, 0, + "stage4.1.fuse_layers.0.1.0", "stage4.1.fuse_layers.0.1.1", true); + auto add_1236 = convBnUpAdd(network, weightMap, *relu_1164->getOutput(0), *add_1214->getOutput(0), width, 1, 1, 0, + "stage4.1.fuse_layers.0.2.0", "stage4.1.fuse_layers.0.2.1", true); + auto add_1258 = convBnUpAdd(network, weightMap, *relu_1192->getOutput(0), *add_1236->getOutput(0), width, 1, 1, 0, + "stage4.1.fuse_layers.0.3.0", "stage4.1.fuse_layers.0.3.1", true); + auto relu_1259 = network->addActivation(*add_1258->getOutput(0), ActivationType::kRELU); + + auto add_1262 = convBnUpAdd(network, weightMap, *relu_1108->getOutput(0), *relu_1136->getOutput(0), width * 2, 3, 2, 1, + "stage4.1.fuse_layers.1.0.0.0", "stage4.1.fuse_layers.1.0.0.1", false); + auto add_1284 = convBnUpAdd(network, weightMap, *relu_1164->getOutput(0), *add_1262->getOutput(0), width * 2, 1, 1, 0, + "stage4.1.fuse_layers.1.2.0", "stage4.1.fuse_layers.1.2.1", true); + auto add_1306 = convBnUpAdd(network, weightMap, *relu_1192->getOutput(0), *add_1284->getOutput(0), width * 2, 1, 1, 0, + "stage4.1.fuse_layers.1.3.0", "stage4.1.fuse_layers.1.3.1", true); + auto relu_1307 = network->addActivation(*add_1306->getOutput(0), ActivationType::kRELU); + + auto relu_1310 = convBnRelu(network, weightMap, *relu_1108->getOutput(0), width, 3, 2, 1, "stage4.1.fuse_layers.2.0.0.0", "stage4.1.fuse_layers.2.0.0.1"); + auto bn_1312 = convBnRelu(network, weightMap, *relu_1310->getOutput(0), width * 4, 3, 2, 1, "stage4.1.fuse_layers.2.0.1.0", "stage4.1.fuse_layers.2.0.1.1", false); + auto add_1315 = convBnUpAdd(network, weightMap, *relu_1136->getOutput(0), *bn_1312->getOutput(0), width * 4, 3, 2, 1, + "stage4.1.fuse_layers.2.1.0.0", "stage4.1.fuse_layers.2.1.0.1", false); + auto add_1316 = network->addElementWise(*relu_1164->getOutput(0), *add_1315->getOutput(0), ElementWiseOperation::kSUM); + auto add_1338 = convBnUpAdd(network, weightMap, *relu_1192->getOutput(0), *add_1316->getOutput(0), width * 4, 1, 1, 0, + "stage4.1.fuse_layers.2.3.0", "stage4.1.fuse_layers.2.3.1", true); + auto relu_1339 = network->addActivation(*add_1338->getOutput(0), ActivationType::kRELU); + + auto relu_1342 = convBnRelu(network, weightMap, *relu_1108->getOutput(0), width, 3, 2, 1, "stage4.1.fuse_layers.3.0.0.0", "stage4.1.fuse_layers.3.0.0.1"); + auto relu_1345 = convBnRelu(network, weightMap, *relu_1342->getOutput(0), width, 3, 2, 1, "stage4.1.fuse_layers.3.0.1.0", "stage4.1.fuse_layers.3.0.1.1"); + auto bn_1347 = convBnRelu(network, weightMap, *relu_1345->getOutput(0), width * 8, 3, 2, 1, "stage4.1.fuse_layers.3.0.2.0", "stage4.1.fuse_layers.3.0.2.1", false); + auto relu_1350 = convBnRelu(network, weightMap, *relu_1136->getOutput(0), width * 2, 3, 2, 1, "stage4.1.fuse_layers.3.1.0.0", "stage4.1.fuse_layers.3.1.0.1"); + auto add_1353 = convBnUpAdd(network, weightMap, *relu_1350->getOutput(0), *bn_1347->getOutput(0), width * 8, 3, 2, 1, + "stage4.1.fuse_layers.3.1.1.0", "stage4.1.fuse_layers.3.1.1.1", false); + auto add_1356 = convBnUpAdd(network, weightMap, *relu_1164->getOutput(0), *add_1353->getOutput(0), width * 8, 3, 2, 1, + "stage4.1.fuse_layers.3.2.0.0", "stage4.1.fuse_layers.3.2.0.1", false); + auto add_1357 = network->addElementWise(*relu_1192->getOutput(0), *add_1356->getOutput(0), ElementWiseOperation::kSUM); + auto relu_1358 = network->addActivation(*add_1357->getOutput(0), ActivationType::kRELU); + + auto relu_1365 = liteResBlock(network, weightMap, *relu_1259->getOutput(0), width, "stage4.2.branches.0.0"); + auto relu_1372 = liteResBlock(network, weightMap, *relu_1365->getOutput(0), width, "stage4.2.branches.0.1"); + auto relu_1379 = liteResBlock(network, weightMap, *relu_1372->getOutput(0), width, "stage4.2.branches.0.2"); + auto relu_1386 = liteResBlock(network, weightMap, *relu_1379->getOutput(0), width, "stage4.2.branches.0.3"); + + auto relu_1393 = liteResBlock(network, weightMap, *relu_1307->getOutput(0), width * 2, "stage4.2.branches.1.0"); + auto relu_1400 = liteResBlock(network, weightMap, *relu_1393->getOutput(0), width * 2, "stage4.2.branches.1.1"); + auto relu_1407 = liteResBlock(network, weightMap, *relu_1400->getOutput(0), width * 2, "stage4.2.branches.1.2"); + auto relu_1414 = liteResBlock(network, weightMap, *relu_1407->getOutput(0), width * 2, "stage4.2.branches.1.3"); + + auto relu_1421 = liteResBlock(network, weightMap, *relu_1339->getOutput(0), width * 4, "stage4.2.branches.2.0"); + auto relu_1428 = liteResBlock(network, weightMap, *relu_1421->getOutput(0), width * 4, "stage4.2.branches.2.1"); + auto relu_1435 = liteResBlock(network, weightMap, *relu_1428->getOutput(0), width * 4, "stage4.2.branches.2.2"); + auto relu_1442 = liteResBlock(network, weightMap, *relu_1435->getOutput(0), width * 4, "stage4.2.branches.2.3"); + + auto relu_1449 = liteResBlock(network, weightMap, *relu_1358->getOutput(0), width * 8, "stage4.2.branches.3.0"); + auto relu_1456 = liteResBlock(network, weightMap, *relu_1449->getOutput(0), width * 8, "stage4.2.branches.3.1"); + auto relu_1463 = liteResBlock(network, weightMap, *relu_1456->getOutput(0), width * 8, "stage4.2.branches.3.2"); + auto relu_1470 = liteResBlock(network, weightMap, *relu_1463->getOutput(0), width * 8, "stage4.2.branches.3.3"); + + auto add_1492 = convBnUpAdd(network, weightMap, *relu_1414->getOutput(0), *relu_1386->getOutput(0), width, 1, 1, 0, + "stage4.2.fuse_layers.0.1.0", "stage4.2.fuse_layers.0.1.1", true); + auto add_1514 = convBnUpAdd(network, weightMap, *relu_1442->getOutput(0), *add_1492->getOutput(0), width, 1, 1, 0, + "stage4.2.fuse_layers.0.2.0", "stage4.2.fuse_layers.0.2.1", true); + + auto add_1536 = convBnUpAdd(network, weightMap, *relu_1470->getOutput(0), *add_1514->getOutput(0), width, 1, 1, 0, + "stage4.2.fuse_layers.0.3.0", "stage4.2.fuse_layers.0.3.1", true); + auto relu_1537 = network->addActivation(*add_1536->getOutput(0), ActivationType::kRELU); + + auto add_1540 = convBnUpAdd(network, weightMap, *relu_1386->getOutput(0), *relu_1414->getOutput(0), + width * 2, 3, 2, 1, "stage4.2.fuse_layers.1.0.0.0", "stage4.2.fuse_layers.1.0.0.1", false); + auto add_1562 = convBnUpAdd(network, weightMap, *relu_1442->getOutput(0), *add_1540->getOutput(0), + width * 2, 1, 1, 0, "stage4.2.fuse_layers.1.2.0", "stage4.2.fuse_layers.1.2.1", true); + auto add_1584 = convBnUpAdd(network, weightMap, *relu_1470->getOutput(0), *add_1562->getOutput(0), + width * 2, 1, 1, 0, "stage4.2.fuse_layers.1.3.0", "stage4.2.fuse_layers.1.3.1", true); + auto relu_1585 = network->addActivation(*add_1584->getOutput(0), ActivationType::kRELU); + + auto relu_1588 = convBnRelu(network, weightMap, *relu_1386->getOutput(0), width, 3, 2, 1, "stage4.2.fuse_layers.2.0.0.0", "stage4.2.fuse_layers.2.0.0.1"); + auto bn_1590 = convBnRelu(network, weightMap, *relu_1588->getOutput(0), width * 4, 3, 2, 1, "stage4.2.fuse_layers.2.0.1.0", "stage4.2.fuse_layers.2.0.1.1", false); + auto add_1593 = convBnUpAdd(network, weightMap, *relu_1414->getOutput(0), *bn_1590->getOutput(0), width * 4, 3, 2, 1, + "stage4.2.fuse_layers.2.1.0.0", "stage4.2.fuse_layers.2.1.0.1", false); + auto add_1594 = network->addElementWise(*relu_1442->getOutput(0), *add_1593->getOutput(0), ElementWiseOperation::kSUM); + auto add_1616 = convBnUpAdd(network, weightMap, *relu_1470->getOutput(0), *add_1594->getOutput(0), width * 4, 1, 1, 0, + "stage4.2.fuse_layers.2.3.0", "stage4.2.fuse_layers.2.3.1", true); + auto relu_1617 = network->addActivation(*add_1616->getOutput(0), ActivationType::kRELU); + + auto relu_1620 = convBnRelu(network, weightMap, *relu_1386->getOutput(0), width, 3, 2, 1, "stage4.2.fuse_layers.3.0.0.0", "stage4.2.fuse_layers.3.0.0.1"); + auto relu_1623 = convBnRelu(network, weightMap, *relu_1620->getOutput(0), width, 3, 2, 1, "stage4.2.fuse_layers.3.0.1.0", "stage4.2.fuse_layers.3.0.1.1"); + auto bn_1625 = convBnRelu(network, weightMap, *relu_1623->getOutput(0), width * 8, 3, 2, 1, "stage4.2.fuse_layers.3.0.2.0", "stage4.2.fuse_layers.3.0.2.1", false); + auto relu_1628 = convBnRelu(network, weightMap, *relu_1414->getOutput(0), width * 2, 3, 2, 1, "stage4.2.fuse_layers.3.1.0.0", "stage4.2.fuse_layers.3.1.0.1"); + auto add_1631 = convBnUpAdd(network, weightMap, *relu_1628->getOutput(0), *bn_1625->getOutput(0), width * 8, 3, 2, 1, + "stage4.2.fuse_layers.3.1.1.0", "stage4.2.fuse_layers.3.1.1.1", false); + auto add_1634 = convBnUpAdd(network, weightMap, *relu_1442->getOutput(0), *add_1631->getOutput(0), width * 8, 3, 2, 1, + "stage4.2.fuse_layers.3.2.0.0", "stage4.2.fuse_layers.3.2.0.1", false); + auto add_1635 = network->addElementWise(*relu_1470->getOutput(0), *add_1634->getOutput(0), ElementWiseOperation::kSUM); + auto relu_1636 = network->addActivation(*add_1635->getOutput(0), ActivationType::kRELU); + + nvinfer1::Dims dim = relu_1537->getOutput(0)->getDimensions(); + dim.d[0] = relu_1585->getOutput(0)->getDimensions().d[0]; + auto resize_1655 = netAddUpsampleBi(network, relu_1585->getOutput(0), dim); + dim.d[0] = relu_1617->getOutput(0)->getDimensions().d[0]; + auto resize_1668 = netAddUpsampleBi(network, relu_1617->getOutput(0), dim); + dim.d[0] = relu_1636->getOutput(0)->getDimensions().d[0]; + auto resize_1681 = netAddUpsampleBi(network, relu_1636->getOutput(0), dim); + + ITensor *concatTensors[] = {relu_1537->getOutput(0), resize_1655->getOutput(0), resize_1668->getOutput(0), resize_1681->getOutput(0)}; + auto concat_1682 = network->addConcatenation(concatTensors, 4); + concat_1682->setAxis(0); + auto relu_1685 = convBnRelu(network, weightMap, *concat_1682->getOutput(0), width * 15, 1, 1, 0, "aux_head.0", "aux_head.1", true, true); + auto conv_1686 = network->addConvolutionNd(*relu_1685->getOutput(0), NUM_CLASSES, DimsHW{1, 1}, weightMap["aux_head.3.weight"], weightMap["aux_head.3.bias"]); + conv_1686->setStrideNd(DimsHW{1, 1}); + conv_1686->setPaddingNd(DimsHW{0, 0}); + auto reshape_1701 = network->addShuffle(*conv_1686->getOutput(0)); + nvinfer1::Dims reshape_dim; + reshape_dim.nbDims = 2; + reshape_dim.d[0] = NUM_CLASSES; + reshape_dim.d[1] = -1; + reshape_1701->setReshapeDimensions(reshape_dim); + + auto softmax_1714 = network->addSoftMax(*reshape_1701->getOutput(0)); + softmax_1714->setAxes(2); + + auto relu_1689 = convBnRelu(network, weightMap, *concat_1682->getOutput(0), 512, 3, 1, 1, "conv3x3_ocr.0", "conv3x3_ocr.1", true, true); + + auto reshape_1710 = network->addShuffle(*relu_1689->getOutput(0)); + nvinfer1::Dims reshape_dim1; + reshape_dim1.nbDims = 2; + reshape_dim1.d[0] = 512; + reshape_dim1.d[1] = -1; + reshape_1710->setReshapeDimensions(reshape_dim1); + nvinfer1::Permutation permutation1; + permutation1.order[0] = 1; + permutation1.order[1] = 0; + reshape_1710->setSecondTranspose(permutation1); + + auto matmul_1715 = network->addMatrixMultiply(*softmax_1714->getOutput(0), MatrixOperation::kNONE, + *reshape_1710->getOutput(0), MatrixOperation::kNONE); + + auto transpose_1716 = network->addShuffle(*matmul_1715->getOutput(0)); + nvinfer1::Permutation permutation2; + permutation2.order[0] = 1; + permutation2.order[1] = 0; + transpose_1716->setFirstTranspose(permutation2); + + auto unsqueeze_1717 = network->addShuffle(*transpose_1716->getOutput(0)); + nvinfer1::Dims reshape_dim3; + reshape_dim3.nbDims = 3; + reshape_dim3.d[0] = 512; + reshape_dim3.d[1] = NUM_CLASSES; + reshape_dim3.d[2] = 1; + unsqueeze_1717->setReshapeDimensions(reshape_dim3); + + auto relu_1737 = convBnRelu(network, weightMap, *unsqueeze_1717->getOutput(0), 256, 1, 1, 0, "ocr_distri_head.object_context_block.f_object.0", "ocr_distri_head.object_context_block.f_object.1.0", true, true); + + auto relu_1740 = convBnRelu(network, weightMap, *relu_1737->getOutput(0), 256, 1, 1, 0, "ocr_distri_head.object_context_block.f_object.2", "ocr_distri_head.object_context_block.f_object.3.0", true, true); + + auto reshape_1747 = network->addShuffle(*relu_1740->getOutput(0)); + nvinfer1::Dims reshape_dim4; + reshape_dim4.nbDims = 2; + reshape_dim4.d[0] = 256; + reshape_dim4.d[1] = -1; + reshape_1747->setReshapeDimensions(reshape_dim4); + + auto relu_1723 = convBnRelu(network, weightMap, *relu_1689->getOutput(0), 256, 1, 1, 0, "ocr_distri_head.object_context_block.f_pixel.0", "ocr_distri_head.object_context_block.f_pixel.1.0", true, true); + auto relu_1726 = convBnRelu(network, weightMap, *relu_1723->getOutput(0), 256, 1, 1, 0, "ocr_distri_head.object_context_block.f_pixel.2", "ocr_distri_head.object_context_block.f_pixel.3.0", true, true); + + auto reshape_1733 = network->addShuffle(*relu_1726->getOutput(0)); + nvinfer1::Dims reshape_dim5; + reshape_dim5.nbDims = 2; + reshape_dim5.d[0] = 256; + reshape_dim5.d[1] = -1; + reshape_1733->setReshapeDimensions(reshape_dim5); + nvinfer1::Permutation permutation3; + permutation3.order[0] = 1; + permutation3.order[1] = 0; + reshape_1733->setSecondTranspose(permutation3); + + auto matmul_1759 = network->addMatrixMultiply(*reshape_1733->getOutput(0), MatrixOperation::kNONE, *reshape_1747->getOutput(0), MatrixOperation::kNONE); + nvinfer1::Dims constant_dim; + constant_dim.nbDims = 2; + int allNum = INPUT_H * INPUT_W / 16; + constant_dim.d[0] = INPUT_H * INPUT_W / 16; + constant_dim.d[1] = 1; + Weights wgt{DataType::kFLOAT, nullptr, allNum}; + float *w = new float[allNum]; + for (int i = 0; i < allNum; i++) + { + w[i] = 0.0625; + } + wgt.values = w; + auto constant_1761 = network->addConstant(constant_dim, wgt); + + auto mul_1761 = network->addElementWise(*constant_1761->getOutput(0), *matmul_1759->getOutput(0), ElementWiseOperation::kPROD); + + auto softmax_1762 = network->addSoftMax(*mul_1761->getOutput(0)); + softmax_1762->setAxes(2); + + auto relu_1750 = convBnRelu(network, weightMap, *unsqueeze_1717->getOutput(0), 256, 1, 1, 0, "ocr_distri_head.object_context_block.f_down.0", "ocr_distri_head.object_context_block.f_down.1.0", true, true); + + auto reshape_1757 = network->addShuffle(*relu_1750->getOutput(0)); + nvinfer1::Dims reshape_dim6; + reshape_dim6.nbDims = 2; + reshape_dim6.d[0] = 256; + reshape_dim6.d[1] = -1; + reshape_1757->setReshapeDimensions(reshape_dim6); + nvinfer1::Permutation permutation4; + permutation4.order[0] = 1; + permutation4.order[1] = 0; + reshape_1757->setSecondTranspose(permutation4); + + auto matmul_1763 = network->addMatrixMultiply(*softmax_1762->getOutput(0), MatrixOperation::kNONE, *reshape_1757->getOutput(0), MatrixOperation::kNONE); + + auto reshape_1777 = network->addShuffle(*matmul_1763->getOutput(0)); + nvinfer1::Dims reshape_dim7; + reshape_dim7.nbDims = 3; + reshape_dim7.d[0] = 256; + reshape_dim7.d[1] = INPUT_H / 4; + reshape_dim7.d[2] = INPUT_W / 4; + reshape_1777->setReshapeDimensions(reshape_dim7); + nvinfer1::Permutation permutation5; + permutation5.order[0] = 1; + permutation5.order[1] = 0; + reshape_1777->setFirstTranspose(permutation5); + + auto relu_1780 = convBnRelu(network, weightMap, *reshape_1777->getOutput(0), 512, 1, 1, 0, "ocr_distri_head.object_context_block.f_up.0", "ocr_distri_head.object_context_block.f_up.1.0", true, true); + + ITensor *concatTensors1[] = {relu_1780->getOutput(0), relu_1689->getOutput(0)}; + auto concat_1781 = network->addConcatenation(concatTensors1, 2); + + auto relu_1784 = convBnRelu(network, weightMap, *concat_1781->getOutput(0), 512, 1, 1, 0, "ocr_distri_head.conv_bn_dropout.0", "ocr_distri_head.conv_bn_dropout.1.0", true, true); + + auto conv_1785 = network->addConvolutionNd(*relu_1784->getOutput(0), NUM_CLASSES, DimsHW{1, 1}, weightMap["cls_head.weight"], weightMap["cls_head.bias"]); + debug_print(conv_1785->getOutput(0), "cls_head"); + dim.nbDims = 3; + dim.d[0] = NUM_CLASSES; + dim.d[1] = INPUT_H; + dim.d[2] = INPUT_W; + auto feature_map = netAddUpsampleBi(network, conv_1785->getOutput(0), dim); + debug_print(feature_map->getOutput(0), "upsample"); + auto topk = network->addTopK(*feature_map->getOutput(0), TopKOperation::kMAX, 1, 0X01); + + debug_print(topk->getOutput(0), "topk"); + + std::cout << "set name out" << std::endl; + topk->getOutput(1)->setName(OUTPUT_BLOB_NAME); + network->markOutput(*topk->getOutput(1)); + builder->setMaxBatchSize(maxBatchSize); + config->setMaxWorkspaceSize((1 << 30)); // 1G +#ifdef USE_FP16 + std::cout << "use fp16" << std::endl; + config->setFlag(BuilderFlag::kFP16); +#endif + ICudaEngine *engine = builder->buildEngineWithConfig(*network, *config); + std::cout << "build success!" << std::endl; + network->destroy(); + for (auto &mem : weightMap) + { + free((void *)(mem.second.values)); + } + return engine; +} +void APIToModel(unsigned int maxBatchSize, IHostMemory **modelStream, std::string wtsPath, int width) +{ + IBuilder *builder = createInferBuilder(gLogger); + IBuilderConfig *config = builder->createBuilderConfig(); + ICudaEngine *engine = createEngine(maxBatchSize, builder, config, DataType::kFLOAT, wtsPath, width); + assert(engine != nullptr); + (*modelStream) = engine->serialize(); + engine->destroy(); + builder->destroy(); +} + +bool parse_args(int argc, char **argv, std::string &wts, std::string &engine, int &width, std::string &img_dir) +{ + if (std::string(argv[1]) == "-s" && argc == 5) + { + wts = std::string(argv[2]); + engine = std::string(argv[3]); + width = std::stoi(argv[4]); + } + else if (std::string(argv[1]) == "-d" && argc == 4) + { + engine = std::string(argv[2]); + img_dir = std::string(argv[3]); + } + else + { + return false; + } + return true; +} +void doInference(IExecutionContext &context, cudaStream_t &stream, void **buffers, int batchSize) +{ + context.enqueue(batchSize, buffers, stream, nullptr); + cudaStreamSynchronize(stream); + cudaDeviceSynchronize(); +} + +int main(int argc, char **argv) +{ + cudaSetDevice(DEVICE); + std::string wtsPath = ""; + std::string engine_name = ""; + int width; + std::string img_dir; + // parse args + if (!parse_args(argc, argv, wtsPath, engine_name, width, img_dir)) + { + std::cerr << "arguments not right!" << std::endl; + std::cerr << "./hrnet_ocr -s [.wts] [.engine] [18 or 32 or 48] // serialize model to plan file" << std::endl; + std::cerr << "./hrnet_ocr -d [.engine] ../samples // deserialize plan file and run inference" << std::endl; + return -1; + } + // create a model using the API directly and serialize it to a stream + if (!wtsPath.empty()) + { + IHostMemory *modelStream{nullptr}; + APIToModel(BATCH_SIZE, &modelStream, wtsPath, width); + assert(modelStream != nullptr); + std::ofstream p(engine_name, std::ios::binary); + if (!p) + { + std::cerr << "could not open plan output file" << std::endl; + return -1; + } + p.write(reinterpret_cast(modelStream->data()), modelStream->size()); + modelStream->destroy(); + return 0; + } + + // deserialize the .engine and run inference + char *trtModelStream{nullptr}; + size_t size{0}; + std::ifstream file(engine_name, std::ios::binary); + if (file.good()) + { + file.seekg(0, file.end); + size = file.tellg(); + file.seekg(0, file.beg); + trtModelStream = new char[size]; + assert(trtModelStream); + file.read(trtModelStream, size); + file.close(); + } + else + { + std::cerr << "could not open plan file" << std::endl; + } + + std::vector file_names; + if (read_files_in_dir(img_dir.c_str(), file_names) < 0) + { + std::cout << "read_files_in_dir failed." << std::endl; + return -1; + } + // prepare input data --------------------------- + cudaSetDeviceFlags(cudaDeviceMapHost); + float *data; + int *prob; // using int. output is index + CHECK(cudaHostAlloc((void **)&data, BATCH_SIZE * 3 * INPUT_H * INPUT_W * sizeof(float), cudaHostAllocMapped)); + CHECK(cudaHostAlloc((void **)&prob, BATCH_SIZE * OUTPUT_SIZE * sizeof(int), cudaHostAllocMapped)); + + IRuntime *runtime = createInferRuntime(gLogger); + assert(runtime != nullptr); + ICudaEngine *engine = runtime->deserializeCudaEngine(trtModelStream, size); + assert(engine != nullptr); + IExecutionContext *context = engine->createExecutionContext(); + assert(context != nullptr); + delete[] trtModelStream; + void *buffers[2]; + // In order to bind the buffers, we need to know the names of the input and output tensors. + // Note that indices are guaranteed to be less than IEngine::getNbBindings() + const int inputIndex = engine->getBindingIndex(INPUT_BLOB_NAME); + const int outputIndex = engine->getBindingIndex(OUTPUT_BLOB_NAME); + assert(inputIndex == 0); + assert(outputIndex == 1); + cudaStream_t stream; + CHECK(cudaStreamCreate(&stream)); + + for (int f = 0; f < (int)file_names.size(); f++) + { + std::cout << file_names[f] << std::endl; + cv::Mat pr_img; + cv::Mat img_BGR = cv::imread(img_dir + "/" + file_names[f], 1); // BGR + cv::Mat img; + cv::cvtColor(img_BGR, img, cv::COLOR_BGR2RGB); + if (img.empty()) + continue; + cv::resize(img, pr_img, cv::Size(INPUT_W, INPUT_H)); + img = pr_img.clone(); // for img show + pr_img.convertTo(pr_img, CV_32FC3); + if (!pr_img.isContinuous()) + { + pr_img = pr_img.clone(); + } + std::memcpy(data, pr_img.data, BATCH_SIZE * 3 * INPUT_W * INPUT_H * sizeof(float)); + + cudaHostGetDevicePointer((void **)&buffers[inputIndex], (void *)data, 0); // buffers[inputIndex]-->data + cudaHostGetDevicePointer((void **)&buffers[outputIndex], (void *)prob, 0); // buffers[outputIndex] --> prob + + // Run inference + auto start = std::chrono::high_resolution_clock::now(); + doInference(*context, stream, buffers, BATCH_SIZE); + auto end = std::chrono::high_resolution_clock::now(); + std::cout << "infer time: " << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; + + cv::Mat outimg(INPUT_H, INPUT_W, CV_8UC1); + for (int row = 0; row < INPUT_H; ++row) + { + uchar *uc_pixel = outimg.data + row * outimg.step; + for (int col = 0; col < INPUT_W; ++col) + { + uc_pixel[col] = (uchar)prob[row * INPUT_W + col]; + } + } + cv::Mat im_color; + cv::cvtColor(outimg, im_color, cv::COLOR_GRAY2RGB); + cv::Mat lut = createLTU(NUM_CLASSES); + cv::LUT(im_color, lut, im_color); + // false color + cv::cvtColor(im_color, im_color, cv::COLOR_RGB2GRAY); + cv::applyColorMap(im_color, im_color, cv::COLORMAP_HOT); + // cv::imshow("False Color Map", im_color); + cv::imwrite(std::to_string(f) + "_false_color_map.png", im_color); + //fusion + cv::Mat fusionImg; + cv::addWeighted(img, 1, im_color, 0.8, 1, fusionImg); + // cv::imshow("Fusion Img", fusionImg); + // cv::waitKey(0); + cv::imwrite(std::to_string(f) + "_fusion_img.png", fusionImg); + } + + // Release stream and buffers + cudaStreamDestroy(stream); + CHECK(cudaFreeHost(buffers[inputIndex])); + CHECK(cudaFreeHost(buffers[outputIndex])); + // Destroy the engine + context->destroy(); + engine->destroy(); + runtime->destroy(); + return 0; +} diff --git a/hrnet/hrnet-semantic-segmentation/hrnetseg.cpp b/hrnet/hrnet-semantic-segmentation/hrnetseg.cpp deleted file mode 100644 index dd32d61..0000000 --- a/hrnet/hrnet-semantic-segmentation/hrnetseg.cpp +++ /dev/null @@ -1,404 +0,0 @@ -#include -#include -#include -#include -#include -#include -#include "common.hpp" -#include "logging.h" - -static Logger gLogger; -#define USE_FP16 -#define DEVICE 0 // GPU id -#define BATCH_SIZE 1 // only support 1 - -const char* INPUT_BLOB_NAME = "image"; -const char* OUTPUT_BLOB_NAME = "output"; -static const int INPUT_H = 512; -static const int INPUT_W = 1024; -static const int NUM_CLASSES = 19; -static const int OUTPUT_SIZE = INPUT_H * INPUT_W; - -// Creat the engine using only the API and not any parser. -ICudaEngine* createEngine(unsigned int maxBatchSize, IBuilder* builder, IBuilderConfig* config, DataType dt) { - 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{INPUT_H, INPUT_W, 3 }); - assert(data); - - // hwc to chw - auto ps = network->addShuffle(*data); - ps->setFirstTranspose(nvinfer1::Permutation{ 2, 0, 1 }); - // mean = [0.485, 0.456, 0.406] - // std = [0.229, 0.224, 0.225] - float mean[3] = { 0.406, 0.456, 0.485 }; - float std[3] = { 0.225, 0.224, 0.229 }; - ITensor* preinput = MeanStd(network, ps->getOutput(0), mean, std, true); - - // BGR to RGB - ISliceLayer *B = network->addSlice(*preinput, Dims3{ 0, 0, 0 }, Dims3{ 1, INPUT_H, INPUT_W }, Dims3{ 1, 1, 1 }); - ISliceLayer *G = network->addSlice(*preinput, Dims3{ 1, 0, 0 }, Dims3{ 1, INPUT_H, INPUT_W }, Dims3{ 1, 1, 1 }); - ISliceLayer *R = network->addSlice(*preinput, Dims3{ 2, 0, 0 }, Dims3{ 1, INPUT_H, INPUT_W }, Dims3{ 1, 1, 1 }); - - ITensor* inputTensors[] = { R->getOutput(0), G->getOutput(0), B->getOutput(0) }; - auto inputcat = network->addConcatenation(inputTensors, 3); - - std::map weightMap = loadWeights("../HRNetSeg.wts"); - Weights emptywts{ DataType::kFLOAT, nullptr, 0 }; - - auto id_876 = convBnLeaky(network, weightMap, *inputcat->getOutput(0), 64, 3, 2, 1, "conv1", "bn1"); - auto id_879 = convBnLeaky(network, weightMap, *id_876->getOutput(0), 64, 3, 2, 1, "conv2", "bn2"); //Res - auto id_891 = ResBlock2Conv(network, weightMap, *id_879->getOutput(0), 64, 256, 1, "layer1.0"); - auto id_901 = ResBlock(network, weightMap, *id_891->getOutput(0), 256, 64, 1, "layer1.1"); - nvinfer1::Dims dim1 = id_901->getOutput(0)->getDimensions(); - - auto id_904 = convBnLeaky(network, weightMap, *id_901->getOutput(0), 18, 3, 1, 1, "transition1.0.0", "transition1.0.1"); - auto id_914 = liteResBlock(network, weightMap, *id_904->getOutput(0), 18, "stage2.0.branches.0.0"); - auto id_921 = liteResBlock(network, weightMap, *id_914->getOutput(0), 18, "stage2.0.branches.0.1"); - - auto id_907 = convBnLeaky(network, weightMap, *id_901->getOutput(0), 36, 3, 2, 1, "transition1.1.0.0", "transition1.1.0.1"); - auto id_928 = liteResBlock(network, weightMap, *id_907->getOutput(0), 36, "stage2.0.branches.1.0"); - auto id_935 = liteResBlock(network, weightMap, *id_928->getOutput(0), 36, "stage2.0.branches.1.1"); - - auto id_957 = convBnUpAdd(network, weightMap, *id_935->getOutput(0), *id_921->getOutput(0), 18, 1, 1, 0, "stage2.0.fuse_layers.0.1.0", "stage2.0.fuse_layers.0.1.1", true); - auto id_958 = network->addActivation(*id_957->getOutput(0), ActivationType::kRELU); - - dim1 = id_935->getOutput(0)->getDimensions(); - dim1 = id_921->getOutput(0)->getDimensions(); - auto id_961 = convBnUpAdd(network, weightMap, *id_921->getOutput(0), *id_935->getOutput(0), 36, 3, 2, 1, "stage2.0.fuse_layers.1.0.0.0", "stage2.0.fuse_layers.1.0.0.1", false); - auto id_962 = network->addActivation(*id_961->getOutput(0), ActivationType::kRELU); - dim1 = id_962->getOutput(0)->getDimensions(); - - auto id_972 = liteResBlock(network, weightMap, *id_958->getOutput(0), 18, "stage3.0.branches.0.0"); - auto id_979 = liteResBlock(network, weightMap, *id_972->getOutput(0), 18, "stage3.0.branches.0.1"); - - auto id_986 = liteResBlock(network, weightMap, *id_962->getOutput(0), 36, "stage3.0.branches.1.0"); - auto id_993 = liteResBlock(network, weightMap, *id_986->getOutput(0), 36, "stage3.0.branches.1.1"); - - auto id_963 = convBnLeaky(network, weightMap, *id_962->getOutput(0), 72, 3, 2, 1, "transition2.2.0.0", "transition2.2.0.1"); - auto id_1000 = liteResBlock(network, weightMap, *id_963->getOutput(0), 72, "stage3.0.branches.2.0"); - auto id_1007 = liteResBlock(network, weightMap, *id_1000->getOutput(0), 72, "stage3.0.branches.2.1"); - - auto id_1029 = convBnUpAdd(network, weightMap, *id_993->getOutput(0), *id_979->getOutput(0), 18, 1, 1, 0, "stage3.0.fuse_layers.0.1.0", "stage3.0.fuse_layers.0.1.1", true); - auto id_1051 = convBnUpAdd(network, weightMap, *id_1007->getOutput(0), *id_1029->getOutput(0), 18, 1, 1, 0, "stage3.0.fuse_layers.0.2.0", "stage3.0.fuse_layers.0.2.1", true); - auto id_1052 = network->addActivation(*id_1051->getOutput(0), ActivationType::kRELU); - - auto id_1055 = convBnUpAdd(network, weightMap, *id_979->getOutput(0), *id_993->getOutput(0), 36, 3, 2, 1, "stage3.0.fuse_layers.1.0.0.0", "stage3.0.fuse_layers.1.0.0.1", false); - auto id_1077 = convBnUpAdd(network, weightMap, *id_1007->getOutput(0), *id_1055->getOutput(0), 36, 1, 1, 0, "stage3.0.fuse_layers.1.2.0", "stage3.0.fuse_layers.1.2.1", true); - auto id_1078 = network->addActivation(*id_1077->getOutput(0), ActivationType::kRELU); - - auto id_1081 = convBnLeaky(network, weightMap, *id_979->getOutput(0), 18, 3, 2, 1, "stage3.0.fuse_layers.2.0.0.0", "stage3.0.fuse_layers.2.0.0.1"); - auto id_1083 = convBnLeaky(network, weightMap, *id_1081->getOutput(0), 72, 3, 2, 1, "stage3.0.fuse_layers.2.0.1.0", "stage3.0.fuse_layers.2.0.1.1",false); - auto id_1086= convBnUpAdd(network, weightMap, *id_993->getOutput(0), *id_1083->getOutput(0), 72, 3, 2, 1, "stage3.0.fuse_layers.2.1.0.0", "stage3.0.fuse_layers.2.1.0.1", false); - auto id_1087 = network->addElementWise(*id_1086->getOutput(0), *id_1007->getOutput(0), ElementWiseOperation::kSUM); - auto id_1088 = network->addActivation(*id_1087->getOutput(0), ActivationType::kRELU); - - auto id_1095 = liteResBlock(network, weightMap, *id_1052->getOutput(0), 18, "stage3.1.branches.0.0"); - auto id_1102 = liteResBlock(network, weightMap, *id_1095->getOutput(0), 18, "stage3.1.branches.0.1"); - auto id_1109 = liteResBlock(network, weightMap, *id_1078->getOutput(0), 36, "stage3.1.branches.1.0"); - auto id_1116 = liteResBlock(network, weightMap, *id_1109->getOutput(0), 36, "stage3.1.branches.1.1"); - auto id_1123 = liteResBlock(network, weightMap, *id_1088->getOutput(0), 72, "stage3.1.branches.2.0"); - auto id_1130 = liteResBlock(network, weightMap, *id_1123->getOutput(0), 72, "stage3.1.branches.2.1"); - - auto id_1152 = convBnUpAdd(network, weightMap, *id_1116->getOutput(0), *id_1102->getOutput(0), 18, 1, 1, 0, "stage3.1.fuse_layers.0.1.0", "stage3.1.fuse_layers.0.1.1", true); - auto id_1174 = convBnUpAdd(network, weightMap, *id_1130->getOutput(0), *id_1152->getOutput(0), 18, 1, 1, 0, "stage3.1.fuse_layers.0.2.0", "stage3.1.fuse_layers.0.2.1", true); - auto id_1175 = network->addActivation(*id_1174->getOutput(0), ActivationType::kRELU); - - auto id_1178 = convBnUpAdd(network, weightMap, *id_1102->getOutput(0), *id_1116->getOutput(0), 36, 3, 2, 1, "stage3.1.fuse_layers.1.0.0.0", "stage3.1.fuse_layers.1.0.0.1", false); - auto id_1200 = convBnUpAdd(network, weightMap, *id_1130->getOutput(0), *id_1178->getOutput(0), 36, 1, 1, 0, "stage3.1.fuse_layers.1.2.0", "stage3.1.fuse_layers.1.2.1", true); - auto id_1201 = network->addActivation(*id_1200->getOutput(0), ActivationType::kRELU); - - auto id_1204 = convBnLeaky(network, weightMap, *id_1102->getOutput(0), 18, 3, 2, 1, "stage3.1.fuse_layers.2.0.0.0", "stage3.1.fuse_layers.2.0.0.1"); - auto id_1206 = convBnLeaky(network, weightMap, *id_1204->getOutput(0), 72, 3, 2, 1, "stage3.1.fuse_layers.2.0.1.0", "stage3.1.fuse_layers.2.0.1.1", false); - auto id_1209 = convBnUpAdd(network, weightMap, *id_1116->getOutput(0), *id_1206->getOutput(0), 72, 3, 2, 1, "stage3.1.fuse_layers.2.1.0.0", "stage3.1.fuse_layers.2.1.0.1", false); - auto id_1210 = network->addElementWise(*id_1209->getOutput(0), *id_1130->getOutput(0), ElementWiseOperation::kSUM); - auto id_1211 = network->addActivation(*id_1210->getOutput(0), ActivationType::kRELU); - - auto id_1218 = liteResBlock(network, weightMap, *id_1175->getOutput(0), 18, "stage3.2.branches.0.0"); - auto id_1225 = liteResBlock(network, weightMap, *id_1218->getOutput(0), 18, "stage3.2.branches.0.1"); - auto id_1232 = liteResBlock(network, weightMap, *id_1201->getOutput(0), 36, "stage3.2.branches.1.0"); - auto id_1239 = liteResBlock(network, weightMap, *id_1232->getOutput(0), 36, "stage3.2.branches.1.1"); - auto id_1246 = liteResBlock(network, weightMap, *id_1211->getOutput(0), 72, "stage3.2.branches.2.0"); - auto id_1253 = liteResBlock(network, weightMap, *id_1246->getOutput(0), 72, "stage3.2.branches.2.1"); - - auto id_1275 = convBnUpAdd(network, weightMap, *id_1239->getOutput(0), *id_1225->getOutput(0), 18, 1, 1, 0, "stage3.2.fuse_layers.0.1.0", "stage3.2.fuse_layers.0.1.1", true); - auto id_1297 = convBnUpAdd(network, weightMap, *id_1253->getOutput(0), *id_1275->getOutput(0), 18, 1, 1, 0, "stage3.2.fuse_layers.0.2.0", "stage3.2.fuse_layers.0.2.1", true); - auto id_1298 = network->addActivation(*id_1297->getOutput(0), ActivationType::kRELU); - - auto id_1301 = convBnUpAdd(network, weightMap, *id_1225->getOutput(0), *id_1239->getOutput(0), 36, 3, 2, 1, "stage3.2.fuse_layers.1.0.0.0", "stage3.2.fuse_layers.1.0.0.1", false); - auto id_1323 = convBnUpAdd(network, weightMap, *id_1253->getOutput(0), *id_1301->getOutput(0), 36, 1, 1, 0, "stage3.2.fuse_layers.1.2.0", "stage3.2.fuse_layers.1.2.1", true); - auto id_1324 = network->addActivation(*id_1323->getOutput(0), ActivationType::kRELU); - - auto id_1327 = convBnLeaky(network, weightMap, *id_1225->getOutput(0), 18, 3, 2, 1, "stage3.2.fuse_layers.2.0.0.0", "stage3.2.fuse_layers.2.0.0.1"); - auto id_1329 = convBnLeaky(network, weightMap, *id_1327->getOutput(0), 72, 3, 2, 1, "stage3.2.fuse_layers.2.0.1.0", "stage3.2.fuse_layers.2.0.1.1", false); - auto id_1332 = convBnUpAdd(network, weightMap, *id_1239->getOutput(0), *id_1329->getOutput(0), 72, 3, 2, 1, "stage3.2.fuse_layers.2.1.0.0", "stage3.2.fuse_layers.2.1.0.1", false); - auto id_1333 = network->addElementWise(*id_1332->getOutput(0), *id_1253->getOutput(0), ElementWiseOperation::kSUM); - auto id_1334 = network->addActivation(*id_1333->getOutput(0), ActivationType::kRELU); - - auto id_1344 = liteResBlock(network, weightMap, *id_1298->getOutput(0), 18, "stage4.0.branches.0.0"); - auto id_1351 = liteResBlock(network, weightMap, *id_1344->getOutput(0), 18, "stage4.0.branches.0.1"); - auto id_1358 = liteResBlock(network, weightMap, *id_1324->getOutput(0), 36, "stage4.0.branches.1.0"); - auto id_1365 = liteResBlock(network, weightMap, *id_1358->getOutput(0), 36, "stage4.0.branches.1.1"); - auto id_1372 = liteResBlock(network, weightMap, *id_1334->getOutput(0), 72, "stage4.0.branches.2.0"); - auto id_1379 = liteResBlock(network, weightMap, *id_1372->getOutput(0), 72, "stage4.0.branches.2.1"); - - auto id_1337 = convBnLeaky(network, weightMap, *id_1334->getOutput(0), 144, 3, 2, 1, "transition3.3.0.0", "transition3.3.0.1"); - auto id_1386 = liteResBlock(network, weightMap, *id_1337->getOutput(0), 144, "stage4.0.branches.3.0"); - auto id_1393 = liteResBlock(network, weightMap, *id_1386->getOutput(0), 144, "stage4.0.branches.3.1"); - - auto id_1415 = convBnUpAdd(network, weightMap, *id_1365->getOutput(0), *id_1351->getOutput(0), 18, 1, 1, 0, "stage4.0.fuse_layers.0.1.0", "stage4.0.fuse_layers.0.1.1", true); - auto id_1437 = convBnUpAdd(network, weightMap, *id_1379->getOutput(0), *id_1415->getOutput(0), 18, 1, 1, 0, "stage4.0.fuse_layers.0.2.0", "stage4.0.fuse_layers.0.2.1", true); - auto id_1459 = convBnUpAdd(network, weightMap, *id_1393->getOutput(0), *id_1437->getOutput(0), 18, 1, 1, 0, "stage4.0.fuse_layers.0.3.0", "stage4.0.fuse_layers.0.3.1", true); - auto id_1460 = network->addActivation(*id_1459->getOutput(0), ActivationType::kRELU); - - auto id_1463 = convBnUpAdd(network, weightMap, *id_1351->getOutput(0), *id_1365->getOutput(0), 36, 3, 2, 1, "stage4.0.fuse_layers.1.0.0.0", "stage4.0.fuse_layers.1.0.0.1", false); - auto id_1458 = convBnUpAdd(network, weightMap, *id_1379->getOutput(0), *id_1463->getOutput(0), 36, 1, 1, 0, "stage4.0.fuse_layers.1.2.0", "stage4.0.fuse_layers.1.2.1", true); - auto id_1507 = convBnUpAdd(network, weightMap, *id_1393->getOutput(0), *id_1458->getOutput(0), 36, 1, 1, 0, "stage4.0.fuse_layers.1.3.0", "stage4.0.fuse_layers.1.3.1", true); - auto id_1508 = network->addActivation(*id_1507->getOutput(0), ActivationType::kRELU); - - auto id_1511 = convBnLeaky(network, weightMap, *id_1351->getOutput(0), 18, 3, 2, 1, "stage4.0.fuse_layers.2.0.0.0", "stage4.0.fuse_layers.2.0.0.1"); - auto id_1513 = convBnLeaky(network, weightMap, *id_1511->getOutput(0), 72, 3, 2, 1, "stage4.0.fuse_layers.2.0.1.0", "stage4.0.fuse_layers.2.0.1.1", false); - auto id_1516 = convBnUpAdd(network, weightMap, *id_1365->getOutput(0), *id_1513->getOutput(0), 72, 3, 2, 1, "stage4.0.fuse_layers.2.1.0.0", "stage4.0.fuse_layers.2.1.0.1", true); - auto id_1517 = network->addElementWise(*id_1516->getOutput(0), *id_1379->getOutput(0), ElementWiseOperation::kSUM); - auto id_1539= convBnUpAdd(network, weightMap, *id_1393->getOutput(0), *id_1517->getOutput(0), 72, 1, 1, 0, "stage4.0.fuse_layers.2.3.0", "stage4.0.fuse_layers.2.3.1", true); - auto id_1540 = network->addActivation(*id_1539->getOutput(0), ActivationType::kRELU); - - auto id_1543 = convBnLeaky(network, weightMap, *id_1351->getOutput(0), 18, 3, 2, 1, "stage4.0.fuse_layers.3.0.0.0", "stage4.0.fuse_layers.3.0.0.1"); - auto id_1546 = convBnLeaky(network, weightMap, *id_1543->getOutput(0), 18, 3, 2, 1, "stage4.0.fuse_layers.3.0.1.0", "stage4.0.fuse_layers.3.0.1.1"); - auto id_1548 = convBnLeaky(network, weightMap, *id_1546->getOutput(0), 144, 3, 2, 1, "stage4.0.fuse_layers.3.0.2.0", "stage4.0.fuse_layers.3.0.2.1", false); - auto id_1551 = convBnLeaky(network, weightMap, *id_1365->getOutput(0), 36, 3, 2, 1, "stage4.0.fuse_layers.3.1.0.0", "stage4.0.fuse_layers.3.1.0.1"); - auto id_1554 = convBnUpAdd(network, weightMap, *id_1551->getOutput(0), *id_1548->getOutput(0), 144, 3, 2, 1, "stage4.0.fuse_layers.3.1.1.0", "stage4.0.fuse_layers.3.1.1.1", false); - auto id_1557 = convBnUpAdd(network, weightMap, *id_1379->getOutput(0), *id_1554->getOutput(0), 144, 3, 2, 1, "stage4.0.fuse_layers.3.2.0.0", "stage4.0.fuse_layers.3.2.0.1", false); - auto id_1558 = network->addElementWise(*id_1557->getOutput(0), *id_1393->getOutput(0), ElementWiseOperation::kSUM); - auto id_1559 = network->addActivation(*id_1558->getOutput(0), ActivationType::kRELU); - - auto id_1566= liteResBlock(network, weightMap, *id_1460->getOutput(0), 18, "stage4.1.branches.0.0"); - auto id_1573 = liteResBlock(network, weightMap, *id_1566->getOutput(0), 18, "stage4.1.branches.0.1"); - auto id_1580 = liteResBlock(network, weightMap, *id_1508->getOutput(0), 36, "stage4.1.branches.1.0"); - auto id_1587 = liteResBlock(network, weightMap, *id_1580->getOutput(0), 36, "stage4.1.branches.1.1"); - auto id_1594 = liteResBlock(network, weightMap, *id_1540->getOutput(0), 72, "stage4.1.branches.2.0"); - auto id_1601 = liteResBlock(network, weightMap, *id_1594->getOutput(0), 72, "stage4.1.branches.2.1"); - auto id_1608 = liteResBlock(network, weightMap, *id_1559->getOutput(0), 144, "stage4.1.branches.3.0"); - auto id_1615 = liteResBlock(network, weightMap, *id_1608->getOutput(0), 144, "stage4.1.branches.3.1"); - - auto id_1637 = convBnUpAdd(network, weightMap, *id_1587->getOutput(0), *id_1573->getOutput(0), 18, 1, 1, 0, "stage4.1.fuse_layers.0.1.0", "stage4.1.fuse_layers.0.1.1", true); - auto id_1659 = convBnUpAdd(network, weightMap, *id_1601->getOutput(0), *id_1637->getOutput(0), 18, 1, 1, 0, "stage4.1.fuse_layers.0.2.0", "stage4.1.fuse_layers.0.2.1", true); - auto id_1681 = convBnUpAdd(network, weightMap, *id_1615->getOutput(0), *id_1659->getOutput(0), 18, 1, 1, 0, "stage4.1.fuse_layers.0.3.0", "stage4.1.fuse_layers.0.3.1", true); - auto id_1682 = network->addActivation(*id_1681->getOutput(0), ActivationType::kRELU); - - auto id_1685 = convBnUpAdd(network, weightMap, *id_1573->getOutput(0), *id_1587->getOutput(0), 36, 3, 2, 1, "stage4.1.fuse_layers.1.0.0.0", "stage4.1.fuse_layers.1.0.0.1", false); - auto id_1707 = convBnUpAdd(network, weightMap, *id_1601->getOutput(0), *id_1685->getOutput(0), 36, 1, 1, 0, "stage4.1.fuse_layers.1.2.0", "stage4.1.fuse_layers.1.2.1", true); - auto id_1729 = convBnUpAdd(network, weightMap, *id_1615->getOutput(0), *id_1707->getOutput(0), 36, 1, 1, 0, "stage4.1.fuse_layers.1.3.0", "stage4.1.fuse_layers.1.3.1", true); - auto id_1730 = network->addActivation(*id_1729->getOutput(0), ActivationType::kRELU); - - auto id_1733 = convBnLeaky(network, weightMap, *id_1573->getOutput(0), 18, 3, 2, 1, "stage4.1.fuse_layers.2.0.0.0", "stage4.1.fuse_layers.2.0.0.1"); - auto id_1735 = convBnLeaky(network, weightMap, *id_1733->getOutput(0), 72, 3, 2, 1, "stage4.1.fuse_layers.2.0.1.0", "stage4.1.fuse_layers.2.0.1.1", false); - auto id_1738 = convBnUpAdd(network, weightMap, *id_1587->getOutput(0), *id_1735->getOutput(0), 72, 3, 2, 1, "stage4.1.fuse_layers.2.1.0.0", "stage4.1.fuse_layers.2.1.0.1", false); - auto id_1739 = network->addElementWise(*id_1601->getOutput(0), *id_1738->getOutput(0), ElementWiseOperation::kSUM); - auto id_1761 = convBnUpAdd(network, weightMap, *id_1615->getOutput(0), *id_1739->getOutput(0), 72, 1, 1, 0, "stage4.1.fuse_layers.2.3.0", "stage4.1.fuse_layers.2.3.1", true); - auto id_1762 = network->addActivation(*id_1761->getOutput(0), ActivationType::kRELU); - - auto id_1765 = convBnLeaky(network, weightMap, *id_1573->getOutput(0), 18, 3, 2, 1, "stage4.1.fuse_layers.3.0.0.0", "stage4.1.fuse_layers.3.0.0.1"); - auto id_1768 = convBnLeaky(network, weightMap, *id_1765->getOutput(0), 18, 3, 2, 1, "stage4.1.fuse_layers.3.0.1.0", "stage4.1.fuse_layers.3.0.1.1"); - auto id_1770 = convBnLeaky(network, weightMap, *id_1768->getOutput(0), 144, 3, 2, 1, "stage4.1.fuse_layers.3.0.2.0", "stage4.1.fuse_layers.3.0.2.1",false); - auto id_1773 = convBnLeaky(network, weightMap, *id_1587->getOutput(0), 36, 3, 2, 1, "stage4.1.fuse_layers.3.1.0.0", "stage4.1.fuse_layers.3.1.0.1"); - auto id_1776 = convBnUpAdd(network, weightMap, *id_1773->getOutput(0), *id_1770->getOutput(0), 144, 3, 2, 1, "stage4.1.fuse_layers.3.1.1.0", "stage4.1.fuse_layers.3.1.1.1", false); - auto id_1779 = convBnUpAdd(network, weightMap, *id_1601->getOutput(0), *id_1776->getOutput(0), 144, 3, 2, 1, "stage4.1.fuse_layers.3.2.0.0", "stage4.1.fuse_layers.3.2.0.1", false); - auto id_1780 = network->addElementWise(*id_1779->getOutput(0), *id_1615->getOutput(0), ElementWiseOperation::kSUM); - auto id_1781 = network->addActivation(*id_1780->getOutput(0), ActivationType::kRELU); - - nvinfer1::Dims dim = id_1682->getOutput(0)->getDimensions(); - dim.d[0] = id_1730->getOutput(0)->getDimensions().d[0]; - auto id_1730_up = netAddUpsampleBi(network, id_1730->getOutput(0), dim); - dim.d[0] = id_1762->getOutput(0)->getDimensions().d[0]; - auto id_1762_up = netAddUpsampleBi(network, id_1762->getOutput(0), dim); - dim.d[0] = id_1781->getOutput(0)->getDimensions().d[0]; - auto id_1781_up = netAddUpsampleBi(network, id_1781->getOutput(0), dim); - - ITensor* concatTensors[] = { id_1682->getOutput(0), id_1730_up ->getOutput(0), id_1762_up->getOutput(0), id_1781_up->getOutput(0) }; - auto id_1827 = network->addConcatenation(concatTensors, 4); - - dim1 = id_1827->getOutput(0)->getDimensions(); - auto id_1830 = convBnLeaky(network, weightMap, *id_1827->getOutput(0), 270, 1, 1, 0, "last_layer.0", "last_layer.1", true, true); - auto id_1831 = network->addConvolutionNd(*id_1830->getOutput(0), NUM_CLASSES, DimsHW{ 1,1 },weightMap["last_layer.3.weight"],weightMap["last_layer.3.bias"]); - id_1831->setStrideNd(DimsHW{ 1, 1 }); - id_1831->setPaddingNd(DimsHW{ 0, 0 }); - - - dim.d[0] = NUM_CLASSES; - dim.d[1] = INPUT_H; - dim.d[2] = INPUT_W; - auto id_1832 = netAddUpsampleBi(network, id_1831->getOutput(0), dim); - auto id_1833 = network->addTopK(*id_1832->getOutput(0), TopKOperation::kMAX, 1, 0X01); - // id_1833->getOutput(1) 1 is index - id_1833->getOutput(1)->setName(OUTPUT_BLOB_NAME); - std::cout << "set name out" << std::endl; - network->markOutput(*id_1833->getOutput(1)); - - // Build engine - builder->setMaxBatchSize(maxBatchSize); - config->setMaxWorkspaceSize((1 << 30)); // 1G -#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; -} -void APIToModel(unsigned int maxBatchSize, IHostMemory** modelStream) { - // 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 = createEngine(maxBatchSize, builder, config, DataType::kFLOAT); - assert(engine != nullptr); - - // Serialize the engine - (*modelStream) = engine->serialize(); - - // Close everything down - engine->destroy(); - builder->destroy(); -} - -void doInference(IExecutionContext& context, cudaStream_t& stream, void** buffers, int batchSize) { - const ICudaEngine& engine = context.getEngine(); - context.enqueue(batchSize, buffers, stream, nullptr); - cudaStreamSynchronize(stream); - cudaDeviceSynchronize(); -} - -int main(int argc, char** argv) { - - cudaSetDevice(DEVICE); - // create a model using the API directly and serialize it to a stream - char *trtModelStream{ nullptr }; - size_t size{ 0 }; - std::string engine_name = "hrnet_seg.engine"; - if (argc == 2 && std::string(argv[1]) == "-s") { - IHostMemory* modelStream{ nullptr }; - APIToModel(BATCH_SIZE, &modelStream); - assert(modelStream != nullptr); - std::ofstream p(engine_name, std::ios::binary); - if (!p) { - std::cerr << "could not open plan output file" << std::endl; - return -1; - } - p.write(reinterpret_cast(modelStream->data()), modelStream->size()); - modelStream->destroy(); - return 0; - } - else if (argc == 3 && std::string(argv[1]) == "-d") { - std::ifstream file(engine_name, std::ios::binary); - if (file.good()) { - file.seekg(0, file.end); - size = file.tellg(); - file.seekg(0, file.beg); - trtModelStream = new char[size]; - assert(trtModelStream); - file.read(trtModelStream, size); - file.close(); - } - } - else { - std::cerr << "arguments not right!" << std::endl; - std::cerr << "./hrnetseg -s // serialize model to plan file" << std::endl; - std::cerr << "./hrnetseg -d ../samples // deserialize plan file and run inference" << std::endl; - return -1; - } - - std::vector file_names; - if (read_files_in_dir(argv[2], file_names) < 0) { - std::cout << "read_files_in_dir failed." << std::endl; - return -1; - } - // prepare input data --------------------------- - cudaSetDeviceFlags(cudaDeviceMapHost); - float* data; - int* prob; // using int. output is index - CHECK(cudaHostAlloc((void **)&data, BATCH_SIZE * 3 * INPUT_H * INPUT_W * sizeof(float), cudaHostAllocMapped)); - CHECK(cudaHostAlloc((void **)&prob, BATCH_SIZE * OUTPUT_SIZE * sizeof(int), cudaHostAllocMapped)); - - IRuntime* runtime = createInferRuntime(gLogger); - assert(runtime != nullptr); - ICudaEngine* engine = runtime->deserializeCudaEngine(trtModelStream, size); - assert(engine != nullptr); - IExecutionContext* context = engine->createExecutionContext(); - assert(context != nullptr); - delete[] trtModelStream; - void* buffers[2]; - // In order to bind the buffers, we need to know the names of the input and output tensors. - // Note that indices are guaranteed to be less than IEngine::getNbBindings() - const int inputIndex = engine->getBindingIndex(INPUT_BLOB_NAME); - const int outputIndex = engine->getBindingIndex(OUTPUT_BLOB_NAME); - assert(inputIndex == 0); - assert(outputIndex == 1); - cudaStream_t stream; - CHECK(cudaStreamCreate(&stream)); - - for (int f = 0; f < (int)file_names.size(); f++) { - cv::Mat pr_img; - cv::Mat img = cv::imread(std::string(argv[2]) + "/" + file_names[f]); // BGR - if (img.empty()) continue; - cv::resize(img, pr_img, cv::Size(INPUT_W, INPUT_H)); - img = pr_img.clone(); // for img show - pr_img.convertTo(pr_img, CV_32FC3); - if (!pr_img.isContinuous()) - { - pr_img = pr_img.clone(); - } - std::memcpy(data, pr_img.data, BATCH_SIZE * 3 * INPUT_W * INPUT_H * sizeof(float)); - - cudaHostGetDevicePointer((void **)&buffers[inputIndex], (void *)data, 0); // buffers[inputIndex]-->data - cudaHostGetDevicePointer((void **)&buffers[outputIndex], (void *)prob, 0); // buffers[outputIndex] --> prob - - // Run inference - auto start = std::chrono::high_resolution_clock::now(); - doInference(*context, stream, buffers, BATCH_SIZE); - auto end = std::chrono::high_resolution_clock::now(); - std::cout << "infer time: " << std::chrono::duration_cast(end - start).count() << "ms" << std::endl; - - cv::Mat outimg(INPUT_H, INPUT_W, CV_8UC1); - for (int row = 0; row destroy(); - engine->destroy(); - runtime->destroy(); - return 0; -} \ No newline at end of file