Onnx crop intermediate node

I have recently encountered the need to onnx The middle nodes of the model are cut out ,mxnet Converted caff Model , Preprocessing （ Subtract the mean divided by the variance ） It is written in the model structure , The pretreatment layer shall be taken out separately during quantification , There are two ways to do this , One that can be loaded mxnet Trained models , Remove the pretreatment before saving ; In the other case, it will directly onnx The corresponding two preprocessing nodes in are cut off , Here is the specific implementation of the second method

import onnx_graphsurgeon as gs
import numpy as np
import onnx

# remove preprocessing node
# (x - mean)/std
def remove_preprocessing_node(onnx_path):
    onnx_model = onnx.load(onnx_path) # load onnx Model 
    graph = onnx_model.graph
    old_nodes = graph.node
    new_nodes = old_nodes[2:] # Get rid of data,sub,mul The first three nodes 
    del onnx_model.graph.node[:] #  Delete the current onnx All of the models node
    onnx_model.graph.node.extend(new_nodes) # extend New nodes 
    conv0_node = onnx_model.graph.node[0]
    conv0_node.input[0] = 'data' # Set the input to the convolution node of the first layer data node 
    # graph = onnx_model.graph
    # print(graph.node)
    onnx.checker.check_model(onnx_model)
    onnx.save(onnx_model, "./janus_mask_occ_nopre.onnx")

def crop_middle_nodes(onnx_path):


if __name__ == "__main__":
    onnx_path = "J:/code/mxnet_mask/janus_mask_occ.onnx"

    remove_preprocessing_node(onnx_path)

Realization effect ：