TRT wrappers for Vista3D
Description
Adds TRT acceleration. Depends on https://github.com/Project-MONAI/MONAI/pull/7990
Types of changes
- [x] Non-breaking change (fix or new feature that would not break existing functionality).
- [ ] Breaking change (fix or new feature that would cause existing functionality to change).
- [ ] New tests added to cover the changes.
- [ ] In-line docstrings updated.
@heyufan1995 @Nic-Ma Please help review this draft PR. Vista3D TRT end-to-end is about 10% faster than PyTorch.
@yiheng-wang-nv please also help review it. Thanks.
The perf gain from TRT end-to-end is about 20% on infer_everything sample I ran (18s vs 22s): python3 -m scripts.infer --config_file 'configs/infer.yaml' - infer_everything --image_file 'example-1.nii.gz'
Hi @heyufan1995 ,
Could you please help review this PR?
Thanks in advance.
Hi, @bfomitchev, I notice that you have added TRT wrapper in latest version.
Could you please help adopt previous method for both supports of TensorRT and PyTorch.
For instance,
PyTorch uses
python3 -m scripts.infer --config_file 'configs/infer.yaml' - infer_everything --image_file 'example-1.nii.gz'
TensorRT uses
python3 -m scripts.export --config_file 'configs/infer.yaml' - infer_everything --image_file 'example-1.nii.gz'
It's fine for users to infer even without the TensorRT environment.
Thanks.
@mingxueg-nv : I have removed export.py in the latest version to avoid duplication. TRT use is being controlled by "trt" config flag in infer.yaml. I will add graceful fallback for the case when TRT is not installed.
@mingxueg-nv : I have removed export.py in the latest version to avoid duplication. TRT use is being controlled by "trt" config flag in infer.yaml. I will add graceful fallback for the case when TRT is not installed.
Thanks.
I will measure accuracy today to supplement more benchmark details.
we can enable trt version in model zoo after this PR merge cc @heyufan1995
Now depends on https://github.com/Project-MONAI/MONAI/pull/7990