MONAI
MONAI copied to clipboard
Add TensorRT conversion support for diffusion unet and hovernet
Is your feature request related to a problem? Please describe. MONAI supported the TensorRT model conversion since 1.2. However, when tried to export the diffusion unet and hovernet to the TensorRT engine-based torchscript, there are some issues that some features need to be added to MONAI core to fix.
- For the diffusion unet, the forward function needs two inputs, a tensor and an extra time step. While the trt_export API only supports one input shape and one dynamic batchsize now, these two parameters should be extended to support the multiple inputs case.
- For the hovernet, the output of the forward function is a dict, which is not supported very well by the torchscript and TensorRT conversion. Therefore, a optional list output with a predefined order would be a good solution. There is a fixed similar issue of the detection network.
Describe the solution you'd like
- Support the multiple inputs case by extending the
input_shapeanddynamic_batchsizeparameters fromSequence[int]toSequence[int] | Sequence[Sequence[int]]type, where each element in these two parameters stands for a shape requirement and a batch size range of one input. - Add a list output support option parameter for the hovernet like this PR and update the MONAI model zoo bundle accordingly.
Hi, its been a while since the last update on this. Has this been abandoned?