title: how do you generate an ONNX representation of a pytorch bert pretrained neural network?
body: <p>I am trying to generate an ONNX file for the <a href="https://github.com/huggingface/pytorch-pretrained-BERT" rel="nofollow noreferrer">pytorch-pretrained-bert</a> run_classifier.py example. </p>
  
  <p>In this case, I'm running it with the following parameters as per the main README.md:</p>
  
  <pre class="lang-py prettyprint-override"><code>export GLUE_DIR=/tmp/glue_data
  
  python run_classifier.py \
    --task_name MRPC \
    --do_train \
    --do_eval \
    --do_lower_case \
    --data_dir $GLUE_DIR/MRPC/ \
    --bert_model bert-base-uncased \
    --max_seq_length 128 \
    --train_batch_size 32 \
    --learning_rate 2e-5 \
    --num_train_epochs 3.0 \
    --output_dir /tmp/mrpc_output/
  </code></pre>
  
  <p>The following code is modified/added at line 552:</p>
  
  <pre class="lang-py prettyprint-override"><code>    # Save a trained model
      model_to_save = model.module if hasattr(model, 'module') else model  # Only save the model it-self
      output_model_file = os.path.join(args.output_dir, "pytorch_model.bin")
      if args.do_train:
          torch.save(model_to_save.state_dict(), output_model_file)
  
      # Save ONNX
      msl = args.max_seq_length
      dummy_input = torch.randn(1, msl, msl, msl, num_labels, device="cpu")
      output_onnx_file = os.path.join(args.output_dir, "classifier.onnx")
      torch.onnx.export(model, dummy_input, output_onnx_file)
  </code></pre>
  
  <p>The dummy_input is supposed to correspond to the bert pretrained model inputs. A sample_batch_size of 1 seems suitable for my needs, I think.</p>
  
  <p>A few suggest that the arguments should match the arguments to the models forward() method. In this case:</p>
  
  <pre class="lang-py prettyprint-override"><code>class BertForSequenceClassification(PreTrainedBertModel):
      def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None):
  </code></pre>
  
  <p>In this case, I believe the ranks are:</p>
  
  <pre><code>input_ids: 1 x 128 &lt;the max_seq_length specified in the args&gt;
  token_type_ids: 1 x max_seq_length
  attention_mask: 1 x max_seq_length
  labels: 1 x 2 &lt;the number of labels for MRPC&gt;
  </code></pre>
  
  <p>so the effective call is:</p>
  
  <pre class="lang-py prettyprint-override"><code>    dummy_input = torch.randn(1, 128, 128, 128, 2, device="cpu")
  </code></pre>
  
  <p>Unfortunately this produces an error:</p>
  
  <pre><code>Exception has occurred: RuntimeError
  The expanded size of the tensor (2) must match the existing size (128) at non-singleton dimension 4.  Target sizes: [1, 128, 128, 128, 2].  Tensor sizes: [1, 128]
  </code></pre>
  
  <p>It seems likely that it's something fairly straightforward. Suggestions appreciated!</p>