ONNX Export with Multiple Arguments: Only tuples, lists and Variables are supported as JIT inputs/outputs

ashraf_g: <p>Hi, everyone!</p>
<h3>
<a name="issue-summary" class="anchor" href="#issue-summary"></a>Issue Summary:</h3>
<p>I’m trying to convert the ESM-1b protein transformer model from PyTorch to ONNX. I run into an issue when I want to provide an extra argument to the model.</p>
<h3>
<a name="conversion-steps-code-and-error-message" class="anchor" href="#conversion-steps-code-and-error-message"></a>Conversion Steps, Code, and Error Message:</h3>
<p>Here is my conversion script (named <code>convert_onnx_esm.py</code>):</p>
<pre><code class="lang-python">import os
import torch
import torch.onnx
import argparse
from esm.pretrained import load_model_and_alphabet_local


parser = argparse.ArgumentParser()

parser.add_argument("--model-path", type=str, required=True)
parser.add_argument("--converted-model-path", type=str, required=True)
commandline_args = parser.parse_args()

model, alphabet = load_model_and_alphabet_local(commandline_args.model_path)
batch_converter = alphabet.get_batch_converter()

data = [
    ("protein1", "VLAGG"),
    ("protein2", "KALTARQ"),
]
batch_labels, batch_strs, batch_tokens = batch_converter(data)

# an example forward pass would be: model(batch_tokens, repr_layers=[33])
# the conversion works if (batch_tokens, [33]) is changed to just: batch_tokens
with torch.no_grad():
    torch.onnx.export(model,
        (batch_tokens, [33]),  
        commandline_args.converted_model_path,
        use_external_data_format=True,
        opset_version=11,
        do_constant_folding=True,
        input_names=["inputs"],
        output_names=["outputs"],
        dynamic_axes={"inputs": [0, 1]}
    )
</code></pre>
<p>Which I run like follows:</p>
<pre><code class="lang-bash">export MODEL_PATH=/tmp/models/esm/esm1b_t33_650M_UR50S.pt
export CONVERTED_GRAPH_PATH=/tmp/models/onnx_esm/graph.onnx

mkdir -p $(dirname $MODEL_PATH) $(dirname $CONVERTED_GRAPH_PATH)
curl https://dl.fbaipublicfiles.com/fair-esm/models/esm1b_t33_650M_UR50S.pt --output /tmp/models/esm/esm1b_t33_650M_UR50S.pt  # this is a 7Gb file

python convert_onnx_esm.py --model-path $MODEL_PATH --converted-model-path $CONVERTED_GRAPH_PATH
</code></pre>
<p>And the error I get after running the conversion script:</p>
<pre><code class="lang-auto">RuntimeError: Only tuples, lists and Variables are supported as JIT inputs/outputs. Dictionaries and strings are also accepted, but their usage is not recommended. Here, received an input of unsupported type: int
</code></pre>
<h3>
<a name="environment" class="anchor" href="#environment"></a>Environment:</h3>
<p>Cuda 11.2, CudNN 8.1.1.33, and Python 3.8.5 with packages:</p>
<pre><code class="lang-python">fair-esm==0.3.0
onnx==1.8.1
onnxconverter-common==1.6.0
onnxruntime-gpu==1.7.0
onnxruntime-tools==1.6.0
torch==1.9.0.dev20210318
</code></pre>
<h3>
<a name="additional-info" class="anchor" href="#additional-info"></a>Additional Info:</h3>
<p>To no avail, I also tried changing the arguments <code>(batch_tokens, [33])</code> to the dictionary format:</p>
<pre><code class="lang-python">(batch_tokens, {"tokens": batch_tokens, "repr_layers": [33]})
</code></pre>
<p>And finally, in case it’s helpful, the <code>nn.Module.forward()</code> method definition starts like this:</p>
<pre><code class="lang-python">def forward(self, tokens, repr_layers=[], need_head_weights=False, return_contacts=False)
</code></pre>
<p>Thank you for any tips or pointers!</p>

ashraf_g: <p>I realized that passing the argument as a list is not valid. A tensor is needed, like this:</p>
<pre><code class="lang-python">torch.onnx.export(model,
    (batch_tokens, torch.tensor([33])),
    converted_model_path,
    use_external_data_format=True,
    ...
)
</code></pre>
<p>Unfortunately, <code>torch.tensor([33]))</code> is as good as <code>[]</code> as far as the model’s behavior is concerned.</p>