您最多选择25个主题
主题必须以中文或者字母或数字开头，可以包含连字符 (-)，并且长度不得超过35个字符
1034 行
40 KiB
1034 行
40 KiB
from __future__ import print_function
|
|
import numpy as np
|
|
import struct # convert from Python values and C structs
|
|
import tensorflow as tf
|
|
import re
|
|
#import barracuda
|
|
#from barracuda import Struct
|
|
from mlagents.trainers import barracuda
|
|
from mlagents.trainers.barracuda import Struct
|
|
from google.protobuf import descriptor
|
|
from google.protobuf.json_format import MessageToJson
|
|
|
|
|
|
if __name__ == '__main__':
|
|
# Handle command line argumengts
|
|
args = barracuda.parse_args(
|
|
description = 'Convert Tensorflow model to Barracuda binary',
|
|
source_extension = '.pb',
|
|
help = 'input Tensorflow serialized .pb file')
|
|
# Te following code can be used as an example of API used from another module
|
|
# convert() is the main entry point for converter
|
|
import tensorflow_to_barracuda as tf2bc
|
|
tf2bc.convert(args.source_file, args.target_file, args.trim_unused_by_output, args)
|
|
|
|
|
|
# TODO: support more than 1 LSTM layer per model - prepend scope to names and inputs
|
|
# TODO: support different activation functions in LSTM
|
|
# TODO: strip output Identity node, instead patch upstream layer names
|
|
# TODO: use ScaleBias and Pow with alpha when input is constant Tensor
|
|
# TODO: support all data format types (curretly only NHWC)
|
|
# TODO: support all data types (currently only FLOAT, INT32, BOOL)
|
|
# TODO: implement FusedResizeAndPadConv2D
|
|
|
|
# Important ProtoBuf definitions:
|
|
# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/types.proto
|
|
# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/tensor.proto
|
|
# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/node_def.proto
|
|
#
|
|
# Node descriptions:
|
|
# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/ops/nn_ops.cc
|
|
# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/ops/math_ops.cc
|
|
# https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/ops/random_ops.cc
|
|
#
|
|
# Class doc:
|
|
# https://www.tensorflow.org/api_docs/cc/
|
|
#
|
|
known_classes = {
|
|
'Dense': Struct(
|
|
id = 1,
|
|
out_shapes = lambda shapes: [
|
|
[shapes[0][0], 1, 1, shapes[0][1]], # W
|
|
[1, 1, 1, shapes[-1][-1]] # B
|
|
],
|
|
patch_data = lambda data: [
|
|
data[0],
|
|
data[1]
|
|
]),
|
|
'MatMul': Struct(
|
|
id = 1,
|
|
out_shapes = lambda shapes: [
|
|
[shapes[0][0], 1, 1, shapes[0][1]], # W
|
|
[1, 1, 1, shapes[0][1]] # B
|
|
],
|
|
patch_data = lambda data: [
|
|
data[0],
|
|
np.zeros(np.shape(data[1]))
|
|
]),
|
|
'BiasAdd': Struct(
|
|
id = 51, # implemented as ScaleBias
|
|
out_shapes = lambda shapes: [
|
|
[1, 1, 1, shapes[0][0]], # ONE
|
|
[1, 1, 1, shapes[0][0]], # B
|
|
],
|
|
patch_data = lambda data: [
|
|
np.ones(np.shape(data[0])),
|
|
data[0]
|
|
]),
|
|
|
|
# TODO: NCHW
|
|
'Conv2D': Struct(
|
|
id = 20,
|
|
out_shapes = lambda shapes: [
|
|
shapes[0], # K
|
|
[1, 1, 1, shapes[-1][-1]] # B
|
|
],
|
|
patch_data = lambda data: [
|
|
data[0],
|
|
data[1]
|
|
]),
|
|
'DepthwiseConv2dNative': Struct( # DepthwiseConv2D
|
|
id = 21,
|
|
out_shapes = lambda s: [
|
|
[s[0][0], s[0][1], s[0][3], s[0][2]], # K TF:[H, W, in_channels, channel_multiplier] => [H, W, 1, in_channels]
|
|
[1, 1, 1, s[-1][-1]] if len(s) > 1 else
|
|
[1, 1, 1, s[0][2]] # B
|
|
],
|
|
patch_data = lambda data: [
|
|
np.transpose(data[0], (0,1,3,2)),
|
|
data[1]
|
|
]),
|
|
'Conv2DBackpropInput': Struct( # Conv2DTranspose
|
|
id = 22,
|
|
out_shapes = lambda shapes: [
|
|
shapes[0], # K
|
|
[1, 1, 1, shapes[-1][-1]] # B
|
|
],
|
|
patch_data = lambda data: [
|
|
data[0],
|
|
data[1]
|
|
]),
|
|
|
|
# TODO: 3D
|
|
|
|
'ResizeNearestNeighbor':
|
|
23, # implemented as Upsample2D
|
|
'ResizeBilinear': 23, # implemented as Upsample2D
|
|
'ResizeBicubic': 23, # implemented as Upsample2D
|
|
'MaxPool': 25,
|
|
'AvgPool': 26,
|
|
|
|
'GlobalAveragePool':28,
|
|
|
|
'Activation': 50,
|
|
|
|
'BatchNormalization': Struct(
|
|
id = 51, # after fusion implemented as ScaleBias
|
|
out_shapes = lambda shapes: [
|
|
[1, 1, 1, shapes[0][0]], # S
|
|
[1, 1, 1, shapes[0][0]], # B
|
|
],
|
|
patch_data = lambda data:
|
|
# fuse [gamma, beta, mean, var, epsilon] => [scale, bias]
|
|
# TODO: double-check if epsilon is the last data argument and not the 1st?
|
|
barracuda.fuse_batchnorm_weights(data[0], data[1], data[2], data[3], data[4]) if len(data) == 5 else
|
|
# fuse [ONE, beta, mean, var, epsilon] => [scale, bias]
|
|
# TODO: double-check if epsilon is the last data argument and not the 1st?
|
|
barracuda.fuse_batchnorm_weights(np.ones(np.shape(data[0])), data[0], data[1], data[2], data[3])
|
|
),
|
|
'FusedBatchNorm': Struct(
|
|
id = 51, # after fusion implemented as ScaleBias
|
|
out_shapes = lambda shapes: [
|
|
[1, 1, 1, shapes[0][0]], # S
|
|
[1, 1, 1, shapes[0][0]], # B
|
|
],
|
|
patch_data = lambda data, layer:
|
|
# fuse [gamma, beta, mean, var, epsilon] => [scale, bias]
|
|
barracuda.fuse_batchnorm_weights(data[0], data[1], data[2], data[3], get_epsilon(layer))
|
|
),
|
|
'LRN': 53,
|
|
|
|
'RandomStandardNormal':
|
|
64,
|
|
'RandomUniform': 65,
|
|
'Multinomial': 66,
|
|
'OneHot': 67,
|
|
|
|
# Broadcast ops
|
|
'Add': 100,
|
|
'Sub': 101,
|
|
'Mul': 102,
|
|
'RealDiv': 103,
|
|
'Pow': 104,
|
|
'Minimum': 110,
|
|
'Maximum': 111,
|
|
|
|
# Reduce ops
|
|
'Max': 124,
|
|
'Mean': 125,
|
|
'Min': 126,
|
|
'Prod': 127,
|
|
'Sum': 128,
|
|
|
|
'Flatten': 200,
|
|
'Reshape': 201,
|
|
#'Squeeze': 203,
|
|
#'Unsqueeze': 204,
|
|
'Concat': 210,
|
|
'StridedSlice': 211,
|
|
}
|
|
|
|
requires_runtime_flag = {
|
|
'Dropout' : 'DropoutRuntime',
|
|
'BatchNormalization' : 'BatchNormalizationRuntime',
|
|
}
|
|
|
|
known_activations = {
|
|
'Linear' : 0,
|
|
'Relu' : 1,
|
|
'Softmax' : 2,
|
|
'Tanh' : 3,
|
|
'Sigmoid' : 4,
|
|
'Elu' : 5,
|
|
'Relu6' : 6,
|
|
'LeakyRelu' : 7,
|
|
'Selu' : 8,
|
|
'Swish' : 9,
|
|
|
|
'LogSoftmax' : 10,
|
|
'Softplus' : 11,
|
|
'Softsign' : 12,
|
|
|
|
'Abs' : 100,
|
|
'Neg' : 101,
|
|
'Ceil' : 102,
|
|
'Floor' : 104,
|
|
|
|
'Sqrt' : 111,
|
|
'Exp' : 113,
|
|
'Log' : 114,
|
|
|
|
'Acos' : 200,
|
|
'Acosh' : 201,
|
|
'Asin' : 202,
|
|
'Asinh' : 203,
|
|
'Atan' : 204,
|
|
'Atanh' : 205,
|
|
'Cos' : 206,
|
|
'Cosh' : 207,
|
|
'Sin' : 208,
|
|
'Sinh' : 209,
|
|
'Tan' : 210
|
|
}
|
|
|
|
known_paddings = {
|
|
'VALID' : [0,0,0,0],
|
|
'SAME' : [-1] # SameUpper
|
|
}
|
|
|
|
supported_data_formats = {
|
|
'NHWC'
|
|
}
|
|
|
|
known_patterns = {
|
|
# TODO: Flatten pattern using namespace regexp
|
|
repr(['Shape', 'StridedSlice', 'Pack', 'Reshape']) : "Flatten",
|
|
repr(['Shape', 'StridedSlice', 'Prod', 'Pack', 'Reshape']) : "Flatten",
|
|
repr(['Shape', 'Slice', 'Slice', 'Prod',
|
|
'ExpandDims', 'ConcatV2', 'Reshape']) : "Flatten",
|
|
repr(['Const', 'Reshape']) : 'Reshape',
|
|
|
|
repr(['Add', 'Rsqrt', 'Mul', 'Mul', 'Sub', 'Add']) : 'BatchNormalization',
|
|
repr(['Add', 'Rsqrt', 'Mul', 'Mul', 'Mul', 'Sub', 'Add']) : 'BatchNormalization',
|
|
|
|
repr(['ConcatV2']) : 'ConcatV2',
|
|
repr(['Mean']) : 'Mean',
|
|
repr(['Multinomial']) : 'Multinomial',
|
|
repr(['OneHot']) : 'OneHot',
|
|
repr(['Square']) : 'Square',
|
|
|
|
repr(['MatMul', 'BiasAdd']) : 'Dense',
|
|
repr(['Conv2D', 'BiasAdd']) : 'Conv2D',
|
|
repr(['DepthwiseConv2dNative', 'BiasAdd']) : 'DepthwiseConv2dNative',
|
|
repr(['Conv2DBackpropInput', 'BiasAdd']) : 'Conv2DBackpropInput',
|
|
|
|
|
|
repr(['Pack', 'Reshape']) : 'Flatten$', # for now we assume that this combination is trivial Flatten
|
|
# for exmaple it is used in ML-agents LSTM nets with sequence_length==1
|
|
|
|
repr(['StridedSlice', 'Reshape',
|
|
re.compile('^lstm/'),
|
|
'Reshape', 'ConcatV2', 'Identity']) : 'BasicLSTM',
|
|
|
|
repr([re.compile('^lstm/'),
|
|
'Reshape', 'ConcatV2', 'Identity']) : 'BasicLSTM',
|
|
|
|
repr(['Sigmoid', 'Mul']) : "Swish",
|
|
|
|
# TODO: FusedResizeAndPadConv2D
|
|
}
|
|
|
|
def by_name(args, name):
|
|
for a in args:
|
|
if a.name.endswith(name):
|
|
return a
|
|
|
|
def by_op(args, op):
|
|
for a in args:
|
|
if a.op == op:
|
|
return a
|
|
|
|
def order_by(args, names):
|
|
ordered = []
|
|
arg_count = len(args)
|
|
for name in names:
|
|
ordered += [a for a in args if a.endswith(name)]
|
|
args = [a for a in args if not a.endswith(name)]
|
|
ordered += args # append what is left
|
|
assert(len(ordered) == arg_count)
|
|
return ordered
|
|
|
|
transform_patterns = {
|
|
'Flatten' : lambda nodes, inputs, tensors, _:
|
|
Struct(
|
|
op = 'Flatten',
|
|
input = inputs
|
|
),
|
|
'Flatten$' : lambda nodes, inputs, tensors, _:
|
|
Struct(
|
|
op = 'Flatten',
|
|
input = [inputs[-1]] # take only the last input, assume all other arguments are trivial (like sequence_length==1 always in ML-agents LSTM nets)
|
|
),
|
|
'Reshape' : lambda nodes, inputs, tensors, _:
|
|
Struct(
|
|
op = 'Reshape',
|
|
input = inputs,
|
|
shape = [tensors[0].data[0], tensors[0].data[1], tensors[0].data[2], tensors[0].data[3]] if len(tensors[0].data) == 4 else
|
|
[tensors[0].data[0], 1, tensors[0].data[1], tensors[0].data[2]] if len(tensors[0].data) == 3 else
|
|
[tensors[0].data[0], 1, 1, tensors[0].data[1]]
|
|
# tensor.name = 'shape'
|
|
),
|
|
'Multinomial' : lambda nodes, inputs, tensors, _:
|
|
Struct(
|
|
op = 'Multinomial',
|
|
input = inputs,
|
|
shape = [int(by_name(tensors, '/num_samples').data[0])],
|
|
#seed = get_attr(nodes[0], 'seed'),
|
|
),
|
|
'OneHot' : lambda nodes, inputs, tensors, _:
|
|
Struct(
|
|
op = 'OneHot',
|
|
input = inputs,
|
|
shape = [int(by_name(tensors, '/depth').data[0])],
|
|
alpha = by_name(tensors, '/on_value').data[0],
|
|
beta = by_name(tensors, '/off_value').data[0],
|
|
),
|
|
'Square' : lambda nodes, inputs, tensors, _:
|
|
Struct(
|
|
op = 'Mul',
|
|
input = [i for i in inputs] + [i for i in inputs], # input * input
|
|
),
|
|
'ConcatV2' : lambda nodes, inputs, tensors, _:
|
|
Struct(
|
|
op = 'Concat',
|
|
input = inputs,
|
|
|
|
# TEMPORARY: until we implemented rank detection and axis remapping (hopefully in exporter)
|
|
# HACK: assume Concat is always for last channel
|
|
axis = int(-1)
|
|
#axis = int(by_name(tensors, '/axis').data[0])
|
|
),
|
|
'BatchNormalization' : lambda nodes, inputs, tensors, _:
|
|
Struct(
|
|
op = 'BatchNormalization',
|
|
input = [i for i in inputs] +
|
|
order_by([t.name for t in tensors], ['gamma', 'beta', 'mean', 'variance']),
|
|
),
|
|
'Mean' : lambda nodes, inputs, tensors, _:
|
|
Struct(
|
|
# TODO: use data_frmt of the input instead of hardcoded [1,2] for HW
|
|
op = 'GlobalAveragePool' if np.array_equal(tensors[0].data, [1,2]) else 'MeanWithUnsupportedReductionTensor',
|
|
input = [i for i in inputs],
|
|
),
|
|
'Dense' : lambda nodes, inputs, tensors, _:
|
|
Struct(
|
|
op = 'Dense',
|
|
input = [i for i in inputs] + [t.name for t in tensors],
|
|
data_frmt = get_attr(by_op(nodes, 'Dense') or by_op(nodes, 'MatMul'), 'data_format'),
|
|
),
|
|
'Conv2D' : lambda nodes, inputs, tensors, _:
|
|
Struct(
|
|
op = 'Conv2D',
|
|
input = [i for i in inputs] + [t.name for t in tensors],
|
|
padding = get_attr(by_op(nodes, 'Conv2D'), 'padding'),
|
|
strides = get_attr(by_op(nodes, 'Conv2D'), 'strides'),
|
|
dilations = get_attr(by_op(nodes, 'Conv2D'), 'dilations'),
|
|
data_frmt = get_attr(by_op(nodes, 'Conv2D'), 'data_format'),
|
|
),
|
|
'DepthwiseConv2dNative' : lambda nodes, inputs, tensors, _:
|
|
Struct(
|
|
op = 'DepthwiseConv2dNative',
|
|
input = [i for i in inputs] + [t.name for t in tensors],
|
|
padding = get_attr(by_op(nodes, 'DepthwiseConv2dNative'), 'padding'),
|
|
strides = get_attr(by_op(nodes, 'DepthwiseConv2dNative'), 'strides'),
|
|
dilations = get_attr(by_op(nodes, 'DepthwiseConv2dNative'), 'dilations'),
|
|
data_frmt = get_attr(by_op(nodes, 'DepthwiseConv2dNative'), 'data_format'),
|
|
),
|
|
'Conv2DBackpropInput' : lambda nodes, inputs, tensors, _:
|
|
Struct(
|
|
op = 'Conv2DBackpropInput',
|
|
input = [i for i in inputs] + [t.name for t in tensors],
|
|
padding = get_attr(by_op(nodes, 'Conv2DBackpropInput'), 'padding'),
|
|
strides = get_attr(by_op(nodes, 'Conv2DBackpropInput'), 'strides'),
|
|
dilations = get_attr(by_op(nodes, 'Conv2DBackpropInput'), 'dilations'),
|
|
data_frmt = get_attr(by_op(nodes, 'Conv2DBackpropInput'), 'data_format'),
|
|
),
|
|
'BasicLSTM' : lambda nodes, inputs, tensors, context:
|
|
basic_lstm(nodes, inputs, tensors, context),
|
|
|
|
'Swish' : lambda nodes, inputs, tensors, _:
|
|
Struct(
|
|
op = 'Swish',
|
|
input = inputs
|
|
),
|
|
|
|
# TODO:'Round'
|
|
# TODO:'Rsqrt'
|
|
}
|
|
|
|
# Parse
|
|
def get_attr(node, attr_name, default=None):
|
|
if type(node) == Struct:
|
|
if hasattr(node, attr_name):
|
|
return getattr(node, attr_name)
|
|
else:
|
|
return default
|
|
|
|
# See: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/framework/attr_value.proto
|
|
val = node.attr[attr_name]
|
|
|
|
if val.HasField("list"):
|
|
return val.list.i
|
|
# NOTE: can't find way to identify type of list BUT it is almost always list(int)
|
|
# except list(float) in FractionalAvg/MaxPool
|
|
if val.HasField("b"):
|
|
return val.b
|
|
if val.HasField("i"):
|
|
return val.i
|
|
if val.HasField("f"):
|
|
return val.f
|
|
if val.HasField("s"):
|
|
return val.s.decode("utf-8")
|
|
if val.HasField("shape"):
|
|
return val.shape
|
|
if val.HasField("tensor"):
|
|
return val.tensor
|
|
return default
|
|
|
|
def get_epsilon(layer):
|
|
return get_attr(layer, 'epsilon', default=0.001) # default epsilon taken from tf.layers.batch_normalization
|
|
|
|
def get_layer_shape(layer):
|
|
shape = get_attr(layer, 'shape')
|
|
if not shape:
|
|
return [-1, -1, -1, -1]
|
|
shape = [dim.size for dim in shape.dim]
|
|
if len(shape) == 1:
|
|
return [1, 1, 1, shape[0]]
|
|
if len(shape) == 2:
|
|
return [shape[0], 1, 1, shape[1]]
|
|
return shape
|
|
|
|
def get_tensor_dims(tensor):
|
|
if isinstance(tensor, np.ndarray):
|
|
return np.shape(tensor)
|
|
|
|
dims = []
|
|
if tensor.tensor_shape:
|
|
dims = [v.size for v in tensor.tensor_shape.dim]
|
|
if tensor.float_val:
|
|
dims = np.shape(tensor.float_val)
|
|
if tensor.int_val:
|
|
dims = np.shape(tensor.int_val)
|
|
if tensor.bool_val:
|
|
dims = np.shape(tensor.bool_val)
|
|
return dims
|
|
|
|
def get_tensor_dtype(tensor):
|
|
if isinstance(tensor, np.ndarray):
|
|
return tensor.dtype
|
|
|
|
dataType = ''
|
|
fields = tensor.ListFields()
|
|
|
|
for field, value in fields:
|
|
if field.name == 'dtype' and field.cpp_type == descriptor.FieldDescriptor.CPPTYPE_ENUM:
|
|
dataType = field.enum_type.values_by_number.get(value, None).name
|
|
|
|
return dataType
|
|
|
|
def get_tensor_data(tensor):
|
|
if isinstance(tensor, np.ndarray):
|
|
return tensor.astype(float)
|
|
|
|
dims = get_tensor_dims(tensor)
|
|
elems = np.product(dims)
|
|
|
|
if tensor.tensor_content:
|
|
# TODO: support other types
|
|
dataType = get_tensor_dtype(tensor)
|
|
if dataType == "DT_FLOAT":
|
|
data = struct.unpack('<'+str(elems)+'f', tensor.tensor_content)
|
|
elif dataType == "DT_INT32":
|
|
data = struct.unpack('<'+str(elems)+'i', tensor.tensor_content)
|
|
elif dataType == "DT_BOOL":
|
|
data = struct.unpack('<'+str(elems)+'?', tensor.tensor_content)
|
|
else:
|
|
print('UNSUPPORTED: data type', dataType)
|
|
if tensor.float_val:
|
|
data = tensor.float_val
|
|
if tensor.int_val:
|
|
data = np.array(tensor.int_val, dtype=float)
|
|
if tensor.bool_val:
|
|
data = np.array(tensor.bool_val, dtype=float)
|
|
return np.array(data).reshape(dims)
|
|
|
|
def flatten(items,enter=lambda x:isinstance(x, list)):
|
|
# http://stackoverflow.com/a/40857703
|
|
# https://github.com/ctmakro/canton/blob/master/canton/misc.py
|
|
"""Yield items from any nested iterable; see REF."""
|
|
for x in items:
|
|
if enter(x):
|
|
yield from flatten(x)
|
|
else:
|
|
yield x
|
|
|
|
def replace_strings_in_list(array_of_strigs, replace_with_strings):
|
|
"A value in replace_with_strings can be either single string or list of strings"
|
|
potentially_nested_list = [replace_with_strings.get(s) or s for s in array_of_strigs]
|
|
return list(flatten(potentially_nested_list))
|
|
|
|
def remove_duplicates_from_list(array):
|
|
"Preserves the order of elements in the list"
|
|
output = []
|
|
unique = set()
|
|
for a in array:
|
|
if a not in unique:
|
|
unique.add(a)
|
|
output.append(a)
|
|
return output
|
|
|
|
#########################################################
|
|
|
|
def pool_to_HW(shape, data_frmt):
|
|
""" Convert from NHWC|NCHW => HW
|
|
"""
|
|
if len(shape) != 4:
|
|
return shape # Not NHWC|NCHW, return as is
|
|
if data_frmt == 'NCHW':
|
|
return [shape[2], shape[3]]
|
|
return [shape[1], shape[2]]
|
|
|
|
def strides_to_HW(shape, format):
|
|
return pool_to_HW(shape, format)
|
|
|
|
#########################################################
|
|
|
|
def gru(nodes, inputs, tensors, context):
|
|
assert(len(inputs) == 2)
|
|
|
|
def find_tensor_by_name(name, default=None):
|
|
nonlocal tensors
|
|
candidates = [t for t in tensors if t.name.endswith(name)]
|
|
return candidates[0].data if candidates else default
|
|
|
|
input = inputs[-1]
|
|
state = inputs[0]
|
|
gates_kernel = find_tensor_by_name('/gates/kernel')
|
|
gates_bias = find_tensor_by_name('/gates/bias', default=np.zeros(np.shape(gates_kernel)[-1]))
|
|
candidate_kernel = find_tensor_by_name('/candidate/kernel')
|
|
candidate_bias = find_tensor_by_name('/candidate/bias', default=np.zeros(np.shape(candidate_kernel)[-1]))
|
|
new_state = nodes[-1].name + '_h'
|
|
|
|
assert(np.shape(gates_kernel)[-1] == np.shape(gates_bias)[-1])
|
|
assert(np.shape(candidate_kernel)[-1] == np.shape(candidate_bias)[-1])
|
|
|
|
num_gates = 2
|
|
seq_length = 1
|
|
hidden_size = np.shape(gates_kernel)[-1] // num_gates
|
|
|
|
gate_kernels = np.split(gates_kernel, num_gates, axis=-1)
|
|
gate_biases = np.split(gates_bias, num_gates, axis=-1)
|
|
|
|
context.model_tensors['kernel_r'] = gate_kernels[0]
|
|
context.model_tensors['kernel_u'] = gate_kernels[1]
|
|
context.model_tensors['kernel_c'] = candidate_kernel
|
|
context.model_tensors['bias_r'] = gate_biases[0]
|
|
context.model_tensors['bias_u'] = gate_biases[1]
|
|
context.model_tensors['bias_c'] = candidate_bias
|
|
|
|
new_layers = barracuda.gru('gru', input, state,
|
|
'kernel_r', 'kernel_u', 'kernel_c',
|
|
'bias_r', 'bias_u', 'bias_c',
|
|
new_state)
|
|
|
|
state_shape = [1, 1, seq_length, hidden_size]
|
|
context.model_memories += [state_shape, state, new_state]
|
|
|
|
# map exptected output of the replaced pattern to output from our GRU cell
|
|
actual_output_node = nodes[-4]
|
|
assert(actual_output_node.op == 'Reshape')
|
|
context.map_ignored_layer_to_its_input[actual_output_node.name] = new_state
|
|
|
|
return new_layers
|
|
|
|
def basic_lstm(nodes, inputs, tensors, context):
|
|
assert(len(inputs) == 2)
|
|
|
|
def find_tensor_by_name(name, default=None):
|
|
nonlocal tensors
|
|
candidates = [t for t in tensors if t.name.endswith(name)]
|
|
return candidates[0].data if candidates else default
|
|
|
|
def find_forget_bias():
|
|
nonlocal nodes
|
|
nonlocal tensors
|
|
# TODO: make it more fault-tolerant
|
|
# search for scalar float constant that is input to Add node
|
|
# and hope it is not a constant for some complex activation function
|
|
for t in tensors:
|
|
if np.prod(t.shape) == 1 and get_tensor_dtype(t.obj) == "DT_FLOAT":
|
|
for n in nodes:
|
|
if n.op == 'Add' and t.name in n.input:
|
|
return t.data
|
|
return np.zeros(1)
|
|
|
|
input = inputs[-1]
|
|
state_c = inputs[0] + '_c'
|
|
state_h = inputs[0] + '_h'
|
|
kernel = find_tensor_by_name('/kernel')
|
|
bias = find_tensor_by_name('/bias', default=np.zeros(np.shape(kernel)[-1]))
|
|
forget_bias = find_forget_bias()
|
|
new_state_c = nodes[-1].name + '_c'
|
|
new_state_h = nodes[-1].name + '_h'
|
|
|
|
assert(np.shape(kernel)[-1] == np.shape(bias)[-1])
|
|
|
|
num_gates = 4
|
|
seq_length = 1
|
|
hidden_size = np.shape(kernel)[-1] // num_gates
|
|
|
|
kernels = np.split(kernel, num_gates, axis=-1)
|
|
biases = np.split(bias, num_gates, axis=-1)
|
|
|
|
context.model_tensors['kernel_i'] = kernels[0]
|
|
context.model_tensors['kernel_j'] = kernels[1]
|
|
context.model_tensors['kernel_f'] = kernels[2]
|
|
context.model_tensors['kernel_o'] = kernels[3]
|
|
context.model_tensors['bias_i'] = biases[0]
|
|
context.model_tensors['bias_j'] = biases[1]
|
|
context.model_tensors['bias_f'] = biases[2] + forget_bias
|
|
context.model_tensors['bias_o'] = biases[3]
|
|
|
|
new_layers = barracuda.lstm('lstm', input, state_c, state_h,
|
|
'kernel_i', 'kernel_j', 'kernel_f', 'kernel_o',
|
|
'bias_i', 'bias_j', 'bias_f', 'bias_o',
|
|
new_state_c, new_state_h)
|
|
|
|
state_shape = [1, 1, seq_length, hidden_size]
|
|
context.model_memories += [state_shape, state_c, new_state_c]
|
|
context.model_memories += [state_shape, state_h, new_state_h]
|
|
|
|
# map expected output of the replaced pattern to output from our LSTM cell
|
|
actual_output_node = nodes[-4]
|
|
assert(actual_output_node.op == 'Reshape')
|
|
context.map_ignored_layer_to_its_input[actual_output_node.name] = new_state_h
|
|
|
|
return new_layers
|
|
|
|
#########################################################
|
|
|
|
def process_layer(layer, context, args):
|
|
model_tensors = context.model_tensors
|
|
input_shapes = context.input_shapes
|
|
map_ignored_layer_to_its_input = context.map_ignored_layer_to_its_input
|
|
|
|
name = layer.name
|
|
class_name = layer.op
|
|
inputs = layer.input # Tensorflow inputs are always explicit, but in case of Keras we had 'inputs = layer.input or [prev_layer_name]'
|
|
inputs = replace_strings_in_list(inputs, map_ignored_layer_to_its_input)
|
|
|
|
if class_name == 'Const':
|
|
model_tensors[name] = layer.attr["value"].tensor
|
|
return
|
|
|
|
if class_name == 'Placeholder':
|
|
assert(inputs == [])
|
|
map_ignored_layer_to_its_input[name] = inputs
|
|
input_shapes[name] = get_layer_shape(layer)
|
|
return
|
|
|
|
if class_name == 'Identity':
|
|
connected_to_const = len(inputs) == 1 and inputs[0] in model_tensors
|
|
if connected_to_const:
|
|
map_ignored_layer_to_its_input[name] = inputs
|
|
return
|
|
else:
|
|
# treat Identity layer that are connected to processing nodes
|
|
# as output from the network
|
|
class_name = 'Linear'
|
|
|
|
# TEMPORARY: until we implemented rank detection and StidedSlice at runtime
|
|
# HACK: skips trivial StridedSlices for rank=2 tensors
|
|
if class_name == 'StridedSlice' and get_attr(layer, 'begin_mask') == 1 and get_attr(layer, 'end_mask') == 1:
|
|
map_ignored_layer_to_its_input[name] = inputs[0]
|
|
return
|
|
|
|
if args.print_layers or args.verbose:
|
|
var_tensors = [i for i in inputs if i not in model_tensors]
|
|
const_tensors = [i for i in inputs if i in model_tensors]
|
|
print("'%s' %s Vars:%s Const:%s" % (name, class_name, var_tensors, const_tensors))
|
|
|
|
if class_name in known_activations:
|
|
activation = class_name
|
|
class_name = 'Activation'
|
|
else:
|
|
activation = 'Linear'
|
|
|
|
if not class_name in known_classes:
|
|
if class_name in requires_runtime_flag:
|
|
print('SKIP:', class_name, 'layer is used only for training')
|
|
else:
|
|
print('IGNORED:', class_name, 'unknown layer')
|
|
map_ignored_layer_to_its_input[name] = inputs
|
|
return
|
|
|
|
klass = known_classes[class_name]
|
|
if type(klass) == int:
|
|
klass = Struct(id = klass)
|
|
|
|
o_l = Struct()
|
|
o_l.type = klass.id
|
|
o_l.class_name = class_name
|
|
o_l.name = name
|
|
|
|
padding = get_attr(layer, 'padding') # layer.attr['padding'].s.decode("utf-8")
|
|
strides = get_attr(layer, 'strides') # layer.attr['strides'].list.i
|
|
dilations = get_attr(layer, 'dilations') # layer.attr['dilations'].list.i
|
|
pool_size = get_attr(layer, 'ksize') # layer.attr['ksize'].list.i
|
|
shape = get_attr(layer, 'shape', default=[])
|
|
data_frmt = get_attr(layer, 'data_format') # layer.attr['data_format'].s.decode("utf-8")
|
|
axis = get_attr(layer, 'axis')
|
|
alpha = get_attr(layer, 'alpha')
|
|
beta = get_attr(layer, 'beta')
|
|
|
|
if activation and not activation in known_activations:
|
|
print('IGNORED: unknown activation', activation)
|
|
if padding and not padding in known_paddings:
|
|
print('IGNORED: unknown padding', padding)
|
|
if data_frmt and not data_frmt in supported_data_formats:
|
|
print('UNSUPPORTED: data format', data_frmt)
|
|
|
|
o_l.activation = known_activations.get(activation) or 0
|
|
o_l.pads = known_paddings.get(padding) or [0,0,0,0]
|
|
o_l.strides = strides_to_HW(strides, data_frmt) if strides else []
|
|
o_l.pool_size = pool_to_HW(pool_size, data_frmt) if pool_size else shape
|
|
o_l.axis = axis or -1
|
|
o_l.alpha = alpha or 1
|
|
o_l.beta = beta or 0
|
|
|
|
tensor_names = [i for i in inputs if i in model_tensors]
|
|
o_l.tensors = [Struct(name = x, shape = get_tensor_dims(model_tensors[x]), data = get_tensor_data(model_tensors[x]))
|
|
for x in tensor_names]
|
|
# Patch shapes & data
|
|
layer_has_model_tensors = len(o_l.tensors) > 0
|
|
if hasattr(klass, 'out_shapes') and layer_has_model_tensors:
|
|
shapes = klass.out_shapes([x.shape for x in o_l.tensors])
|
|
|
|
# if we have more shapes than actual tensors,
|
|
# then create & fill missing tensors with zeros
|
|
in_tensor_num = len(o_l.tensors)
|
|
for index, new_shape in enumerate(shapes):
|
|
if index >= in_tensor_num:
|
|
new_tensor = Struct(name = ('%s/patch:%i') % (name, index-in_tensor_num),
|
|
shape = new_shape,
|
|
data = np.zeros(new_shape))
|
|
o_l.tensors.append(new_tensor)
|
|
assert(len(shapes) <= len(o_l.tensors))
|
|
|
|
if hasattr(klass, 'patch_data'):
|
|
data = [x.data for x in o_l.tensors]
|
|
|
|
patch_data_fn = klass.patch_data
|
|
patch_data_expected_arg_count = patch_data_fn.__code__.co_argcount
|
|
patch_data_args = (data, layer) if patch_data_expected_arg_count > 1 else (data,)
|
|
tensor_data = patch_data_fn(*patch_data_args)
|
|
o_l.tensors = o_l.tensors[:len(tensor_data)] # resize tensor array to match patched data - patching might reduce number of tensors
|
|
for x, data in zip(o_l.tensors, tensor_data):
|
|
x.data = data
|
|
|
|
# after this point we should have equal amount of shapes and tensors
|
|
assert(len(o_l.tensors) == len(shapes))
|
|
|
|
for x, shape in zip(o_l.tensors, shapes):
|
|
x.shape = shape
|
|
|
|
o_l.inputs = [i for i in inputs if i not in model_tensors]
|
|
|
|
else:
|
|
# no 'patch_data' lambda was specified, op does not require tensor args
|
|
o_l.tensors = []
|
|
o_l.inputs = inputs
|
|
|
|
# Force all tensors to float32
|
|
for x in o_l.tensors:
|
|
x.data = x.data.astype(np.float32)
|
|
|
|
# Layer is ready
|
|
context.layers.append(o_l)
|
|
|
|
class ModelBuilderContext:
|
|
def __init__(self):
|
|
self.layers = []
|
|
self.input_shapes = {}
|
|
self.model_tensors = {}
|
|
self.model_memories = []
|
|
self.map_ignored_layer_to_its_input = {}
|
|
|
|
def process_model(model, args):
|
|
o_context = ModelBuilderContext()
|
|
|
|
# Find node patterns
|
|
nodes_as_array = [node for node in model.node]
|
|
|
|
node_index = 0
|
|
while node_index < len(nodes_as_array):
|
|
node = nodes_as_array[node_index]
|
|
match = False
|
|
for pattern_repr, pattern_name in known_patterns.items():
|
|
pattern = eval(pattern_repr)
|
|
if node_index + len(pattern) > len(nodes_as_array):
|
|
continue # pattern too long, skip
|
|
|
|
require_exact_match = (pattern[0] == 'Const' or pattern[0] == 'Identity')
|
|
pattern_end = node_index
|
|
|
|
def match_node(node, pattern):
|
|
return node.op == pattern or (hasattr(pattern, 'match') and pattern.match(node.name))
|
|
|
|
for p in pattern:
|
|
if not require_exact_match:
|
|
while pattern_end < len(nodes_as_array) and nodes_as_array[pattern_end].op != p and (
|
|
nodes_as_array[pattern_end].op == 'Const' or
|
|
nodes_as_array[pattern_end].op == 'Identity'):
|
|
pattern_end += 1
|
|
if pattern_end >= len(nodes_as_array):
|
|
break
|
|
|
|
match = False
|
|
if (hasattr(p, 'match')): # regexp
|
|
while pattern_end < len(nodes_as_array) and p.match(nodes_as_array[pattern_end].name):
|
|
match = True
|
|
pattern_end += 1
|
|
else: # exact string
|
|
match = nodes_as_array[pattern_end].op == p
|
|
pattern_end += 1
|
|
|
|
if not match:
|
|
break
|
|
|
|
def get_tensors(pattern_nodes):
|
|
nonlocal o_context
|
|
map_ignored_layer_to_its_input = o_context.map_ignored_layer_to_its_input
|
|
|
|
# tensors <= all Const nodes within this pattern
|
|
tensor_nodes = [n for n in pattern_nodes if n.op == 'Const']
|
|
tensors = [Struct(name = n.name, obj = n.attr["value"].tensor, shape = get_tensor_dims(n.attr["value"].tensor), data = get_tensor_data(n.attr["value"].tensor))
|
|
for n in tensor_nodes]
|
|
|
|
# TODO: unify / reuse code from process_layer
|
|
identity_nodes = [n for n in pattern_nodes if n.op == 'Identity']
|
|
for i in identity_nodes:
|
|
inputs = replace_strings_in_list(i.input, map_ignored_layer_to_its_input)
|
|
map_ignored_layer_to_its_input[i.name] = inputs
|
|
|
|
# gather inputs from Op nodes (not Const, not Identity)
|
|
op_nodes = [n for n in pattern_nodes if n not in tensor_nodes and n not in identity_nodes]
|
|
inputs_to_op_nodes = list(flatten([list(flatten(n.input)) for n in op_nodes]))
|
|
inputs_to_op_nodes = replace_strings_in_list(inputs_to_op_nodes, map_ignored_layer_to_its_input)
|
|
inputs_to_op_nodes = [i.split(':')[0] for i in inputs_to_op_nodes]
|
|
|
|
# filter only inputs that are coming from nodes that are outside this pattern
|
|
# preserve the order
|
|
pattern_nodes = [n.name for n in pattern_nodes]
|
|
#inputs_from_outside_pattern = remove_duplicates_from_list([i for i in inputs_to_op_nodes if nodes_by_name[i] not in pattern_nodes])
|
|
inputs_from_outside_pattern = remove_duplicates_from_list([i for i in inputs_to_op_nodes if i not in pattern_nodes])
|
|
|
|
return inputs_from_outside_pattern, tensors
|
|
|
|
if match:
|
|
nodes = nodes_as_array[node_index:pattern_end]
|
|
name = nodes[-1].name
|
|
var_tensors, const_tensors = get_tensors(nodes)
|
|
if args.print_patterns or args.verbose:
|
|
print('PATTERN:', name, '~~', pattern_name, pattern, '<-', var_tensors, '+', [t.name for t in const_tensors])
|
|
for n in nodes:
|
|
if n.op == 'Const' or n.op == 'Identity':
|
|
process_layer(n, o_context, args)
|
|
|
|
new_layers = transform_patterns[pattern_name](nodes, var_tensors, const_tensors, o_context)
|
|
if not isinstance(new_layers, list):
|
|
if not hasattr(new_layers, name): new_layers.name = name
|
|
new_layers = [new_layers]
|
|
|
|
for l in new_layers:
|
|
# TODO: prefix new layer names with scope, patch inputs
|
|
#l.name = name + '/' + l.name
|
|
process_layer(l, o_context, args)
|
|
|
|
node_index = pattern_end
|
|
break # pattern found & processed
|
|
|
|
if not match:
|
|
# TODO: gather tensors in the same way as patterns do
|
|
process_layer(node, o_context, args)
|
|
node_index += 1
|
|
|
|
return o_context.layers, o_context.input_shapes, o_context.model_tensors, o_context.model_memories
|
|
|
|
#########################################################
|
|
|
|
def convert(source_file, target_file, trim_unused_by_output="", verbose=False, compress_f16=False):
|
|
"""
|
|
Converts a TensorFlow model into a Barracuda model.
|
|
:param source_file: The TensorFlow Model
|
|
:param target_file: The name of the file the converted model will be saved to
|
|
:param trim_unused_by_output: The regexp to match output nodes to remain in the model. All other uconnected nodes will be removed.
|
|
:param verbose: If True, will display debug messages
|
|
:param compress_f16: If true, the float values will be converted to f16
|
|
:return:
|
|
"""
|
|
if (type(verbose)==bool):
|
|
args = Struct()
|
|
args.verbose = verbose
|
|
args.print_layers = verbose
|
|
args.print_source_json = verbose
|
|
args.print_barracuda_json = verbose
|
|
args.print_layer_links = verbose
|
|
args.print_patterns = verbose
|
|
args.print_tensors = verbose
|
|
else:
|
|
args = verbose
|
|
|
|
# Load Tensorflow model
|
|
print("Converting %s to %s" % (source_file, target_file))
|
|
f = open(source_file, 'rb')
|
|
i_model = tf.GraphDef()
|
|
i_model.ParseFromString(f.read())
|
|
|
|
if args.verbose:
|
|
print('OP_TYPES:', {layer.op for layer in i_model.node})
|
|
|
|
if args.print_source_json or args.verbose:
|
|
for layer in i_model.node:
|
|
if not layer.op == 'Const':
|
|
print('MODEL:', MessageToJson(layer) + ",")
|
|
|
|
# Convert
|
|
o_model = barracuda.Model()
|
|
o_model.layers, o_input_shapes, o_model.tensors, o_model.memories = \
|
|
process_model(i_model, args)
|
|
|
|
# Cleanup unconnected Identities (they might linger after processing complex node patterns like LSTM)
|
|
def cleanup_layers(layers):
|
|
all_layers = {l.name for l in layers}
|
|
all_inputs = {i for l in layers for i in l.inputs}
|
|
|
|
def is_unconnected_identity(layer):
|
|
if layer.class_name == 'Activation' and layer.activation == 0: # Identity
|
|
assert(len(layer.inputs) == 1)
|
|
if layer.inputs[0] not in all_layers and layer.name not in all_inputs:
|
|
return True;
|
|
return False;
|
|
|
|
return [l for l in layers if not is_unconnected_identity(l)]
|
|
o_model.layers = cleanup_layers(o_model.layers)
|
|
|
|
all_inputs = {i for l in o_model.layers for i in l.inputs}
|
|
embedded_tensors = {t.name for l in o_model.layers for t in l.tensors}
|
|
|
|
# Find global tensors
|
|
def dims_to_barracuda_shape(dims):
|
|
shape = list(dims)
|
|
while len(shape) < 4:
|
|
shape = [1] + shape
|
|
return shape
|
|
o_model.globals = [t for t in o_model.tensors if t not in all_inputs and t not in embedded_tensors]
|
|
#for x in global_tensors:
|
|
# shape = dims_to_barracuda_shape(get_tensor_dims(o_model.tensors[x]))
|
|
# o_globals += [Struct(
|
|
# name = x,
|
|
# shape = shape,
|
|
# data = np.reshape(get_tensor_data(o_model.tensors[x]), shape).astype(np.float32))]
|
|
|
|
# Trim
|
|
if trim_unused_by_output:
|
|
o_model.layers = barracuda.trim(o_model.layers, trim_unused_by_output, args.verbose)
|
|
|
|
# Create load layers for constants
|
|
const_tensors = [i for i in all_inputs if i in o_model.tensors]
|
|
const_tensors += o_model.globals
|
|
for x in const_tensors:
|
|
shape = dims_to_barracuda_shape(get_tensor_dims(o_model.tensors[x]))
|
|
|
|
o_l = Struct(
|
|
type = 255, # Load
|
|
class_name = "Const",
|
|
name = x,
|
|
pads = [0,0,0,0],
|
|
strides = [],
|
|
pool_size = [],
|
|
axis = -1,
|
|
alpha = 1,
|
|
beta = 0,
|
|
activation = 0,
|
|
inputs = [],
|
|
tensors = [Struct(
|
|
name = x,
|
|
shape = shape,
|
|
data = np.reshape(get_tensor_data(o_model.tensors[x]), shape).astype(np.float32))]
|
|
)
|
|
o_model.layers.insert(0, o_l)
|
|
|
|
# Find model inputs & outputs
|
|
all_layers = {l.name for l in o_model.layers}
|
|
# global inputs => are inputs that are NOT connected to any layer in the network
|
|
# global outputs => are outputs that are NOT feeding any layer in the network OR are coming from Identity layers
|
|
o_model.inputs = {i:o_input_shapes[i] for l in o_model.layers for i in l.inputs if i not in all_layers and i not in o_model.memories}
|
|
|
|
def is_output_layer(layer):
|
|
if layer.class_name == 'Const': # Constants never count as global output even when unconnected
|
|
return False;
|
|
if layer.name not in all_inputs: # this layer is not inputing to any other layer
|
|
return True
|
|
if layer.class_name == 'Activation' and layer.activation == 0: # Identity marks global output
|
|
return True
|
|
return False
|
|
o_model.outputs = [l.name for l in o_model.layers if is_output_layer(l)]
|
|
|
|
# Compress
|
|
if compress_f16:
|
|
o_model = barracuda.compress(o_model)
|
|
|
|
# Sort model so that layer inputs are always ready upfront
|
|
o_model.layers = barracuda.sort(o_model.layers, o_model.inputs, o_model.memories, args.verbose)
|
|
|
|
# Summary
|
|
barracuda.summary(o_model,
|
|
print_layer_links = args.print_layer_links or args.verbose,
|
|
print_barracuda_json = args.print_barracuda_json or args.verbose,
|
|
print_tensors = args.print_tensors or args.verbose)
|
|
|
|
# Write to file
|
|
barracuda.write(o_model, target_file)
|
|
print('DONE: wrote', target_file, 'file.')
|