0

我参考Bert Quantization with ONNX Runtime on CPU,并尝试使用以下代码将dialogpt转换为 onnx:

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model = AutoModelForCausalLM.from_pretrained("DialoGPT-medium/")
model.to("cpu")

def export_onnx_model(args, model, onnx_model_path):
    # with torch.no_grad():
    model.eval()
    inputs = {'input_ids':      torch.ones(1,128, dtype=torch.int64),
                'attention_mask': torch.ones(1,128, dtype=torch.int64),
                'token_type_ids': torch.ones(1,128, dtype=torch.int64)}
    outputs = model(**inputs)

    symbolic_names = {0: 'batch_size', 1: 'max_seq_len'}
    torch.onnx.export(model,                                            
                (inputs['input_ids'],                                             
                inputs['attention_mask'], 
                inputs['token_type_ids']),                                                         
                onnx_model_path,                                                
                opset_version=10,                                                 
                do_constant_folding=True,                                         
                input_names=['input_ids',                        
                            'input_mask', 
                            'segment_ids'],
                output_names=['output'],                    
                dynamic_axes={'input_ids': symbolic_names,                            
                            'input_mask' : symbolic_names,
                            'segment_ids' : symbolic_names})

export_onnx_model(configs, model, "dialogpt.onnx")

但我遇到错误:

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

我怎样才能解决这个问题?

4

0 回答 0