PyTorch NNAPI integration prototype (#46780)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/46780
This is in prototype status, but pretty functional. There are two major
parts.
- Model converter. This is a pure Python component that consumes a
model in TorchScript format, converts the operations into NNAPI
semantics, and serializes the model in a custom format. It then wraps
the result in a new TorchScript model that can invoke NNAPI under the
hood.
- Runtime. This is a TorchBind object that deserializes the model and
sends the result to NNAPI. This is fairly simple since the serialized
format is basically just a list of NNAPI calls to make, so most of the
code is spent on bounds checking.
A few notes on the design.
- Currently, all tensor sizes need to be fixed, and those fixed sizes
are burned directly into the serialized model. This will probably
need to change. NNAPI supports variable-sized tensors, but the
important hardware backends do not. However, we're seeing use cases
crop up where the input size is not known until around the time that
the model is loaded (for example, it might depend on the camera aspect
ratio). I think the proper fix here is to remove the code in the
converter that eagerly calculates the sizes of the intermediate
tensors and replace it with a code generator that will generate some
TorchScript code that will perform those calculations at model load
time. This way, we will be able to support models that have
variable-sized inputs while still only showing fixed-sized operands to
NNAPI.
- The important hardware backends want operands to be in NHWC order, but
PyTorch natively represents all tensors and NCHW. The strategy for
this is to keep NCHW during most of the conversion process, but track
and additional value per operand representing the "dimension order".
The dimension order gets propagated through convolutions and pointwise
ops. When we're ready to serialize the model, we reorder the
dimensions for "channels last" operands to NHWC.
Test Plan:
Some local testing with FB prod models. I'll need to add some examples
and automated tests.
Reviewed By: iseeyuan
Differential Revision: D24574040
Pulled By: dreiss
fbshipit-source-id: 6adc8571b234877ee3666ec0c0de24da35c38a1f
