mmdeploy/apis/extract_model.py
def extract_model(model: Union[str, onnx.ModelProto],
                  start: Union[str, Iterable[str]],
                  end: Union[str, Iterable[str]],
                  start_name_map: Optional[Dict[str, str]] = None,
                  end_name_map: Optional[Dict[str, str]] = None,
                  dynamic_axes: Optional[Dict[str, Dict[int, str]]] = None,
                  save_file: Optional[str] = None) -> onnx.ModelProto:
    """Extract sub-model from an ONNX model.

    The sub-model is defined by the names of the input and output tensors
    exactly.

    Args:
        model (str | onnx.ModelProto): Input ONNX model to be extracted.
        start (str | Sequence[str]): Start marker(s) to extract.
        end (str | Sequence[str]): End marker(s) to extract.
        start_name_map (Dict[str, str]): A mapping of start names, defaults to
            `None`.
        end_name_map (Dict[str, str]): A mapping of end names, defaults to
            `None`.
        dynamic_axes (Dict[str, Dict[int, str]]): A dictionary to specify
            dynamic axes of input/output, defaults to `None`.
        save_file (str): A file to save the extracted model, defaults to
            `None`.

    Returns:
        onnx.ModelProto: The extracted sub-model.
    """
    if isinstance(model, str):
        model = onnx.load(model)

    num_value_info = len(model.graph.value_info)
    inputs = []
    outputs = []
    if not isinstance(start, (list, tuple)):
        start = [start]
    for s in start:
        start_name, func_id, start_type = parse_extractor_io_string(s)
        for node in model.graph.node:
            if node.op_type == 'Mark':
                attr = attribute_to_dict(node.attribute)
                if attr['func'] == start_name and attr[
                        'type'] == start_type and attr['func_id'] == func_id:
                    name = node.input[0]
                    if name not in inputs:
                        new_name = get_new_name(
                            attr, mark_name=s, name_map=start_name_map)
                        rename_value(model, name, new_name)
                        if not any([
                                v_info.name == new_name
                                for v_info in model.graph.value_info
                        ]):
                            new_val_info = onnx.helper.make_tensor_value_info(
                                new_name, attr['dtype'], attr['shape'])
                            model.graph.value_info.append(new_val_info)
                        inputs.append(new_name)

    logging.info(f'inputs: {", ".join(inputs)}')

    # collect outputs
    if not isinstance(end, (list, tuple)):
        end = [end]
    for e in end:
        end_name, func_id, end_type = parse_extractor_io_string(e)
        for node in model.graph.node:
            if node.op_type == 'Mark':
                attr = attribute_to_dict(node.attribute)
                if attr['func'] == end_name and attr[
                        'type'] == end_type and attr['func_id'] == func_id:
                    name = node.output[0]
                    if name not in outputs:
                        new_name = get_new_name(
                            attr, mark_name=e, name_map=end_name_map)
                        rename_value(model, name, new_name)
                        if not any([
                                v_info.name == new_name
                                for v_info in model.graph.value_info
                        ]):
                            new_val_info = onnx.helper.make_tensor_value_info(
                                new_name, attr['dtype'], attr['shape'])
                            model.graph.value_info.append(new_val_info)
                        outputs.append(new_name)

    logging.info(f'outputs: {", ".join(outputs)}')

    # replace Mark with Identity
    for node in model.graph.node:
        if node.op_type == 'Mark':
            del node.attribute[:]
            node.domain = ''
            node.op_type = 'Identity'

    extractor = create_extractor(model)
    extracted_model = extractor.extract_model(inputs, outputs)

    # remove all Identity, this may be done by onnx simplifier
    remove_identity(extracted_model)

    # collect all used inputs
    used = set()
    for node in extracted_model.graph.node:
        for input in node.input:
            used.add(input)

    for output in extracted_model.graph.output:
        used.add(output.name)

    # delete unused inputs
    success = True
    while success:
        success = False
        for i, input in enumerate(extracted_model.graph.input):
            if input.name not in used:
                del extracted_model.graph.input[i]
                success = True
                break

    # eliminate output without shape
    for xs in [extracted_model.graph.output]:
        for x in xs:
            if not x.type.tensor_type.shape.dim:
                logging.info(f'fixing output shape: {x.name}')
                x.CopyFrom(
                    onnx.helper.make_tensor_value_info(
                        x.name, x.type.tensor_type.elem_type, []))

    # eliminate 0-batch dimension, dirty workaround for two-stage detectors
    for input in extracted_model.graph.input:
        if input.name in inputs:
            if input.type.tensor_type.shape.dim[0].dim_value == 0:
                input.type.tensor_type.shape.dim[0].dim_value = 1

    # eliminate duplicated value_info for inputs
    success = True
    # num_value_info == 0 if dynamic shape
    if num_value_info == 0:
        while len(extracted_model.graph.value_info) > 0:
            extracted_model.graph.value_info.pop()
    while success:
        success = False
        for i, x in enumerate(extracted_model.graph.value_info):
            if x.name in inputs:
                del extracted_model.graph.value_info[i]
                success = True
                break

    # dynamic shape support
    if dynamic_axes is not None:
        for input_node in extracted_model.graph.input:
            if input_node.name in dynamic_axes:
                axes = dynamic_axes[input_node.name]
                for k, v in axes.items():
                    input_node.type.tensor_type.shape.dim[k].dim_value = 0
                    input_node.type.tensor_type.shape.dim[k].dim_param = v
        for output_node in extracted_model.graph.output:
            for idx, dim in enumerate(output_node.type.tensor_type.shape.dim):
                dim.dim_value = 0
                dim.dim_param = f'dim_{idx}'

    # save extract_model if save_file is given
    if save_file is not None:
        onnx.save(extracted_model, save_file)

    return extracted_model
Extract sub-model from an ONNX model.

    The sub-model is defined by the names of the input and output tensors
    exactly.

    Args:
        model (str | onnx.ModelProto): Input ONNX model to be extracted.
        start (str | Sequence[str]): Start marker(s) to extract.
        end (str | Sequence[str]): End marker(s) to extract.
        start_name_map (Dict[str, str]): A mapping of start names, defaults to
            `None`.
        end_name_map (Dict[str, str]): A mapping of end names, defaults to
            `None`.
        dynamic_axes (Dict[str, Dict[int, str]]): A dictionary to specify
            dynamic axes of input/output, defaults to `None`.
        save_file (str): A file to save the extracted model, defaults to
            `None`.

    Returns:
        onnx.ModelProto: The extracted sub-model.
    
Extract sub-model from an ONNX model.
The sub-model is defined by the names of the input and output tensors
exactly.
Extract sub - model from an ONNX model . The sub - model is defined by the names of the input and output tensors exactly .
def extract_model(model: Union[str, onnx.ModelProto],
                  start: Union[str, Iterable[str]],
                  end: Union[str, Iterable[str]],
                  start_name_map: Optional[Dict[str, str]] = None,
                  end_name_map: Optional[Dict[str, str]] = None,
                  dynamic_axes: Optional[Dict[str, Dict[int, str]]] = None,
                  save_file: Optional[str] = None) -> onnx.ModelProto:
    if isinstance(model, str):
        model = onnx.load(model)
    num_value_info = len(model.graph.value_info)
    inputs = []
    outputs = []
    if not isinstance(start, (list, tuple)):
        start = [start]
    for s in start:
        start_name, func_id, start_type = parse_extractor_io_string(s)
        for node in model.graph.node:
            if node.op_type == 'Mark':
                attr = attribute_to_dict(node.attribute)
                if attr['func'] == start_name and attr[
                        'type'] == start_type and attr['func_id'] == func_id:
                    name = node.input[0]
                    if name not in inputs:
                        new_name = get_new_name(
                            attr, mark_name=s, name_map=start_name_map)
                        rename_value(model, name, new_name)
                        if not any([
                                v_info.name == new_name
                                for v_info in model.graph.value_info
                        ]):
                            new_val_info = onnx.helper.make_tensor_value_info(
                                new_name, attr['dtype'], attr['shape'])
                            model.graph.value_info.append(new_val_info)
                        inputs.append(new_name)
    logging.info(f'inputs: {", ".join(inputs)}')
    if not isinstance(end, (list, tuple)):
        end = [end]
    for e in end:
        end_name, func_id, end_type = parse_extractor_io_string(e)
        for node in model.graph.node:
            if node.op_type == 'Mark':
                attr = attribute_to_dict(node.attribute)
                if attr['func'] == end_name and attr[
                        'type'] == end_type and attr['func_id'] == func_id:
                    name = node.output[0]
                    if name not in outputs:
                        new_name = get_new_name(
                            attr, mark_name=e, name_map=end_name_map)
                        rename_value(model, name, new_name)
                        if not any([
                                v_info.name == new_name
                                for v_info in model.graph.value_info
                        ]):
                            new_val_info = onnx.helper.make_tensor_value_info(
                                new_name, attr['dtype'], attr['shape'])
                            model.graph.value_info.append(new_val_info)
                        outputs.append(new_name)
    logging.info(f'outputs: {", ".join(outputs)}')
    for node in model.graph.node:
        if node.op_type == 'Mark':
            del node.attribute[:]
            node.domain = ''
            node.op_type = 'Identity'
    extractor = create_extractor(model)
    extracted_model = extractor.extract_model(inputs, outputs)
    remove_identity(extracted_model)
    used = set()
    for node in extracted_model.graph.node:
        for input in node.input:
            used.add(input)
    for output in extracted_model.graph.output:
        used.add(output.name)
    success = True
    while success:
        success = False
        for i, input in enumerate(extracted_model.graph.input):
            if input.name not in used:
                del extracted_model.graph.input[i]
                success = True
                break
    for xs in [extracted_model.graph.output]:
        for x in xs:
            if not x.type.tensor_type.shape.dim:
                logging.info(f'fixing output shape: {x.name}')
                x.CopyFrom(
                    onnx.helper.make_tensor_value_info(
                        x.name, x.type.tensor_type.elem_type, []))
    for input in extracted_model.graph.input:
        if input.name in inputs:
            if input.type.tensor_type.shape.dim[0].dim_value == 0:
                input.type.tensor_type.shape.dim[0].dim_value = 1
    success = True
    if num_value_info == 0:
        while len(extracted_model.graph.value_info) > 0:
            extracted_model.graph.value_info.pop()
    while success:
        success = False
        for i, x in enumerate(extracted_model.graph.value_info):
            if x.name in inputs:
                del extracted_model.graph.value_info[i]
                success = True
                break
    if dynamic_axes is not None:
        for input_node in extracted_model.graph.input:
            if input_node.name in dynamic_axes:
                axes = dynamic_axes[input_node.name]
                for k, v in axes.items():
                    input_node.type.tensor_type.shape.dim[k].dim_value = 0
                    input_node.type.tensor_type.shape.dim[k].dim_param = v
        for output_node in extracted_model.graph.output:
            for idx, dim in enumerate(output_node.type.tensor_type.shape.dim):
                dim.dim_value = 0
                dim.dim_param = f'dim_{idx}'
    if save_file is not None:
        onnx.save(extracted_model, save_file)
    return extracted_model
def extract_model ( model : Union [ str , onnx . ModelProto ] , start : Union [ str , Iterable [ str ] ] , end : Union [ str , Iterable [ str ] ] , start_name_map : Optional [ Dict [ str , str ] ] = None , end_name_map : Optional [ Dict [ str , str ] ] = None , dynamic_axes : Optional [ Dict [ str , Dict [ int , str ] ] ] = None , save_file : Optional [ str ] = None ) -> onnx . ModelProto : if isinstance ( model , str ) : model = onnx . load ( model ) num_value_info = len ( model . graph . value_info ) inputs = [ ] outputs = [ ] if not isinstance ( start , ( list , tuple ) ) : start = [ start ] for s in start : start_name , func_id , start_type = parse_extractor_io_string ( s ) for node in model . graph . node : if node . op_type == 'Mark' : attr = attribute_to_dict ( node . attribute ) if attr [ 'func' ] == start_name and attr [ 'type' ] == start_type and attr [ 'func_id' ] == func_id : name = node . input [ 0 ] if name not in inputs : new_name = get_new_name ( attr , mark_name = s , name_map = start_name_map ) rename_value ( model , name , new_name ) if not any ( [ v_info . name == new_name for v_info in model . graph . value_info ] ) : new_val_info = onnx . helper . make_tensor_value_info ( new_name , attr [ 'dtype' ] , attr [ 'shape' ] ) model . graph . value_info . append ( new_val_info ) inputs . append ( new_name ) logging . info ( f'inputs: {", ".join(inputs)}' ) if not isinstance ( end , ( list , tuple ) ) : end = [ end ] for e in end : end_name , func_id , end_type = parse_extractor_io_string ( e ) for node in model . graph . node : if node . op_type == 'Mark' : attr = attribute_to_dict ( node . attribute ) if attr [ 'func' ] == end_name and attr [ 'type' ] == end_type and attr [ 'func_id' ] == func_id : name = node . output [ 0 ] if name not in outputs : new_name = get_new_name ( attr , mark_name = e , name_map = end_name_map ) rename_value ( model , name , new_name ) if not any ( [ v_info . name == new_name for v_info in model . graph . value_info ] ) : new_val_info = onnx . helper . make_tensor_value_info ( new_name , attr [ 'dtype' ] , attr [ 'shape' ] ) model . graph . value_info . append ( new_val_info ) outputs . append ( new_name ) logging . info ( f'outputs: {", ".join(outputs)}' ) for node in model . graph . node : if node . op_type == 'Mark' : del node . attribute [ : ] node . domain = '' node . op_type = 'Identity' extractor = create_extractor ( model ) extracted_model = extractor . extract_model ( inputs , outputs ) remove_identity ( extracted_model ) used = set ( ) for node in extracted_model . graph . node : for input in node . input : used . add ( input ) for output in extracted_model . graph . output : used . add ( output . name ) success = True while success : success = False for i , input in enumerate ( extracted_model . graph . input ) : if input . name not in used : del extracted_model . graph . input [ i ] success = True break for xs in [ extracted_model . graph . output ] : for x in xs : if not x . type . tensor_type . shape . dim : logging . info ( f'fixing output shape: {x.name}' ) x . CopyFrom ( onnx . helper . make_tensor_value_info ( x . name , x . type . tensor_type . elem_type , [ ] ) ) for input in extracted_model . graph . input : if input . name in inputs : if input . type . tensor_type . shape . dim [ 0 ] . dim_value == 0 : input . type . tensor_type . shape . dim [ 0 ] . dim_value = 1 success = True if num_value_info == 0 : while len ( extracted_model . graph . value_info ) > 0 : extracted_model . graph . value_info . pop ( ) while success : success = False for i , x in enumerate ( extracted_model . graph . value_info ) : if x . name in inputs : del extracted_model . graph . value_info [ i ] success = True break if dynamic_axes is not None : for input_node in extracted_model . graph . input : if input_node . name in dynamic_axes : axes = dynamic_axes [ input_node . name ] for k , v in axes . items ( ) : input_node . type . tensor_type . shape . dim [ k ] . dim_value = 0 input_node . type . tensor_type . shape . dim [ k ] . dim_param = v for output_node in extracted_model . graph . output : for idx , dim in enumerate ( output_node . type . tensor_type . shape . dim ) : dim . dim_value = 0 dim . dim_param = f'dim_{idx}' if save_file is not None : onnx . save ( extracted_model , save_file ) return extracted_model
Extract sub-model from an ONNX model.
Extract sub - model from an ONNX model .
"""Extract sub-model from an ONNX model.

    The sub-model is defined by the names of the input and output tensors
    exactly.

    Args:
        model (str | onnx.ModelProto): Input ONNX model to be extracted.
        start (str | Sequence[str]): Start marker(s) to extract.
        end (str | Sequence[str]): End marker(s) to extract.
        start_name_map (Dict[str, str]): A mapping of start names, defaults to
            `None`.
        end_name_map (Dict[str, str]): A mapping of end names, defaults to
            `None`.
        dynamic_axes (Dict[str, Dict[int, str]]): A dictionary to specify
            dynamic axes of input/output, defaults to `None`.
        save_file (str): A file to save the extracted model, defaults to
            `None`.

    Returns:
        onnx.ModelProto: The extracted sub-model.
    """ # collect outputs # replace Mark with Identity # remove all Identity, this may be done by onnx simplifier # collect all used inputs # delete unused inputs # eliminate output without shape # eliminate 0-batch dimension, dirty workaround for two-stage detectors # eliminate duplicated value_info for inputs # num_value_info == 0 if dynamic shape # dynamic shape support # save extract_model if save_file is given
The extracted sub-model.
The extracted sub - model .
onnx.ModelProto
model
Union[str, onnx.ModelProto]
Input ONNX model to be extracted.
Input ONNX model to be extracted .
start
Union[str, Iterable[str]]
Start marker(s) to extract.
Start marker ( s ) to extract .
end
Union[str, Iterable[str]]
End marker(s) to extract.
End marker ( s ) to extract .
start_name_map
Optional[Dict[str, str]]
A mapping of start names, defaults to
`None`.
A mapping of start names defaults to ` None ` .
end_name_map
Optional[Dict[str, str]]
A mapping of end names, defaults to
`None`.
A mapping of end names defaults to ` None ` .
dynamic_axes
Optional[Dict[str, Dict[int, str]]]
A dictionary to specify
dynamic axes of input/output, defaults to `None`.
A dictionary to specify dynamic axes of input / output defaults to ` None ` .
save_file
Optional[str]
A file to save the extracted model, defaults to
`None`.
A file to save the extracted model defaults to ` None ` .