I have managed to get a model converted using the conversion script that I modified:
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import os
import argparse
import sys
import time
import struct
import json
import onnx
from onnx import helper, shape_inference
from onnx import AttributeProto, TensorProto, GraphProto
from onnx import numpy_helper
import numpy as np
def to_float16(bin: bytes):
return np.frombuffer(bin, dtype=np.float16)
def to_float(bin: bytes):
return np.frombuffer(bin, dtype=np.float32)
def main():
parser = argparse.ArgumentParser()
parser.add_argument('coreml_folder', help='Input coreml model folder')
args = parser.parse_args()
coreml_folder = args.coreml_folder
layer_bytes = []
net_layer_data = {}
coreml_net = coreml_folder + '/model.espresso.net'
coreml_shape = coreml_folder + '/model.espresso.shape'
coreml_weights = coreml_folder + '/model.espresso.weights'
with open(coreml_net, encoding='utf-8') as f:
net_dict = json.load(f)
net_layers = net_dict['layers']
# print(net_layers[1])
with open(coreml_shape, encoding='utf-8') as f:
net_dict = json.load(f)
net_layer_shapes = net_dict['layer_shapes']
# print(net_layer_shapes[net_layers[1]['bottom']])
# coreml_weights
with open(coreml_weights, 'rb') as f:
# First byte of the file is an integer with how many
# sections there are. This lets us iterate through each section
# and get the map for how to read the rest of the file.
num_layers = struct.unpack('<i', f.read(4))[0]
# print("num_layers: " + str(num_layers))
f.read(4) # padding bytes
# The next section defines the number of bytes each layer contains.
# It has a format of
# | Layer Number | <padding> | Bytes in layer | <padding> |
while len(layer_bytes) < num_layers:
layer_num, _, num_bytes, _ = struct.unpack('<iiii', f.read(16))
layer_bytes.append((layer_num, num_bytes))
# print("layer_num: " + str(layer_num))
# print("num_bytes: " + str(num_bytes))
# print("each layer:\n")
# Read actual layer weights. Weights are floats as far as I can tell.
for layer_num, num_bytes in layer_bytes:
# print("layer_num: " + str(layer_num))
# print("count: " + str(num_bytes / 4))
# data = struct.unpack("f" * int((num_bytes / 4)), f.read(num_bytes))
# more recent versions of neuralhash switched to float16...
raw_data = f.read(num_bytes)
net_layer_data[layer_num] = raw_data
# print(net_layer_data[1])
# tensor_shape = [16]
# tensor = helper.make_tensor('1', TensorProto.FLOAT, tensor_shape, net_layer_data[1])
# print(tensor)
# 构建onnx
# 创建输入 (ValueInfoProto)
net_inputes = []
net_input_names = net_layers[0]['bottom'].split(',')
for net_input_name in net_input_names:
net_input_shape_dict = net_layer_shapes[net_input_name]
net_input = helper.make_tensor_value_info(net_input_name, TensorProto.FLOAT,
[net_input_shape_dict['n'], net_input_shape_dict['k'],
net_input_shape_dict['h'], net_input_shape_dict['w']])
net_inputes.append(net_input)
# 创建输出 (ValueInfoProto)
net_output_shape_dict = net_layer_shapes[net_layers[-1]['top']]
net_output = helper.make_tensor_value_info(net_layers[-1]['top'], TensorProto.FLOAT,
[net_output_shape_dict['n'], net_output_shape_dict['k'],
net_output_shape_dict['h'], net_output_shape_dict['w']])
net_outputes = [net_output]
# check the model.espresso.net file to adjust the index #
print('check the model.espresso.net file to adjust the index')
# net_output_name = net_layers[-2]['top']
net_output_name = net_layers[-1]['top']
if net_output_name.isdigit() != True:
net_output2_shape_dict = net_layer_shapes[net_output_name]
net_output2 = helper.make_tensor_value_info(net_output_name, TensorProto.FLOAT,
[net_output2_shape_dict['n'], net_output2_shape_dict['k'],
net_output2_shape_dict['h'], net_output2_shape_dict['w']])
net_outputes.append(net_output2)
onnx_blob_shapes = []
for blob_name, blob_shape_dict in net_layer_shapes.items():
onnx_blob_shapes.append(helper.make_tensor_value_info(blob_name, TensorProto.FLOAT,
[blob_shape_dict['n'],
blob_shape_dict['k'],
blob_shape_dict['h'],
blob_shape_dict['w']]))
# 创建nodes NodeProto
onnx_net_nodes = []
onnx_net_weights = []
# check the model.espresso.net file to adjust the index #
print('check the model.espresso.net file to adjust the index')
# layer_info = net_layers[1]
layer_info = net_layers[0]
for layer_info in net_layers:
print(layer_info['type'])
if layer_info['type'] == 'convolution':
stride_x = 1
if ('stride_x' in layer_info):
stride_x = layer_info['stride_x']
stride_y = 1
if ('stride_y' in layer_info):
stride_y = layer_info['stride_y']
auto_pad = None
if layer_info['pad_mode'] == 1:
auto_pad = 'SAME_UPPER'
node_inputs = layer_info['bottom'].split(',')
if ('blob_weights' in layer_info):
node_inputs.append(str(layer_info['blob_weights']))
elif ('blob_weights_f16' in layer_info):
node_inputs.append(str(layer_info['blob_weights_f16']))
if ('blob_biases' in layer_info):
node_inputs.append(str(layer_info['blob_biases']))
node_conv_outputs = layer_info['top'].split(',')
node_relu_outputs = []
if layer_info['fused_relu'] == 1:
node_relu_outputs = node_conv_outputs
node_conv_outputs = []
for temp_output in node_relu_outputs:
conv_output_blob_name = 'conv_' + temp_output
node_conv_outputs.append(conv_output_blob_name)
blob_shape_dict = net_layer_shapes[temp_output]
onnx_blob_shapes.append(helper.make_tensor_value_info(conv_output_blob_name, TensorProto.FLOAT,
[blob_shape_dict['n'],
blob_shape_dict['k'],
blob_shape_dict['h'],
blob_shape_dict['w']]))
conv_group_num = layer_info['n_groups']
if auto_pad:
layer_node = helper.make_node('Conv', # node type
node_inputs, # inputs
node_conv_outputs, # outputs
kernel_shape=[layer_info['Nx'], layer_info['Ny']],
strides=[stride_x, stride_y],
auto_pad=auto_pad,
group=conv_group_num,
dilations=[1, 1])
else:
layer_node = helper.make_node('Conv', # node type
node_inputs, # inputs
node_conv_outputs, # outputs
kernel_shape=[layer_info['Nx'], layer_info['Ny']],
strides=[stride_x, stride_y],
pads=[layer_info['pad_l'], layer_info['pad_t'], layer_info['pad_r'],
layer_info['pad_b']],
group=conv_group_num,
dilations=[1, 1])
onnx_net_nodes.append(layer_node)
# weights
weights_shape = [layer_info['C'], int(layer_info['K'] / conv_group_num), layer_info['Nx'], layer_info['Ny']]
onnx_weights_tensor = helper.make_tensor(str(layer_info['blob_weights_f16']), TensorProto.FLOAT, weights_shape,
tuple(to_float16(net_layer_data[layer_info['blob_weights_f16']]).astype(np.float32)))
onnx_net_weights.append(onnx_weights_tensor)
# bias
if ('blob_biases' in layer_info):
bias_shape = [layer_info['C']]
onnx_bias_tensor = helper.make_tensor(str(layer_info['blob_biases']), TensorProto.FLOAT, bias_shape,
tuple(to_float(net_layer_data[layer_info['blob_biases']])))
onnx_net_weights.append(onnx_bias_tensor)
if layer_info['fused_relu'] == 1:
layer_node = helper.make_node('Relu', # node type
node_conv_outputs, # inputs
node_relu_outputs, # outputs
)
onnx_net_nodes.append(layer_node)
elif layer_info['type'] == 'pool':
stride_x = 1
if ('stride_x' in layer_info):
stride_x = layer_info['stride_x']
stride_y = 1
if ('stride_y' in layer_info):
stride_y = layer_info['stride_y']
node_type = 'MaxPool'
if layer_info['avg_or_max'] == 0:
node_type = 'AveragePool'
node_inputs = layer_info['bottom'].split(',')
node_outputs = layer_info['top'].split(',')
layer_node = helper.make_node(node_type, # node type
node_inputs, # inputs
node_outputs, # outputs
kernel_shape=[layer_info['size_x'], layer_info['size_y']],
strides=[stride_x, stride_y],
pads=[layer_info['pad_l'], layer_info['pad_t'], layer_info['pad_r'],
layer_info['pad_b']])
onnx_net_nodes.append(layer_node)
elif layer_info['type'] == 'elementwise':
node_inputs = layer_info['bottom'].split(',')
node_type = ''
node_inputs_extra = []
if layer_info['operation'] == 2 or layer_info['operation'] == 0:
# check
node_type = 'Add'
if len(node_inputs) == 1:
# scales
scales_tensor_name = 'elementwise_' + layer_info['top']
node_inputs_extra.append(scales_tensor_name)
scales = [layer_info['alpha']]
onnx_scales_tensor = helper.make_tensor(scales_tensor_name, TensorProto.FLOAT, [1], scales)
onnx_net_weights.append(onnx_scales_tensor)
elif layer_info['operation'] == 1:
# check 注意如果输如只有1个，需要像取[layer_info['alpha']]值
node_type = 'Mul'
if len(node_inputs) == 1:
# scales
scales_tensor_name = 'elementwise_' + layer_info['top']
node_inputs_extra.append(scales_tensor_name)
scales = [layer_info['alpha']]
onnx_scales_tensor = helper.make_tensor(scales_tensor_name, TensorProto.FLOAT, [1], scales)
onnx_net_weights.append(onnx_scales_tensor)
elif layer_info['operation'] == -999:
node_type = 'Sub'
if len(node_inputs) == 1:
# scales
scales_tensor_name = 'elementwise_' + layer_info['top']
node_inputs_extra.append(scales_tensor_name)
scales = [layer_info['alpha']]
onnx_scales_tensor = helper.make_tensor(scales_tensor_name, TensorProto.FLOAT, [1], scales)
onnx_net_weights.append(onnx_scales_tensor)
elif layer_info['operation'] == 3:
# check
node_type = 'Mul'
if len(node_inputs) == 1:
# scales
scales_tensor_name = 'elementwise_' + layer_info['top']
node_inputs_extra.append(scales_tensor_name)
scales = [layer_info['alpha']]
onnx_scales_tensor = helper.make_tensor(scales_tensor_name, TensorProto.FLOAT, [1], scales)
onnx_net_weights.append(onnx_scales_tensor)
elif layer_info['operation'] == 10:
# check，求倒数，y=1/x
node_type = 'Div'
# scales
scales_tensor_name = 'elementwise_' + layer_info['top']
node_inputs_extra.append(scales_tensor_name)
scales = [layer_info['alpha']]
onnx_scales_tensor = helper.make_tensor(scales_tensor_name, TensorProto.FLOAT, [1], scales)
onnx_net_weights.append(onnx_scales_tensor)
elif layer_info['operation'] == 24:
# check
node_type = 'Abs'
elif layer_info['operation'] == 119:
node_type = 'Clip'
# check
alpha_tensor_name = 'elementwise_' + layer_info['top'] + 'alpha'
beta_tensor_name = 'elementwise_' + layer_info['top'] + 'beta'
node_inputs_extra.append(alpha_tensor_name)
node_inputs_extra.append(beta_tensor_name)
alpha = 0.0
if 'alpha' in layer_info:
alpha = [layer_info['alpha']]
beta = 1.0
if 'beta' in layer_info:
beta = [layer_info['beta']]
onnx_alpha_tensor = helper.make_tensor(alpha_tensor_name, TensorProto.FLOAT, [1], alpha)
onnx_beta_tensor = helper.make_tensor(beta_tensor_name, TensorProto.FLOAT, [1], beta)
onnx_net_weights.append(onnx_alpha_tensor)
onnx_net_weights.append(onnx_beta_tensor)
else:
print('Error: unsupported elementwise operation: ' + str(layer_info['operation']))
assert (0)
node_inputs = layer_info['bottom'].split(',')
node_inputs.extend(node_inputs_extra)
node_outputs = layer_info['top'].split(',')
layer_node = helper.make_node(node_type, # node type
node_inputs, # inputs
node_outputs, # outputs
)
onnx_net_nodes.append(layer_node)
elif layer_info['type'] == 'upsample':
node_type = 'Upsample'
mode = 0
if layer_info['mode'] != 0:
print('Error: unsupported upsample mode: ' + str(layer_info['mode']))
assert (0)
scales_tensor_name = 'upsample_' + layer_info['top']
node_inputs = layer_info['bottom'].split(',')
if node_inputs[0].isdigit() != True:
node_input_shape_dict = net_layer_shapes[node_inputs[0]]
node_input_tensor = helper.make_tensor_value_info(node_inputs[0], TensorProto.FLOAT,
[node_input_shape_dict['n'],
node_input_shape_dict['k'],
node_input_shape_dict['h'],
node_input_shape_dict['w']])
net_inputes.append(node_input_tensor)
node_inputs.append(scales_tensor_name)
node_outputs = layer_info['top'].split(',')
layer_node = helper.make_node(node_type, # node type
node_inputs, # inputs
node_outputs, # outputs
mode='nearest',
)
onnx_net_nodes.append(layer_node)
# scales
scales = [1.0, 1.0, layer_info['scaling_factor_x'], layer_info['scaling_factor_y']]
onnx_scales_tensor = helper.make_tensor(scales_tensor_name, TensorProto.FLOAT, [4], scales)
onnx_net_weights.append(onnx_scales_tensor)
elif layer_info['type'] == 'concat':
node_type = 'Concat'
node_inputs = layer_info['bottom'].split(',')
node_outputs = layer_info['top'].split(',')
layer_node = helper.make_node(node_type, # node type
node_inputs, # inputs
node_outputs, # outputs
axis=1,
)
onnx_net_nodes.append(layer_node)
elif layer_info['type'] == 'activation':
node_inputs = layer_info['bottom'].split(',')
node_outputs = layer_info['top'].split(',')
activation_mode = layer_info['mode']
if activation_mode == 0:
layer_node = helper.make_node('Relu', # node type
node_inputs, # inputs
node_outputs, # outputs
)
elif activation_mode == 1:
layer_node = helper.make_node('Tanh', # node type
node_inputs, # inputs
node_outputs, # outputs
)
elif activation_mode == 2:
layer_node = helper.make_node('LeakyRelu', # node type
node_inputs, # inputs
node_outputs, # outputs
)
elif activation_mode == 3:
layer_node = helper.make_node('Sigmoid', # node type
node_inputs, # inputs
node_outputs, # outputs
)
elif activation_mode == 4:
layer_node = helper.make_node('PRelu', # node type
node_inputs, # inputs
node_outputs, # outputs
)
elif activation_mode == 8:
layer_node = helper.make_node('Elu', # node type
node_inputs, # inputs
node_outputs, # outputs
)
elif activation_mode == 9:
layer_node = helper.make_node('ThresholdedRelu', # node type
node_inputs, # inputs
node_outputs, # outputs
alpha=layer_info['alpha']
)
elif activation_mode == 10:
layer_node = helper.make_node('Softplus', # node type
node_inputs, # inputs
node_outputs, # outputs
)
elif activation_mode == 12:
layer_node = helper.make_node('Softsign', # node type
node_inputs, # inputs
node_outputs, # outputs
)
else:
print('Error: unsupported activation mode: ' + str(activation_mode))
assert (0)
onnx_net_nodes.append(layer_node)
elif layer_info['type'] == 'load_constant':
# constant_blob
print('constant_blob: ' + str(layer_info['constant_blob']))
elif layer_info['type'] == 'batchnorm':
node_inputs = layer_info['bottom'].split(',')
weights_prefix = str(layer_info['blob_batchnorm_params'])
if ('blob_batchnorm_params' in layer_info):
node_inputs.append(weights_prefix + 's')
node_inputs.append(weights_prefix + 'bias')
# node_inputs.append(weights_prefix+'mean')
# node_inputs.append(weights_prefix+'var')
channels = layer_info['C']
node_bn_outputs = layer_info['top'].split(',')
data = to_float(net_layer_data[layer_info['blob_batchnorm_params']])
data = data.reshape([channels, 4])
s = data[:, 0]
bias = data[:, 1]
# mean = data[:,1]
# var = data[:,0]
training_mode = 0
if layer_info['training'] == 1:
# node_bn_outputs.append(layer_info['name']+'mean')
# node_bn_outputs.append(layer_info['name']+'var')
training_mode = 1
layer_node = helper.make_node('InstanceNormalization', # node type
node_inputs, # inputs
node_bn_outputs, # outputs
epsilon=layer_info['training_eps'])
# momentum=layer_info['training_momentum'],
# training_mode=training_mode)
onnx_net_nodes.append(layer_node)
# weights
weights_shape = [layer_info['C']]
onnx_s_tensor = helper.make_tensor(weights_prefix + 's', TensorProto.FLOAT, weights_shape, s)
onnx_net_weights.append(onnx_s_tensor)
onnx_bias_tensor = helper.make_tensor(weights_prefix + 'bias', TensorProto.FLOAT, weights_shape, bias)
onnx_net_weights.append(onnx_bias_tensor)
# onnx_mean_tensor = helper.make_tensor(weights_prefix+'mean', TensorProto.FLOAT, weights_shape, mean)
# onnx_net_weights.append(onnx_mean_tensor)
# onnx_var_tensor = helper.make_tensor(weights_prefix+'var', TensorProto.FLOAT, weights_shape, var)
# onnx_net_weights.append(onnx_var_tensor)
elif layer_info['type'] == 'inner_product':
stride_x = 1
if ('stride_x' in layer_info):
stride_x = layer_info['stride_x']
stride_y = 1
if ('stride_y' in layer_info):
stride_y = layer_info['stride_y']
node_inputs = layer_info['bottom'].split(',')
if ('blob_weights' in layer_info):
node_inputs.append(str(layer_info['blob_weights']))
if ('blob_biases' in layer_info):
node_inputs.append(str(layer_info['blob_biases']))
layer_info['Nx'] = 1
layer_info['Ny'] = 1
node_conv_outputs = layer_info['top'].split(',')
node_relu_outputs = []
if layer_info['has_relu'] == 1:
node_relu_outputs = node_conv_outputs
node_conv_outputs = []
for temp_output in node_relu_outputs:
conv_output_blob_name = 'conv_' + temp_output
node_conv_outputs.append(conv_output_blob_name)
blob_shape_dict = net_layer_shapes[temp_output]
onnx_blob_shapes.append(helper.make_tensor_value_info(conv_output_blob_name, TensorProto.FLOAT,
[blob_shape_dict['n'],
blob_shape_dict['k'],
blob_shape_dict['h'],
blob_shape_dict['w']]))
layer_node = helper.make_node('Conv', # node type
node_inputs, # inputs
node_conv_outputs, # outputs
kernel_shape=[layer_info['Nx'], layer_info['Ny']],
strides=[stride_x, stride_y],
pads=[0, 0, 0, 0],
group=1,
dilations=[1, 1])
onnx_net_nodes.append(layer_node)
# weights
weights_shape = [layer_info['nC'], int(layer_info['nB']), layer_info['Nx'], layer_info['Ny']]
onnx_weights_tensor = helper.make_tensor(str(layer_info['blob_weights_f16']), TensorProto.FLOAT, weights_shape,
tuple(to_float16(net_layer_data[layer_info['blob_weights_f16']]).astype(np.float32)))
onnx_net_weights.append(onnx_weights_tensor)
# bias
if ('blob_biases' in layer_info):
bias_shape = [layer_info['nC']]
onnx_bias_tensor = helper.make_tensor(str(layer_info['blob_biases']), TensorProto.FLOAT, bias_shape,
tuple(to_float(net_layer_data[layer_info['blob_biases']])))
onnx_net_weights.append(onnx_bias_tensor)
if layer_info['has_relu'] == 1:
layer_node = helper.make_node('Relu', # node type
node_conv_outputs, # inputs
node_relu_outputs, # outputs
)
onnx_net_nodes.append(layer_node)
else:
print('Error: unsupported layer type: ' + layer_info['type'])
assert (0)
# 创建graph GraphProto
graph_def = helper.make_graph(
onnx_net_nodes,
'onnx-model',
net_inputes,
net_outputes,
initializer=onnx_net_weights,
value_info=onnx_blob_shapes,
)
# 创建model (ModelProto)
# onnx_model = helper.make_model(graph_def, producer_name='YouTu Tencent')
onnx_model = helper.make_model(graph_def, producer_name='YouTu Tencent',
opset_imports=[helper.make_operatorsetid("", 12)])
# Change model version to support onnxruntime
onnx_model.ir_version = 7
# print('The model is:\n{}'.format(onnx_model))
onnx.checker.check_model(onnx_model)
print('Before shape inference, the shape info of Y is:\n{}'.format(onnx_model.graph.value_info))
# Apply shape inference on the model
inferred_model = shape_inference.infer_shapes(onnx_model)
onnx.checker.check_model(inferred_model)
print('After shape inference, the shape info of Y is:\n{}'.format(inferred_model.graph.value_info))
onnx.save(inferred_model, coreml_folder + '/model.onnx')
if __name__ == '__main__':
main()
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent
