Adding model metadata in TorchScript model file

[CoderHam]

🚀 Feature

Adding model metadata in TorchScript model file including:

1. Number of inputs and outputs
2. Datatype of inputs and outputs
3. Shape of inputs and outputs

It would greatly simply and improve model deployment and error handling for TorchScript models and add to the stability of the C++ Backend.

Motivation

We are working on the Triton Inference Server and support PyTorch (Libtorch) as one of the backends. The TorchScript model file (*.pt) does not contain sufficient metadata about the model such as the number, shape and datatypes of inputs and outputs. Our current workaround is for the user to specify this information manually.

As I understand, the shape data is not retained to allow for the model to be more generic. It would however be useful for us to have as much information about the model available in the model file itself. Without this we require the user to manually specify these details which can often be tedious and prone to error.

Pitch

Both TensorFlow and ONNX models store such information and it makes it easier to verify this information and deploy these models.

In addition to this it might be useful to have an option for shape specific information to be retained. In cases where certain dimensions are allowed to be variable these dimensions could be represented with a dummy value such as -1. All this information would allow us to validate that the data being passed to the model for inference is indeed valid.

Alternatives

I noted that there is a way to add additional metadata manually to the model but that too is a manual process.

Additional context

I had a discussion on the PyTorch Slack where I was informed that there exists function inputs() and outputs() [At least in Python] for the model that contains information about the number and datatype of the inputs however this required the @torch.jit.script decorator. Since there are two ways to export a model to torchscript (the second one using tracing) it could be useful and more user friendly to have this metadata always added to the model file (as in TensorFlow and ONNX) since it would add negligible size.

cc @suo

[suo]

hey! PyTorch does not impose requirements of dtype/shape for inputs, so I'm not sure we can generically capture such information in the model metadata. Number of inputs conceivably we could, although it is also trivially derivable from the deserialized model (graph.inputs().size()).

I suppose we could require that the user supply information on dtype/shape, but it doesn't seem very user friendly, given that the information may not be known to the user upfront, and it is not required to run the model normally.