Support Dict[str, Union[List[str], torch.Tensor]] as torchscript model input type
voe09 opened this issue · 2 comments
Neuropod has supported torch 1.10.2 in its latest release. In Torch 1.10.2, torch.jit introduced the support for Union type. That means it can support Dict[str, Union[List[str], torch.Tensor]] as its model input type.
Currently in neuropod, torchscript backend can only support Dict[str, List[str]] or Dict[str, torch.Tensor]. Related code is https://github.com/uber/neuropod/blob/master/source/neuropod/backends/torchscript/torch_backend.cc#L354
We would like to request support for Dict[str, Union[List[str], torch.Tensor]]
As we talked about offline, this sounds good to me.
Some additional thoughts that you should feel free to ignore (or we should create a separate issue about if they sound useful):
- It would be great to also include the option of a model taking multiple dicts as inputs instead of just one (so if a model takes in a string dict and a tensor dict, Neuropod will supply both). This would have the benefit of supporting older versions of Torch as well. It should be a fairly straightforward change to implement this.
- Neuropod already supports several output types (
Tensor,Dict[str, Tensor], named tuples,Dict[str, List[str]], tuple of dicts to support multiple return types, etc). Adding the ability to returnDict[str, Union[List[str], torch.Tensor]]would be nice if we are going to support that type as an input
These are definitely not required as part of this task (and should not all be put in the same PR), but some things to consider if you're interested in improving the TorchScript model interface.
Permalink version of the code snippet linked in original issue above:
neuropod/source/neuropod/backends/torchscript/torch_backend.cc
Lines 354 to 355 in 1097df8
Implemented in #543