def create_longformer_inputs(onnx_model_path, batch_size, sequence_length, global_length, samples): from onnx import TensorProto from onnx_model import OnnxModel onnx_model = OnnxModel(onnx.load(onnx_model_path)) symbols = {'batch_size': batch_size, 'sequence_length': sequence_length} dummy_inputs = {} for input in onnx_model.get_graph_inputs_excluding_initializers(): shape = get_shape_from_type_proto(input.type) for i, dim in enumerate(shape): if type(dim) == str and dim not in symbols.keys(): raise RuntimeError(f"symbol is not supported: {dim}") else: shape[i] = symbols[dim] elem_type = input.type.tensor_type.elem_type assert elem_type in [ TensorProto.FLOAT, TensorProto.INT32, TensorProto.INT64 ] data_type = numpy.float32 if elem_type == TensorProto.FLOAT else ( numpy.int64 if elem_type == TensorProto.INT64 else numpy.int32) if "global" in input.name: data = numpy.zeros(shape, dtype=data_type) data[:, :global_length] = 1 else: data = numpy.ones(shape, dtype=data_type) dummy_inputs[input.name] = data all_inputs = [dummy_inputs for _ in range(samples)] return all_inputs
def create_gpt2_inputs(onnx_model_path, batch_size, sequence_length, past_sequence_length, samples): from onnx import TensorProto from onnx_model import OnnxModel onnx_model = OnnxModel(onnx.load(onnx_model_path)) # The symbolic name shall be same as those used in Gpt2Helper.export_onnx(...) function. symbols = { 'batch_size': batch_size, 'seq_len': sequence_length, 'past_seq_len': past_sequence_length, 'total_seq_len': sequence_length + past_sequence_length } dummy_inputs = {} for input in onnx_model.get_graph_inputs_excluding_initializers(): shape = get_shape_from_type_proto(input.type) for i, dim in enumerate(shape): if type(dim) == str and dim not in symbols.keys(): raise RuntimeError(f"symbol is not supported: {dim}") else: shape[i] = symbols[dim] elem_type = input.type.tensor_type.elem_type assert elem_type in [ TensorProto.FLOAT, TensorProto.INT32, TensorProto.INT64 ] data_type = numpy.float32 if elem_type == TensorProto.FLOAT else ( numpy.int64 if elem_type == TensorProto.INT64 else numpy.int32) data = numpy.ones(shape, dtype=data_type) dummy_inputs[input.name] = data all_inputs = [dummy_inputs for _ in range(samples)] return all_inputs
def create_dummy_inputs(onnx_model_path, batch_size, sequence_length, samples): from onnx import TensorProto from onnx_model import OnnxModel onnx_model = OnnxModel(onnx.load(onnx_model_path)) dummy_inputs = {} for input in onnx_model.get_graph_inputs_excluding_initializers(): shape = get_shape_from_type_proto(input.type) symbol_dims = [] for i, dim in enumerate(shape): if type(dim) == str: symbol_dims.append(i) # allowed symbolic dimensions: batch_size and sequence_length if len(symbol_dims) > 2: return None if len(symbol_dims) > 0: shape[symbol_dims[0]] = batch_size if len(symbol_dims) > 1: shape[symbol_dims[1]] = sequence_length elem_type = input.type.tensor_type.elem_type assert elem_type in [ TensorProto.FLOAT, TensorProto.INT32, TensorProto.INT64 ] data_type = numpy.float32 if elem_type == TensorProto.FLOAT else ( numpy.int64 if elem_type == TensorProto.INT64 else numpy.int32) data = numpy.ones(shape, dtype=data_type) dummy_inputs[input.name] = data all_inputs = [dummy_inputs for _ in range(samples)] return all_inputs
def get_longformer_inputs(onnx_file, input_ids_name=None, input_mask_name=None, global_mask_name=None): """ Get graph inputs for longformer model. """ model = ModelProto() with open(onnx_file, "rb") as f: model.ParseFromString(f.read()) onnx_model = OnnxModel(model) graph_inputs = onnx_model.get_graph_inputs_excluding_initializers() if input_ids_name is not None: input_ids = onnx_model.find_graph_input(input_ids_name) if input_ids is None: raise ValueError( f"Graph does not have input named {input_ids_name}") input_mask = None if input_mask_name: input_mask = onnx_model.find_graph_input(input_mask_name) if input_mask is None: raise ValueError( f"Graph does not have input named {input_mask_name}") global_mask = None if global_mask_name: global_mask = onnx_model.find_graph_input(global_mask_name) if global_mask is None: raise ValueError( f"Graph does not have input named {global_mask_name}") expected_inputs = 1 + (1 if input_mask else 0) + (1 if global_mask else 0) if len(graph_inputs) != expected_inputs: raise ValueError( f"Expect the graph to have {expected_inputs} inputs. Got {len(graph_inputs)}" ) return input_ids, input_mask, global_mask if len(graph_inputs) != 3: raise ValueError("Expect the graph to have 3 inputs. Got {}".format( len(graph_inputs))) # Try guess the inputs based on naming. input_ids = None input_mask = None global_mask = None for input in graph_inputs: input_name_lower = input.name.lower() if "global" in input_name_lower: global_mask = input elif "mask" in input_name_lower: input_mask = input else: input_ids = input if input_ids and input_mask and global_mask: return input_ids, input_mask, global_mask raise ValueError( "Fail to assign 3 inputs. You might try rename the graph inputs.")
def find_bert_inputs( onnx_model: OnnxModel, input_ids_name: Optional[str] = None, segment_ids_name: Optional[str] = None, input_mask_name: Optional[str] = None, ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]: """Find graph inputs for BERT model. First, we will deduce inputs from EmbedLayerNormalization node. If not found, we will guess the meaning of graph inputs based on naming. Args: onnx_model (OnnxModel): onnx model object input_ids_name (str, optional): Name of graph input for input IDs. Defaults to None. segment_ids_name (str, optional): Name of graph input for segment IDs. Defaults to None. input_mask_name (str, optional): Name of graph input for attention mask. Defaults to None. Raises: ValueError: Graph does not have input named of input_ids_name or segment_ids_name or input_mask_name ValueError: Expected graph input number does not match with specified input_ids_name, segment_ids_name and input_mask_name Returns: Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]: input tensors of input_ids, segment_ids and input_mask """ graph_inputs = onnx_model.get_graph_inputs_excluding_initializers() if input_ids_name is not None: input_ids = onnx_model.find_graph_input(input_ids_name) if input_ids is None: raise ValueError( f"Graph does not have input named {input_ids_name}") segment_ids = None if segment_ids_name: segment_ids = onnx_model.find_graph_input(segment_ids_name) if segment_ids is None: raise ValueError( f"Graph does not have input named {segment_ids_name}") input_mask = None if input_mask_name: input_mask = onnx_model.find_graph_input(input_mask_name) if input_mask is None: raise ValueError( f"Graph does not have input named {input_mask_name}") expected_inputs = 1 + (1 if segment_ids else 0) + (1 if input_mask else 0) if len(graph_inputs) != expected_inputs: raise ValueError( f"Expect the graph to have {expected_inputs} inputs. Got {len(graph_inputs)}" ) return input_ids, segment_ids, input_mask if len(graph_inputs) != 3: raise ValueError("Expect the graph to have 3 inputs. Got {}".format( len(graph_inputs))) embed_nodes = onnx_model.get_nodes_by_op_type("EmbedLayerNormalization") if len(embed_nodes) == 1: embed_node = embed_nodes[0] input_ids = get_graph_input_from_embed_node(onnx_model, embed_node, 0) segment_ids = get_graph_input_from_embed_node(onnx_model, embed_node, 1) input_mask = get_graph_input_from_embed_node(onnx_model, embed_node, 7) if input_mask is None: for input in graph_inputs: input_name_lower = input.name.lower() if "mask" in input_name_lower: input_mask = input if input_mask is None: raise ValueError(f"Failed to find attention mask input") return input_ids, segment_ids, input_mask # Try guess the inputs based on naming. input_ids = None segment_ids = None input_mask = None for input in graph_inputs: input_name_lower = input.name.lower() if "mask" in input_name_lower: # matches input with name like "attention_mask" or "input_mask" input_mask = input elif ( "token" in input_name_lower or "segment" in input_name_lower ): # matches input with name like "segment_ids" or "token_type_ids" segment_ids = input else: input_ids = input if input_ids and segment_ids and input_mask: return input_ids, segment_ids, input_mask raise ValueError( "Fail to assign 3 inputs. You might try rename the graph inputs.")
def get_bert_inputs(onnx_file, input_ids_name=None, segment_ids_name=None, input_mask_name=None): """ Get graph inputs for bert model. First, we will deduce from EmbedLayerNormalization node. If not found, we will guess based on naming. """ model = ModelProto() with open(onnx_file, "rb") as f: model.ParseFromString(f.read()) onnx_model = OnnxModel(model) graph_inputs = onnx_model.get_graph_inputs_excluding_initializers() if input_ids_name is not None: input_ids = onnx_model.find_graph_input(input_ids_name) if input_ids is None: raise ValueError( f"Graph does not have input named {input_ids_name}") segment_ids = None if segment_ids_name: segment_ids = onnx_model.find_graph_input(segment_ids_name) if segment_ids is None: raise ValueError( f"Graph does not have input named {segment_ids_name}") input_mask = None if input_mask_name: input_mask = onnx_model.find_graph_input(input_mask_name) if input_mask is None: raise ValueError( f"Graph does not have input named {input_mask_name}") expected_inputs = 1 + (1 if segment_ids else 0) + (1 if input_mask else 0) if len(graph_inputs) != expected_inputs: raise ValueError( f"Expect the graph to have {expected_inputs} inputs. Got {len(graph_inputs)}" ) return input_ids, segment_ids, input_mask if len(graph_inputs) != 3: raise ValueError("Expect the graph to have 3 inputs. Got {}".format( len(graph_inputs))) embed_nodes = onnx_model.get_nodes_by_op_type('EmbedLayerNormalization') if len(embed_nodes) == 1: embed_node = embed_nodes[0] input_ids = get_graph_input_from_embed_node(onnx_model, embed_node, 0) segment_ids = get_graph_input_from_embed_node(onnx_model, embed_node, 1) input_mask = get_graph_input_from_embed_node(onnx_model, embed_node, 7) return input_ids, segment_ids, input_mask # Try guess the inputs based on naming. input_ids = None segment_ids = None input_mask = None for input in graph_inputs: input_name_lower = input.name.lower() if "mask" in input_name_lower: # matches input with name like "attention_mask" or "input_mask" input_mask = input elif "token" in input_name_lower or "segment" in input_name_lower: # matches input with name like "segment_ids" or "token_type_ids" segment_ids = input else: input_ids = input if input_ids and segment_ids and input_mask: return input_ids, segment_ids, input_mask raise ValueError( "Fail to assign 3 inputs. You might try rename the graph inputs.")