Question

我正在尝试为pytorch-pretrained-bert run_classifier.py示例生成ONNX文件。

在这种情况下，我根据主README.md使用以下参数运行它：

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/

在第552行修改/添加了以下代码：

    # 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)

dummy_input应该对应于bert预训练模型输入。我认为sample_batch_size为1似乎很适合我的需求。

一些建议参数应与模型forward（）方法的参数匹配。在这种情况下：

class BertForSequenceClassification(PreTrainedBertModel):
    def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None):

在这种情况下，我相信排名是：

input_ids: 1 x 128 <the max_seq_length specified in the args>
token_type_ids: 1 x max_seq_length
attention_mask: 1 x max_seq_length
labels: 1 x 2 <the number of labels for MRPC>

所以有效的呼叫是：

    dummy_input = torch.randn(1, 128, 128, 128, 2, device="cpu")

不幸的是，这会产生错误：

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]

这似乎很简单。建议表示赞赏！

如何生成pytorch bert预训练神经网络的ONNX表示形式？

0 个答案: