有时候我们通过pytorch导出onnx模型，需要修改一下onnx的图结构，怎么修改呢？

下面两个Python实例提供了修改思路。
Changing the graph is easier than recreating it with make_graph, just use append, remove and insert.参考https://github.com/onnx/onnx/issues/2259

onnx_model = onnx.load(onnxfile)
graph = onnx_model.graph
pads = onnx.helper.make_tensor('avg_pads', onnx.TensorProto.INT64, [8], np.zeros(8, dtype=int))
graph.initializer.append(pads)
node = graph.node[584]
new_node = onnx.helper.make_node(
    'Pad',
    name='__Pad_584_fixed',
    inputs=['675', 'avg_pads'],
    outputs=['676'],
    mode='constant'
)
graph.node.remove(node)
graph.node.insert(584, new_node)
# Fix Equals (replace with Not)
node = graph.node[322]
new_node = onnx.helper.make_node(
    'Not',
    name='__Not__Equal_322',
    inputs=['412'],
    outputs=['414'],
)
graph.node.remove(node)
graph.node.insert(322, new_node)

onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, onnxfile)

来源：https://github.com/saurabh-shandilya/onnx-utils

# ------------------------------------------------
# ONNX Model Editor and Graph Extractor
# License under The MIT License
# Written by Saurabh Shandilya
# -----------------------------------------------

import onnx
from onnx import helper, checker
from onnx import TensorProto
import re
import argparse

def createGraphMemberMap(graph_member_list):
    member_map=dict();
    for n in graph_member_list:
        member_map[n.name]=n;
    return member_map


def split_io_list(io_list,new_names_all):
    #splits input/output list to identify removed, retained and totally new nodes    
    removed_names=[]
    retained_names=[]
    for n in io_list:
        if n.name not in new_names_all:                
            removed_names.append(n.name)              
        if n.name in new_names_all:                
            retained_names.append(n.name)                      
    new_names=list(set(new_names_all)-set(retained_names)) 
    return [removed_names,retained_names,new_names]
          
def traceDependentNodes(graph,name,node_input_names,node_map, initializer_map):
    # recurisvely traces all dependent nodes for a given output nodes in a graph    
    for n in graph.node:
        for noutput in n.output:       
            if (noutput == name) and (n.name not in node_input_names):
                # give node "name" is node n's output, so add node "n" to node_input_names list 
                node_input_names.append(n.name)
                if n.name in node_map.keys():
                    for ninput in node_map[n.name].input:
                        # trace input node's inputs 
                        node_input_names = traceDependentNodes(graph,ninput,node_input_names,node_map, initializer_map)                                        
    # don't forget the initializers they can be terminal inputs on a path.                    
    if name in initializer_map.keys():
        node_input_names.append(name)                    
    return node_input_names     
    
def onnx_edit(input_model, output_model, new_input_node_names, input_shape_map, new_output_node_names, output_shape_map, verify):
    """ edits and modifies an onnx model to extract a subgraph based on input/output node names and shapes.
    Arguments: 
        input_model: path of input onnx model
        output_model: path of output onnx model    
        new_input_node_names: list of input node names including list of original input nodes if they are to be retained.
            If the list is empty original input nodes are assumed. 
        input_shape_map: dictionary/map of input node names to corresponding shapes. Shapes are needed for model checker to pass.
        new_output_node_names: list of output node names, including list of original output nodes if they are to be retained
            If the list if empty original output nodes are assumed.
        output_shape_map: dictionary/map of output node names to corresponding shape. Shapes are needed for model checker to pass.
        verify: set to true if input and output models need to be verified.
    """
    # LOAD MODEL AND PREP MAPS
    model = onnx.load(input_model)
    graph = model.graph
    if(verify):
        print("input model Errors: ", onnx.checker.check_model(model))
    
    node_map = createGraphMemberMap(graph.node)
    input_map = createGraphMemberMap(graph.input)
    output_map = createGraphMemberMap(graph.output)
    initializer_map = createGraphMemberMap(graph.initializer)
       
    if not new_input_node_names:
        new_input_node_names = list(input_map)
    if not new_output_node_names:
        new_output_node_names = list(output_map)
       
    # MODIFY INPUTS
    # Break the graph based on the new input node names
    [removed_names,retained_names,new_names]=split_io_list(graph.input,new_input_node_names)
    for name in removed_names:
        if name in input_map.keys():
            graph.input.remove(input_map[name])                              
    for name in new_names:
        # If a new input name corresponds to an existing node, it implies that original node in the graph needs to be replaced with an input node
        # Exactly here the graph is broken
        if name in node_map.keys():
            graph.node.remove(node_map[name])
        if(name in input_shape_map.keys()):
            new_nv = helper.make_tensor_value_info(name, TensorProto.FLOAT, input_shape_map[name])
        else:
            new_nv = helper.make_tensor_value_info(name, TensorProto.FLOAT, None)    
        graph.input.extend([new_nv])
    node_map = createGraphMemberMap(graph.node)
    input_map = createGraphMemberMap(graph.input)    

    # MODIFY OUTPUTS
    # Break the graph based on the new output node names   
    [removed_names,retained_names,new_names]=split_io_list(graph.output,new_output_node_names)
    for name in removed_names:
        if name in output_map.keys():
            graph.output.remove(output_map[name])                              
    for name in new_names:
        if(name in output_shape_map.keys()):
            new_nv = helper.make_tensor_value_info(name, TensorProto.FLOAT, output_shape_map[name])
        else:
            new_nv = helper.make_tensor_value_info(name, TensorProto.FLOAT, None)
        graph.output.extend([new_nv])
    output_map = createGraphMemberMap(graph.output)      

    # CLEANUP NODES
    # Trace all dependent nodes for the current set of output nodes defined & prepare a list of invalid nodes
    valid_node_names=[]
    for new_output_node_name in new_output_node_names:
        valid_node_names=traceDependentNodes(graph,new_output_node_name,valid_node_names,node_map, initializer_map)
        valid_node_names=list(set(valid_node_names))
    invalid_node_names = list( (set(node_map.keys()) | set(initializer_map.keys())) - set(valid_node_names))
    # Remove all the invalid nodes from the graph               
    for name in invalid_node_names:
        if name in node_map.keys():
            graph.node.remove(node_map[name])        
        if name in initializer_map.keys():
            graph.initializer.remove(initializer_map[name])
        if name in input_map.keys():
            graph.input.remove(input_map[name])    

    # SAVE MODEL
    if(verify):    
        print("output model Errors: ", onnx.checker.check_model(model))
    onnx.save(model, output_model)

def parse_nodename_and_shape(name):
    # parses node names and shapes from input argument string
    inputs = []
    shapes = {}
    # input takes in most cases the format name:0, where 0 is the output number, and shapes
    # are appended to the same e.g. name:0[1,28,28,3]
    name_pattern = r"(?:([\w\d/\-\._:]+)(\[[\-\d,]+\])?),?"
    
    splits = re.split(name_pattern, name)
    for i in range(1, len(splits), 3):        
        inputs.append(splits[i])
        if splits[i + 1] is not None:
            shapes[splits[i]] = [int(n) for n in splits[i + 1][1:-1].split(",")]
    if not shapes:
        shapes = None
    return inputs, shapes    
    
    
    
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("input", help="input onnx model")
    parser.add_argument("output", help="output onnx model")
    parser.add_argument("--inputs", help="comma separated model input names appended with shapes, e.g. --inputs <nodename>[1,2,3],<nodename1>[1,2,3] ")
    parser.add_argument("--outputs", help="comma separated model output names appended with shapes, e.g. --outputs <nodename>[1,2,3],<nodename1>[1,2,3] ")    
    parser.add_argument('--skipverify', dest='skipverify', action='store_true',
                    help='skip verification of model. Useful if shapes are not known')
    args = parser.parse_args()
        
    if args.inputs:
        new_input_node_names, input_shape_map = parse_nodename_and_shape(args.inputs)
        #print(new_input_node_names)
        #print(input_shape_map)
    else: 
        new_input_node_names = []
        input_shape_map = {}
        
    if args.outputs:
        new_output_node_names, output_shape_map = parse_nodename_and_shape(args.outputs)
        #print(new_output_node_names)
        #print(output_shape_map)
    else:
        new_output_node_names = []
        output_shape_map = {}
        
    onnx_edit(args.input,args.output,new_input_node_names, input_shape_map, new_output_node_names, output_shape_map, not args.skipverify)