History

Wang Xinyu cb9efbdf21 unet: fix bug and coding style (#1171 ) * unet: fix bug * update readme * remove useless code		2022-12-13 22:37:41 +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: fix bug and coding style (#1171 )	2022-12-13 22:37:41 +08:00
unet.cpp	unet: fix bug and coding style (#1171 )	2022-12-13 22:37:41 +08:00

README.md

UNet

Pytorch model from Pytorch-UNet.

Contributors

Requirements

Please use TensorRT 7.x.

There is a bug with TensorRT 8.x, we are working on it.

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.