onnx查询算子并修改内部算子与连接

onnx查询算子并修改内部算子与连接

• 一个神奇的网站 https://fossies.org/linux/onnx/onnx/helper.py

修改onnx模型中的算子连接方式

import onnxruntime
from onnxruntime.datasets import get_example
import onnx
from onnx import helper, checker
from onnx import TensorProto
import re
import argparse

def createGraphMemberMap(graph_member_list):
    member_map=[]
    number=0
    for n in graph_member_list:
        if n.name.find("AveragePool")>-1:
            member_map.append(number)
        number +=1
    return member_map

def findResize(graph_member_list):
    member_map=[]
    number=0
    for n in graph_member_list:
        if n.name.find("Relu_2291")>-1:
            member_map.append(number)
        number +=1
    return member_map

if __name__  == "__main__":
    onnx_model = onnx.load("/root/lhw1/mmdetection/zhengjianzhao_matting/portrait-matting/cc2_sim.onnx")
    graph = onnx_model.graph
    id_list = createGraphMemberMap(graph.node)
    idx_resize = findResize(graph.node)
    for idx in id_list:
        name = "AdaptiveAveragePool_" + graph.node[idx].name.split("_")[1]
        graph.node[idx].name= name
        output = int(120/graph.node[idx].attribute[-1].ints[1])
        dicts={"ceil_model": 1, graph.node[idx].attribute[0].name: graph.node[idx].attribute[0].ints, graph.node[idx].attribute[1].name: graph.node[idx].attribute[1].ints, graph.node[idx].attribute[2].name: graph.node[idx].attribute[2].ints, "output_size": (output, output)}
        new_scale_node = onnx.helper.make_node("AdaptiveAveragePool", graph.node[idx].input, graph.node[idx].output, name, **dicts)
        graph.node.remove(graph.node[idx])
        graph.node.insert(idx, new_scale_node)

        
    for id_re in idx_resize:
        dicts_1 = {"mode": "bilinear"}
        new_scale_node1 = onnx.helper.make_node("Upsample",["2867", "2810"],graph.node[392].output, graph.node[392].name, **dicts_1)
        new_scale_node2 = onnx.helper.make_node("Upsample",["2929", "2810"],graph.node[401].output, graph.node[401].name, **dicts_1)
        new_scale_node3 = onnx.helper.make_node("Upsample",["2991", "2810"],graph.node[410].output, graph.node[410].name, **dicts_1)
        new_scale_node4 = onnx.helper.make_node("Upsample",["3053", "2810"],graph.node[419].output, graph.node[419].name, **dicts_1)
        graph.node.remove(graph.node[392])
        graph.node.insert(392, new_scale_node1)
        graph.node.remove(graph.node[401])
        graph.node.insert(401, new_scale_node2)
        graph.node.remove(graph.node[410])
        graph.node.insert(410, new_scale_node3)
        graph.node.remove(graph.node[419])
        graph.node.insert(419, new_scale_node4)
    print(graph.node)
    onnx.save(onnx_model, "./adaptive_c.onnx")
    

测试onnx的一致性

import onnxruntime
from onnxruntime.datasets import get_example

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
import struct
import torch as t
import math
import numpy as np
if __name__  == "__main__":
    torch.manual_seed(66)
    x = torch.randn(1, 11, 1200, 960)
    x1 = torch.randn(1, 3, 1200, 960)
    x2 = torch.randn(1, 3, 1200, 960)
    x3 = torch.randn(1, 2, 1200, 960)

    net = FBANet(init_state_path="/root/lhw1/mmdetection/zhengjianzhao_matting/portrait-matting/plugins/matting/model/Alpha_FBA.pth", align_corner=False)

    net.eval()
    net_out = net.forward(x,x1,x2,x3).detach().numpy()
    
    torch.onnx.export(net, (x, x1, x2, x3), "cc2.onnx", export_params=True,opset_version=11, do_constant_folding=True, input_names=["input","img_norm","img","trimap"], output_names=["output"])

    example_model = get_example("/root/lhw1/mmdetection/zhengjianzhao_matting/portrait-matting/cc2.onnx")
    sess = onnxruntime.InferenceSession(example_model)
    onnx_out = sess.run(None, {"input": to_numpy(x), "img_norm": to_numpy(x1), "img": to_numpy(x2), "trimap": to_numpy(x3)})

    print(onnx_out)
    print(net_out)

备注：https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cpu/AdaptiveAvgPoolKernel.cpp
有自适应pooling的实现方式，而且没发用avgpooling来代替，网上说有固定公式来计算出来的kernel与stride是错的