Open tangy5 opened 1 year ago
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, glad to see this topic in discussion. My understanding from #5125 is that there is some work around using
torch.jit.trace
ortorch.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
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.itertools
,einops
needs to replaced.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.