Repository: pytorch
Issue: Add the capability to export GradMultiply to ONNX
Labels: module: onnx,triaged

### 🚀 The feature, motivation and pitch

My team is working on a model whose upstream model is `DistilHuBERT`, and we'd like to convert this to ONNX-compatible format. This is to ensure that our work can be implemented to various platforms.

When we used `torch.onnx.export()`, we observe that the operator `GradMultiply` cannot be converted. We'd love to have this conversion as a feature in PyTorch.

Here are some bits and pieces to better outline our situation:

- The principal error message: `RuntimeError: ONNX export failed: Couldn't export Python operator GradMultiply`

- Full stack trace:

``` # start
Traceback (most recent call last):
  File "convert_distulhubert.py", line 29, in <module>
    output_names = ['output'])
  File "/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/onnx/__init__.py", line 276, in export
    custom_opsets, enable_onnx_checker, use_external_data_format)
  File "/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/onnx/utils.py", line 94, in export
    use_external_data_format=use_external_data_format)
  File "/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/onnx/utils.py", line 712, in _export
    val_add_node_names, val_use_external_data_format, model_file_location)
RuntimeError: ONNX export failed: Couldn't export Python operator GradMultiply

Defined at:
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/s3prl/upstream/distiller/model.py(172): forward_feature
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/s3prl/upstream/distiller/model.py(191): forward
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/nn/modules/module.py(860): _slow_forward
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/nn/modules/module.py(887): _call_impl
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/s3prl/upstream/distiller/builder.py(92): _forward
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/s3prl/upstream/distiller/builder.py(141): forward
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/nn/modules/module.py(860): _slow_forward
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/nn/modules/module.py(887): _call_impl
/Users/mathieum/ml/prispeech/api/backend/models/distilhubert.py(55): forward
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/nn/modules/module.py(860): _slow_forward
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/nn/modules/module.py(887): _call_impl
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/jit/_trace.py(116): wrapper
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/jit/_trace.py(130): forward
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/nn/modules/module.py(889): _call_impl
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/jit/_trace.py(1139): _get_trace_graph
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/onnx/utils.py(380): _trace_and_get_graph_from_model
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/onnx/utils.py(420): _create_jit_graph
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/onnx/utils.py(459): _model_to_graph
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/onnx/utils.py(701): _export
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/onnx/utils.py(94): export
/Users/mathieum/opt/anaconda3/envs/pytorch_lite_conversion/lib/python3.7/site-packages/torch/onnx/__init__.py(276): export
convert_distulhubert.py(29): <module>


Graph we tried to export:
graph(%input : Float(1, 1000, strides=[1000, 1], requires_grad=0, device=cpu),
      %model.model.feature_extractor.conv_layers.0.0.weight : Float(512, 1, 10, strides=[10, 10, 1], requires_grad=1, device=cpu),
...
```

- And here is the source code:


``` # start
import torch
import torchvision
import config as config
from models import DownstreamModel
import torchaudio
from torch.utils.mobile_optimizer import optimize_for_mobile
​
distilhubert_upstream = config.MODELS['distilhubert']['model_cls'](
        checkpoint_file=config.MODELS['distilhubert']['checkpoint_file'])
distilhubert_downstream = DownstreamModel(
    num_labels=config.MODELS['distilhubert']['num_labels'],
    upstream_dim=config.MODELS['distilhubert']['upstream_dim'],
    checkpoint_file=config.MODELS['distilhubert']['checkpoint_file']
)
​
model = distilhubert_upstream
#model = distilhubert_downstream
​
model.eval()
​
torch.onnx.export(model,              # model being run
                  torch.rand((1, 1000)),                         # model input (or a tuple for multiple inputs)
                  "distilhubert_upstream.onnx", # where to save the model (can be a file or file-like object)
                  export_params=True,        # store the trained parameter weights inside the model file
                  opset_version=12,          # the ONNX version to export the model to
                  do_constant_folding=True,  # whether to execute constant folding for optimization
                  input_names = ['input'],   # the model's input names
                  output_names = ['output'])
``` 

### Alternatives

We thought of [pytorchjs](https://github.com/raghavmecheri/pytorchjs) and [TorchSharp](https://github.com/dotnet/TorchSharp), but these would serve only to a limited set of use cases, and it would require us to re-define a new architecture every time we consider a new model. Conversion to ONNX through PyTorch appears to be the most practical and reusable implementation.

### Additional context

_No response_