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

UNet support TensorRT 8.x (#1245)

Browse Source 
* TensorRT 8.x supported

* Unet TensorRT 8.x supported, add Contributors nengwp

---------

Co-authored-by: nengwp <nengwp@github.ai>
This commit is contained in:
nengwp 2023-02-28 10:23:40 +08:00 committed by GitHub
parent f92dcf43dc
commit ef22b1d03b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 5 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
unet/README.md
View File

@ -8,12 +8,13 @@ Pytorch model from [Pytorch-UNet](https://github.com/milesial/Pytorch-UNet).
<a href="https://github.com/East-Face"><img src="https://avatars.githubusercontent.com/u/35283869?v=4s=48" width="40px;" alt=""/></a>
<a href="https://github.com/irvingzhang0512"><img src="https://avatars.githubusercontent.com/u/22089207?s=48&v=4" width="40px;" alt=""/></a>
<a href="https://github.com/wang-xinyu"><img src="https://avatars.githubusercontent.com/u/15235574?s=48&v=4" width="40px;" alt=""/></a>
<a href="https://github.com/nengwp"><img src="https://avatars.githubusercontent.com/u/44516353?s=96&v=4" width="40px;" alt=""/></a>
## Requirements
Please use TensorRT 7.x.
There is a bug with TensorRT 8.x, we are working on it.
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

1
unet/unet.cpp
View File

@ -41,6 +41,7 @@ ILayer* doubleConv(INetworkDefinition* network, std::map<std::string, Weights>&
ILayer* down(INetworkDefinition* network, std::map<std::string, Weights>& weightMap, ITensor& input, int outch, int p, std::string lname) {
IPoolingLayer* pool1 = network->addPoolingNd(input, PoolingType::kMAX, DimsHW{ 2, 2 });
pool1->setStrideNd(DimsHW{ 2, 2 });
assert(pool1);
ILayer* dcov1 = doubleConv(network, weightMap, *pool1->getOutput(0), outch, 3, lname + ".maxpool_conv.1", outch);
assert(dcov1);