History

nengwp ef22b1d03b UNet support TensorRT 8.x (#1245 ) * TensorRT 8.x supported * Unet TensorRT 8.x supported, add Contributors nengwp --------- Co-authored-by: nengwp <nengwp@github.ai>		2023-02-28 10:23:40 +08:00
..
CMakeLists.txt	unet: fix bug and coding style (#1171 )	2022-12-13 22:37:41 +08:00
common.hpp	unet: fix bug and coding style (#1171 )	2022-12-13 22:37:41 +08:00
gen_wts.py	unet: fix bug and coding style (#1171 )	2022-12-13 22:37:41 +08:00
logging.h	unet support trt8	2022-11-25 21:13:49 +08:00
macros.h	unet support trt8	2022-11-25 21:13:49 +08:00
README.md	UNet support TensorRT 8.x (#1245 )	2023-02-28 10:23:40 +08:00
unet.cpp	UNet support TensorRT 8.x (#1245 )	2023-02-28 10:23:40 +08:00

README.md

UNet

Pytorch model from Pytorch-UNet.

Contributors

Requirements

Now TensorRT 8.x is supported and you can use it. The key cause of the previous bug is the pooling layer Stride setting problem.

Build and Run

Generate .wts

cp {path-of-tensorrtx}/unet/gen_wts.py Pytorch-UNet/
cd Pytorch-UNet/
wget https://github.com/milesial/Pytorch-UNet/releases/download/v3.0/unet_carvana_scale0.5_epoch2.pth
python gen_wts.py unet_carvana_scale0.5_epoch2.pth

Generate TensorRT engine

cd tensorrtx/unet/
mkdir build
cd build
cmake ..
make
cp {path-of-Pytorch-UNet}/unet.wts .
./unet -s

Run inference

wget https://raw.githubusercontent.com/wang-xinyu/tensorrtx/f60dcc7bec28846cd973fc95ac829c4e57a11395/unet/samples/0cdf5b5d0ce1_01.jpg
./unet -d 0cdf5b5d0ce1_01.jpg

Check result.jpg

Benchmark

Pytorch	TensorRT FP32	TensorRT FP16
816x672	816x672	816x672
58ms	43ms (batchsize 8)	14ms (batchsize 8)

More Information

See the readme in home page.