bert pytorch模型转onnx,并改变输入输出

#Exa

mple Codes: def createNewGraph(inFile): """ https://github.com/microsoft/onnxruntime/issues/11783 """ model = onnx.load(inFile) graph_def = model.graph opset_imports = [onnx.helper.make_opsetid(domain="", version=17), onnx.helper.make_opsetid('com.microsoft', 1)] model_def = onnx.helper.make_model( graph_def, producer_name='test', opset_imports=opset_imports) return model_def def bert(): """ https://huggingface.co/bert-base-uncased git lfs install git clone https://huggingface.co/bert-base-uncased pip install onnxruntime==1.13.1 #3 inputs-->embedlayernormalization pip install onnxruntime==1.10.0 #2 inputs-->embedlayernormalization, 1 input--> pip install transformers==4.20.1 """ import torch # pip install transformers from transformers import BertTokenizer, BertModel from torch.autograd import Variable from onnxruntime import SessionOptions, InferenceSession, GraphOptimizationLevel from onnxruntime.transformers.fusion_options import FusionOptions from onnxruntime.transformers.optimizer import optimize_by_onnxruntime, optimize_model from onnx import helper tokenizer = BertTokenizer.from_pretrained('bert-base-uncased') model = BertModel.from_pretrained("bert-base-uncased") input_ids = Variable(torch.randint(0, 30522, (1,512))) token_types_ids = Variable(torch.zeros((1,512), dtype=int)) attention_mask = Variable(torch.zeros((1,512), dtype=int)) axes = {0: "batch_size", 1: "seq_len"} axes_o1 = {0: "batch_size"} torch.onnx.export( model, (input_ids, attention_mask , token_types_ids), #(input_ids, None, token_types_ids), 'bert-base-uncased.onnx', #export_params=True, do_constant_folding=True, input_names=["input_ids", "attention_mask", "token_type_ids"], #input_names=["input_ids", "token_type_ids"], output_names=["output_start_logits", "output_end_logits"], dynamic_axes={ "input_ids": axes, "attention_mask": axes, "token_type_ids": axes, "output_start_logits": axes, "output_end_logits": axes_o1, }, verbose=True, opset_version=17, ) in_path = 'bert-base-uncased.onnx' fusion_options = FusionOptions('bert') fusion_options.enable_bias_gelu = False fusion_options.enable_skip_layer_norm = False fusion_options.enable_bias_skip_layer_norm = False m = optimize_model( in_path, model_type='bert', num_heads=12, hidden_size=768, opt_level=0, optimization_options=fusion_options, use_gpu=False ) print(m.get_fused_operator_statistics()) m.save_model_to_file('bert-base-uncased_optimized.onnx', use_external_data_format=False) model_final = createNewGraph('bert-base-uncased_optimized.onnx') onnx.save(model_final, 'bert-base-uncased_final.onnx') #----------------------------------write shape: def changeInDim(model, shape): inputs = model.graph.input for i, input in enumerate(inputs): for j, dim in enumerate(input.type.tensor_type.shape.dim): print(f"dim[{i}][{j}]: {dim}--{dim.dim_value}-->{shape[i][j]}") dim.dim_value = shape[i][j] def changeOtDim(model, shape): outputs = model.graph.output for i, input in enumerate(outputs): for j, dim in enumerate(input.type.tensor_type.shape.dim): print(f"dim[{i}][{j}]: {dim}--{dim.dim_value}-->{shape[i][j]}") dim.dim_value = shape[i][j] def run_changeInOut(): inFile = "bert-base-uncased_final.onnx" outFile = "bert-base-uncased_shaped.onnx" model = onnx.load(inFile) changeInDim(model, ([1, 384], [1, 384], [1, 384])) changeOtDim(model, ([1, 384, 768], [1, 768])) inferred_model = onnx.shape_inference.infer_shapes(model) onnx.checker.check_model(inferred_model) onnx.save(inferred_model, outFile)