This repo implemtents HRNet-Semantic-Segmentation-v1.1 with TensorRT7 network definition API
The code refers to https://github.com/wang-xinyu/tensorrtx/tree/master/hrnet
If you want to know more about HRNet, please click HRNet-segmentation的网络结构分析
- generate .wts, use config
experiments/cityscapes/seg_hrnet_w48_train_512x1024_sgd_lr1e-2_wd5e-4_bs_12_epoch484.yaml
and pretrained weighthrnet_w48_cityscapes_cls19_1024x2048_trainset.pth
as example. changePRETRAINED
inexperiments/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
- cmake and make
mkdir build
cd build
cmake ..
make
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
TRT Result:
pytorch result: