Error exporting BERT + classification layer to ONNX

username: Kroshtan
cooked: <p>I have the following model:</p>
<pre><code class="lang-auto">class BertClassifier(nn.Module):
    """
    Class defining the classifier model with a BERT encoder and a single fully connected classifier layer.
    """
    def __init__(self, dropout=0.5, num_labels=24):
        super(BertClassifier, self).__init__()

        self.bert = BertModel.from_pretrained('bert-base-uncased')
        self.dropout = nn.Dropout(dropout)
        self.linear = nn.Linear(768, num_labels)
        self.relu = nn.ReLU()
        self.best_score = 0

    def forward(self, input_id, mask):
        _, pooled_output = self.bert(input_ids=input_id, attention_mask=mask, return_dict=False)
        output = self.relu(self.linear(self.dropout(pooled_output)))

        return output
</code></pre>
<p>Using these inputs:</p>
<pre><code class="lang-auto">    ex_string = "example string"
    inputs = tokenizer(ex_string,
                       padding='max_length', max_length=512, truncation=True,
                       return_tensors="pt")
    input_id = inputs['input_ids'].squeeze(1)
    mask = inputs['attention_mask']
</code></pre>
<p>And I export the model to ONNX using:</p>
<pre><code class="lang-auto">torch.onnx.export(model, (input_id, mask), 'tryout.onnx', export_params=True, do_constant_folding=True)
</code></pre>
<p>Which results in the following stack trace:</p>
<pre><code class="lang-auto">/.local/lib/python3.9/site-packages/torch/onnx/symbolic_helper.py:325: UserWarning: Type cannot be inferred, which might cause exported graph to produce incorrect results.
  warnings.warn("Type cannot be inferred, which might cause exported graph to produce incorrect results.")
[W shape_type_inference.cpp:434] Warning: Constant folding in symbolic shape inference fails: index_select(): Index is supposed to be a vector
Exception raised from index_select_out_cpu_ at ../aten/src/ATen/native/TensorAdvancedIndexing.cpp:887 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x42 (0x7fc288ab8d62 in /.local/lib/python3.9/site-packages/torch/lib/libc10.so)
frame #1: at::native::index_select_out_cpu_(at::Tensor const&amp;, long, at::Tensor const&amp;, at::Tensor&amp;) + 0x3a9 (0x7fc2cd9e5189 in /.local/lib/python3.9/site-packages/torch/lib/libtorch_cpu.so)
frame #2: at::native::index_select_cpu_(at::Tensor const&amp;, long, at::Tensor const&amp;) + 0xe6 (0x7fc2cd9e7146 in /.local/lib/python3.9/site-packages/torch/lib/libtorch_cpu.so)
frame #3: &lt;unknown function&gt; + 0x1d37f12 (0x7fc2ce0def12 in /.local/lib/python3.9/site-packages/torch/lib/libtorch_cpu.so)
frame #4: at::_ops::index_select::redispatch(c10::DispatchKeySet, at::Tensor const&amp;, long, at::Tensor const&amp;) + 0xb9 (0x7fc2cdc7a099 in /.local/lib/python3.9/site-packages/torch/lib/libtorch_cpu.so)
frame #5: &lt;unknown function&gt; + 0x3250ac3 (0x7fc2cf5f7ac3 in /.local/lib/python3.9/site-packages/torch/lib/libtorch_cpu.so)
frame #6: &lt;unknown function&gt; + 0x32510f5 (0x7fc2cf5f80f5 in /.local/lib/python3.9/site-packages/torch/lib/libtorch_cpu.so)
frame #7: at::_ops::index_select::call(at::Tensor const&amp;, long, at::Tensor const&amp;) + 0x166 (0x7fc2cdcf9ce6 in /.local/lib/python3.9/site-packages/torch/lib/libtorch_cpu.so)
frame #8: torch::jit::onnx_constant_fold::runTorchBackendForOnnx(torch::jit::Node const*, std::vector&lt;at::Tensor, std::allocator&lt;at::Tensor&gt; &gt;&amp;, int) + 0x1b5f (0x7fc34fd5d6ff in /.local/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame #9: &lt;unknown function&gt; + 0xbbdc22 (0x7fc34fda4c22 in /.local/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame #10: torch::jit::ONNXShapeTypeInference(torch::jit::Node*, std::map&lt;std::string, c10::IValue, std::less&lt;std::string&gt;, std::allocator&lt;std::pair&lt;std::string const, c10::IValue&gt; &gt; &gt; const&amp;, int) + 0xa8e (0x7fc34fdaa46e in /.local/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame #11: &lt;unknown function&gt; + 0xbc4f74 (0x7fc34fdabf74 in /.local/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame #12: &lt;unknown function&gt; + 0xb35730 (0x7fc34fd1c730 in /.local/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame #13: &lt;unknown function&gt; + 0x2a5d8b (0x7fc34f48cd8b in /.local/lib/python3.9/site-packages/torch/lib/libtorch_python.so)
frame #14: python3() [0x53a8eb]
&lt;omitting python frames&gt;
frame #17: python3() [0x50f5e9]
frame #20: python3() [0x50f5e9]
frame #23: python3() [0x50f5e9]
frame #26: python3() [0x50f5e9]
frame #29: python3() [0x50f5e9]
frame #32: python3() [0x50f5e9]
frame #35: python3() [0x608ebb]
frame #36: python3() [0x603ea4]
frame #37: python3() [0x60834d]
frame #41: &lt;unknown function&gt; + 0x2dfd0 (0x7fc352298fd0 in /lib/x86_64-linux-gnu/libc.so.6)
frame #42: __libc_start_main + 0x7d (0x7fc35229907d in /lib/x86_64-linux-gnu/libc.so.6)
 (function ComputeConstantFolding)
Traceback (most recent call last):
  File "/bert_extraction/bert_onnx.py", line 71, in &lt;module&gt;
    torch.onnx.export(model, (input_id, mask), 'tryout.onnx', export_params=True, do_constant_folding=True)
  File "/.local/lib/python3.9/site-packages/torch/onnx/__init__.py", line 316, in export
    return utils.export(model, args, f, export_params, verbose, training,
  File "/.local/lib/python3.9/site-packages/torch/onnx/utils.py", line 107, in export
    _export(model, args, f, export_params, verbose, training, input_names, output_names,
  File "/.local/lib/python3.9/site-packages/torch/onnx/utils.py", line 724, in _export
    _model_to_graph(model, args, verbose, input_names,
  File "/.local/lib/python3.9/site-packages/torch/onnx/utils.py", line 544, in _model_to_graph
    params_dict = torch._C._jit_pass_onnx_constant_fold(graph, params_dict,
IndexError: index_select(): Index is supposed to be a vector
</code></pre>
<p>I can get the ONNX model to compile when I change the <code>do_constant_folding</code> flag to <code>False</code>, but obviously I don’t want to do that, as I’m trying to optimize the inference-time.<br>
Can anyone shed some light on the error or what I’m doing wrong?</p>