SwinUNETR for TensorRT

[tangy5]

Current SwinUNETR is not scriptable by torch.jit.script, several features has some limitations to make SwinUNETR compatible with TensorRT.

This issue can discuss the potential solution for making SwinUENTR support TensorRT.

A potential solution is to make a 'light' version of this network, so we won't significantly impact users with a transition.

• checkpoint option can be removed, as inference pipleline is mostly used and do not require checkpointing.
• some packages such as itertools, einops needs to replaced.
• Functions not scriptable such as Slice() need to be rewrite.

Forward functions that contains none basic python can be modified to support torch.jit.script. A light version of SwinUNETR can be a subclass of current SwinUENTR network, when users uses SwinUNETR from bundle, it uses the light version which works with TensorRT. The light version SwinUNETR need to be tested and evaluated for comparable performance to the current version.

Welcome comment on other options. There can be a better choice.

[csheaff]

Hi, glad to see this topic in discussion. My understanding from #5125 is that there is some work around using torch.jit.trace or torch.onnx.export. Is this not the case?

[binliunls]

Hi, glad to see this topic in discussion. My understanding from #5125 is that there is some work around using torch.jit.trace or torch.onnx.export. Is this not the case?

Hi @csheaff, Currently, the swin_unetr model in the model-zoo is already supported by the torch.jit.trace conversion way. This ticket is for converting the swin_unetr to a TensorRT engine-based torchscript, which is not supported for now. We are working with the TensorRT team on it and hopefully will support it in the future version.

Thanks, Bin