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

add hrnet semantic segmentation model w18, w32 and w48 (#503)

Browse Source 
* create psenet

create psenet with weight from tensorflow

* delete some useless code

* repalce tab with 4 blanks

* fix network bug, rewrite post-processing pse algorithm

* update readme

* update readme

* add RepVGG

* fix typo

* add hrnetseg w18 w32 w48

* add hrnetseg with ocr w18 w32 w48

* merge hrnet and small, add hrnet_ocr

* fix warning

* change project name
This commit is contained in:
weiwei zhou 2021-04-29 20:26:15 +08:00 committed by GitHub
parent b58a7987bc
commit 1f7672ce5a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
8 changed files with 1623 additions and 790 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
hrnet/hrnet-semantic-segmentation/CMakeLists.txt
View File

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

71
hrnet/hrnet-semantic-segmentation/README.md
View File

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

429
hrnet/hrnet-semantic-segmentation/common.hpp
View File

@ -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<std::string> &file_names) {
@ -45,9 +45,21 @@ int read_files_in_dir(const char *p_dir_name, std::vector<std::string> &file_nam
return 0;
}
// TensorRT weight files have a simple space delimited format:
// [type] [size] <data x size in hex>
std::map<std::string, Weights> 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<std::string, Weights> loadWeights(const std::string file)
{
std::cout << "Loading weights: " << file << std::endl;
std::map<std::string, Weights> weightMap;
@ -62,7 +74,7 @@ std::map<std::string, Weights> 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<std::string, Weights> loadWeights(const std::string file) {
wt.type = DataType::kFLOAT;
// Load blob
uint32_t* val = reinterpret_cast<uint32_t*>(malloc(sizeof(val) * size));
uint32_t *val = reinterpret_cast<uint32_t *>(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<std::string, Weights> 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<float*>(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<float *>(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<std::string, Weights>& 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<std::string, Weights> &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<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
float *scval = reinterpret_cast<float *>(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<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
float *shval = reinterpret_cast<float *>(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<float*>(malloc(sizeof(float) * len));
for (int i = 0; i < len; i++) {
float *pval = reinterpret_cast<float *>(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<std::string, Weights>& 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<std::string, Weights> &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<std::string, Weights>& 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<std::string, Weights> &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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights>& 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<std::string, Weights> &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<std::string, Weights> &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<std::string, Weights> &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<std::string, Weights>& 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<std::string, Weights> &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;
}
}

175
hrnet/hrnet-semantic-segmentation/demo.py
View File

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

60
hrnet/hrnet-semantic-segmentation/gen_wts.py Normal file
View File

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

569
hrnet/hrnet-semantic-segmentation/hrnet.cpp Normal file
View File

@ -0,0 +1,569 @@
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
#include <chrono>
#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<std::string, Weights> 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<const char *>(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<std::string> 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<std::chrono::milliseconds>(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;
}

690
hrnet/hrnet-semantic-segmentation/hrnet_ocr.cpp Normal file
View File

@ -0,0 +1,690 @@
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
#include <chrono>
#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<std::string, Weights> 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<const char *>(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<std::string> 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<std::chrono::milliseconds>(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;
}

404
hrnet/hrnet-semantic-segmentation/hrnetseg.cpp
View File

@ -1,404 +0,0 @@
#include <fstream>
#include <iostream>
#include <map>
#include <sstream>
#include <vector>
#include <chrono>
#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<std::string, Weights> 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<const char*>(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<std::string> 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<std::chrono::milliseconds>(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);
//fusion
cv::Mat fusionImg;
cv::addWeighted(img, 1, im_color, 0.5, 1, fusionImg);
cv::imshow("Fusion Img", fusionImg);
cv::waitKey(0);
}
// 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;
}