Error converting fine-tuned GPT-Neo model to ONNX

nlp
1

Hi, I am trying to convert a fine-tuned GPT-Neo (125M) model to ONNX using the code below:

from transformers import pipeline, convert_graph_to_onnx, GPTNeoForCausalLM, GPT2Tokenizer
from pathlib import Path
import torch

model_name = "EleutherAI/gpt-neo-125M" 
pipeline_name = "text-generation"

tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-125M", bos_token='<|startoftext|>',
                                                          eos_token='<|endoftext|>', pad_token='<|pad|>')
nlp = pipeline(pipeline_name, model=model_name, tokenizer=tokenizer)

with torch.no_grad():
    (
        input_names,
        output_names,
        dynamic_axes,
        tokens,
    ) = convert_graph_to_onnx.infer_shapes(nlp, "pt")
    ordered_input_names, model_args = convert_graph_to_onnx.ensure_valid_input(
        nlp.model, tokens, input_names
    )

model_name = 'gpt-neo'

predictor_path = './' + model_name 

model2 = GPTNeoForCausalLM.from_pretrained(predictor_path)

text = "I feel happy and "
prompt = f'<|startoftext|>Review: {text} Sentiment: <|endoftext|>'
encodings_dict = nlp.tokenizer(prompt, truncation=True, max_length=300, padding="max_length", return_tensors="pt")

torch.onnx.export(
    model2,
    (encodings_dict['input_ids'], encodings_dict['attention_mask']),
    'model_test.onnx',
    input_names=input_names,
    output_names=output_names,
    dynamic_axes=dynamic_axes,
    do_constant_folding=True,
    use_external_data_format=True, # Needed because of model size
    enable_onnx_checker=True,
    opset_version=13
)

But I get this error:


ValueError                                Traceback (most recent call last)
<ipython-input-15-024e093371a4> in <module>()
     43     use_external_data_format=True, # Needed because of model size
     44     enable_onnx_checker=True,
---> 45     opset_version=13
     46 )

3 frames
/usr/local/lib/python3.7/dist-packages/torch/onnx/utils.py in _validate_dynamic_axes(dynamic_axes, model, input_names, output_names)
   1300             for i, x in enumerate(value):
   1301                 if not isinstance(x, int):
-> 1302                     raise ValueError("The type of axis index is expected to be an integer")
   1303                 if x in value_dict:
   1304                     warnings.warn("Duplicate dynamic axis index {} was provided for input {}."

ValueError: The type of axis index is expected to be an integer

But if I remove dynamic_axes, I get this:

IndexError                                Traceback (most recent call last)
<ipython-input-16-89ddee8c10f8> in <module>()
     43     use_external_data_format=True, # Needed because of model size
     44     enable_onnx_checker=True,
---> 45     opset_version=13
     46 )

15 frames
/usr/local/lib/python3.7/dist-packages/transformers/models/gpt_neo/modeling_gpt_neo.py in forward(self, input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict)
    545             past_key_values = tuple([None] * len(self.h))
    546         else:
--> 547             past_length = past_key_values[0][0].size(-2)
    548 
    549         device = input_ids.device if input_ids is not None else inputs_embeds.device

IndexError: dimension specified as -2 but tensor has no dimensions

Can someone help me please?

2

Try FeatureExtractionPipeline instead of a regular one:

from transformers import FeatureExtractionPipeline, AutoModel, AutoTokenizer
model_pipeline = FeatureExtractionPipeline(
    model=AutoModel.from_pretrained(full_model_path_name),
    tokenizer=AutoTokenizer.from_pretrained(
        full_model_path_name, use_fast=True
    ),
    framework="pt",
    device=-1,
)

At least it worked for me.