mirror of
https://github.com/deepinsight/insightface.git
synced 2025-12-30 08:02:27 +00:00
116 lines
3.4 KiB
Python
116 lines
3.4 KiB
Python
import sys
|
|
import os
|
|
import argparse
|
|
import onnx
|
|
import mxnet as mx
|
|
from onnx import helper
|
|
from onnx import TensorProto
|
|
from onnx import numpy_helper
|
|
|
|
print('mxnet version:', mx.__version__)
|
|
print('onnx version:', onnx.__version__)
|
|
|
|
assert mx.__version__ >= '1.8', 'mxnet version should >= 1.8'
|
|
assert onnx.__version__ >= '1.2.1', 'onnx version should >= 1.2.1'
|
|
|
|
import numpy as np
|
|
from mxnet.contrib import onnx as onnx_mxnet
|
|
|
|
def create_map(graph_member_list):
|
|
member_map={}
|
|
for n in graph_member_list:
|
|
member_map[n.name]=n
|
|
return member_map
|
|
|
|
|
|
parser = argparse.ArgumentParser(description='convert arcface models to onnx')
|
|
# general
|
|
parser.add_argument('params', default='./r100a/model-0000.params', help='mxnet params to load.')
|
|
parser.add_argument('output', default='./r100a.onnx', help='path to write onnx model.')
|
|
parser.add_argument('--eps', default=1.0e-8, type=float, help='eps for weights.')
|
|
parser.add_argument('--input-shape', default='3,112,112', help='input shape.')
|
|
args = parser.parse_args()
|
|
input_shape = (1,) + tuple( [int(x) for x in args.input_shape.split(',')] )
|
|
|
|
params_file = args.params
|
|
pos = params_file.rfind('-')
|
|
prefix = params_file[:pos]
|
|
epoch = int(params_file[pos+1:pos+5])
|
|
sym_file = prefix + "-symbol.json"
|
|
assert os.path.exists(sym_file)
|
|
assert os.path.exists(params_file)
|
|
|
|
sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
|
|
eps = args.eps
|
|
|
|
arg = {}
|
|
aux = {}
|
|
invalid = 0
|
|
ac = 0
|
|
for k in arg_params:
|
|
v = arg_params[k]
|
|
nv = v.asnumpy()
|
|
#print(k, nv.dtype)
|
|
nv = nv.astype(np.float32)
|
|
ac += nv.size
|
|
invalid += np.count_nonzero(np.abs(nv)<eps)
|
|
nv[np.abs(nv) < eps] = 0.0
|
|
arg[k] = mx.nd.array(nv, dtype='float32')
|
|
print(invalid, ac)
|
|
arg_params = arg
|
|
invalid = 0
|
|
ac = 0
|
|
for k in aux_params:
|
|
v = aux_params[k]
|
|
nv = v.asnumpy().astype(np.float32)
|
|
ac += nv.size
|
|
invalid += np.count_nonzero(np.abs(nv)<eps)
|
|
nv[np.abs(nv) < eps] = 0.0
|
|
aux[k] = mx.nd.array(nv, dtype='float32')
|
|
print(invalid, ac)
|
|
aux_params = aux
|
|
|
|
all_args = {}
|
|
all_args.update(arg_params)
|
|
all_args.update(aux_params)
|
|
converted_model_path = onnx_mxnet.export_model(sym, all_args, [input_shape], np.float32, args.output, opset_version=11)
|
|
|
|
model = onnx.load(args.output)
|
|
graph = model.graph
|
|
input_map = create_map(graph.input)
|
|
node_map = create_map(graph.node)
|
|
init_map = create_map(graph.initializer)
|
|
|
|
#fix PRelu issue
|
|
for input_name in input_map.keys():
|
|
if input_name.endswith('_gamma'):
|
|
node_name = input_name[:-6]
|
|
if not node_name in node_map:
|
|
continue
|
|
node = node_map[node_name]
|
|
if node.op_type!='PRelu':
|
|
continue
|
|
input_shape = input_map[input_name].type.tensor_type.shape.dim
|
|
input_dim_val=input_shape[0].dim_value
|
|
|
|
graph.initializer.remove(init_map[input_name])
|
|
weight_array = numpy_helper.to_array(init_map[input_name])
|
|
|
|
b=[]
|
|
for w in weight_array:
|
|
b.append(w)
|
|
new_nv = helper.make_tensor(input_name, TensorProto.FLOAT, [input_dim_val,1,1], b)
|
|
graph.initializer.extend([new_nv])
|
|
|
|
for init_name in init_map.keys():
|
|
weight_array = numpy_helper.to_array(init_map[init_name])
|
|
assert weight_array.dtype==np.float32
|
|
if init_name in input_map:
|
|
graph.input.remove(input_map[init_name])
|
|
|
|
#support batch-inference
|
|
graph.input[0].type.tensor_type.shape.dim[0].dim_param = 'None'
|
|
|
|
onnx.save(model, args.output)
|
|
|