def version_9(cls, ctx, node, **kwargs):
node.type = "ConstantOfShape"
# both shape and value in tensorflow are passed as tensor.
# In onnx the value is an attribute so we need to fetch the value as const which
# sooner or later will be a problem for tensorflow-onnx.
# ConstantOfShape in onnxruntime only support int64, so insert cast op
input_dtype_is_int64 = utils.map_onnx_to_numpy_type(ctx.get_dtype(node.input[0])) == np.int64
if not input_dtype_is_int64:
ctx.insert_new_node_on_input(node, "Cast", node.input[0], to=onnx_pb.TensorProto.INT64)
dtype = ctx.get_dtype(node.output[0])
value = np.array([node.inputs[1].get_tensor_value()]).astype(utils.map_onnx_to_numpy_type(dtype))
value_proto = numpy_helper.from_array(value)
node.set_attr("value", value_proto)
ctx.remove_input(node, node.input[1], 1)
def version_7(cls, ctx, node, **kwargs):
node.type = "BatchNormalization"
# tf inputs: x, scale, bias, mean, variance
# tf outputs: y, batch_mean, batch_var
# a: data_format, epsilon, is_training
# onnx inputs: X, scale, B, mean, variance, attributes: epsilon, momentum=0.9, spatial : 1
# output: y, mean, var, savedmean, savedvar,
# detach unused outputs. While we could let the unused outputs dangle,
# some runtimes like pytorch/caffe2 do complain about it.
consumers = [ctx.find_output_consumers(output_name) for output_name in node.output[1:]]
if not any(consumers):
new_output = [node.output[0]]
node.output = new_output
conv_convert_inputs(ctx, node, with_kernel=False)
scale_shape = ctx.get_shape(node.input[1])
mean_shape = ctx.get_shape(node.input[3])
var_shape = ctx.get_shape(node.input[4])
val_type = utils.map_onnx_to_numpy_type(ctx.get_dtype(node.input[1]))
if mean_shape != scale_shape:
new_mean_value = np.array(np.resize(node.inputs[3].get_tensor_value(as_list=False), scale_shape),
dtype=val_type)
new_mean_node_name = utils.make_name(node.name)
ctx.make_const(new_mean_node_name, new_mean_value)
node.input[3] = new_mean_node_name
if var_shape != scale_shape:
new_var_value = np.array(np.resize(node.inputs[4].get_tensor_value(as_list=False), scale_shape),
dtype=val_type)
new_val_node_name = utils.make_name(node.name)
ctx.make_const(new_val_node_name, new_var_value)
node.input[4] = new_val_node_name
def version_1(cls, ctx, node, **kwargs):
# Not exactly nearest even but close enough
np_dtype = utils.map_onnx_to_numpy_type(ctx.get_dtype(node.input[0]))
const_half = ctx.make_const(utils.make_name("const_half"), np.array(0.5, np_dtype)).output[0]
add_node = ctx.make_node("Add", [node.input[0], const_half], op_name_scope=node.name).output[0]
node.type = "Floor"
ctx.replace_inputs(node, [add_node])
def version_9(cls, ctx, node, **kwargs):
node.type = "Div"
np_dtype = utils.map_onnx_to_numpy_type(ctx.get_dtype(node.input[1]))
zero_const = ctx.make_const(utils.make_name("const_zero"), np.array(0, np_dtype)).output[0]
is_zero = ctx.make_node("Equal", [node.input[1], zero_const]).output[0]
where_node = ctx.make_node("Where", [is_zero, zero_const, node.output[0]])
ctx.insert_node_on_output(where_node, node.output[0])
def version_7(cls, ctx, node, **kwargs):
# make subgraph to implement one_hot, idea comes from onehot_op
indices_name = node.input[1]
indices_shape = ctx.get_shape(indices_name)
if len(indices_shape) != 1:
# TODO: this works for rank=1 but tensorflow supports more than this.
# Same principle should work but we need to implement our own eye.
raise ValueError("onehot op: only rank1 is supported")
logit_name = node.input[0]
depth = ctx.get_shape(logit_name)[-1]
# if number of classes is unknown or too large
if depth == utils.ONNX_UNKNOWN_DIMENSION or depth > 20000:
sparse_softmax_cross_entropy_with_logits_op_by_gathernd(ctx, node, **kwargs)
return
logit_dtype = ctx.get_dtype(logit_name)
utils.make_sure(logit_dtype, "Dtype of {} is None".format(logit_name))
dtype = utils.map_onnx_to_numpy_type(logit_dtype)
eye = np.eye(depth).astype(dtype)
const_name = utils.make_name("const_eye")
const_eye = ctx.make_const(name=const_name, np_val=eye)
onehot = ctx.make_node(op_type="Gather", inputs=[const_eye.output[0], indices_name], attr={"axis": 0})
log_softmax = ctx.make_node(op_type="LogSoftmax", inputs=[logit_name])
# implement tf.multiply(np.float32(-1.0), tf.reduce_sum(tf.multiply(one_hot, log_softmax), axis=1))
mul1 = ctx.make_node(op_type="Mul", inputs=[onehot.output[0], log_softmax.output[0]])
reduce_sum = ctx.make_node(op_type="ReduceSum", inputs=[mul1.output[0]], attr={"axes": [1]})
const_name = utils.make_name("const_negative_one")
const_negative_one = ctx.make_const(name=const_name, np_val=np.array(-1).astype(dtype))
mul2 = ctx.make_node(op_type="Mul", inputs=[const_negative_one.output[0], reduce_sum.output[0]])
shapes = node.output_shapes
dtypes = node.output_dtypes
ctx.remove_node(node.name)
ctx.make_node(op_type="Squeeze", inputs=[mul2.output[0]], outputs=[node.output[0]], attr={"axes": [1]},
shapes=[shapes[0]], dtypes=[dtypes[0]])
def rewrite_ragged_variant_shape(g, ops):
pattern1 = \
OpTypePattern('Shape', name='shape', inputs=[
OpTypePattern('RaggedTensorToVariant', name='raggedtovariant')
])
pattern_list = [pattern1]
for pattern in pattern_list:
matcher = GraphMatcher(pattern)
match_results = list(matcher.match_ops(ops))
for match in match_results:
shape = match.get_op('shape')
raggedtovariant = match.get_op('raggedtovariant')
if raggedtovariant.get_attr_value("batched_input") != 1:
continue
if raggedtovariant.get_attr_value("RAGGED_RANK") != 1:
continue
# Shape of batched variant from ragged is same as number of splits minus 1
g.replace_inputs(shape, [raggedtovariant.input[0]])
np_dtype = utils.map_onnx_to_numpy_type(
g.get_dtype(shape.output[0]))
const_one = g.make_const(utils.make_name("const_one"),
np.array(1, np_dtype)).output[0]
g.insert_new_node_on_output("Sub",
shape.output[0],
inputs=[shape.output[0], const_one])
return ops
def version_10(cls, ctx, node, **kwargs):
scale = node.get_attr_value('scale')
zero_point = node.get_attr_value('zero_point')
axis = node.get_attr_value('quantized_dimension')
np_q_type = utils.map_onnx_to_numpy_type(ctx.get_dtype(node.input[0]))
if len(scale) > 1 or len(zero_point) > 1:
utils.make_sure(ctx.opset >= 13,
"Opset 13 is required for per-axis quantization")
node.set_attr("axis", axis)
scale_node = ctx.make_const(utils.make_name("scale"),
np.array(scale, dtype=np.float32))
zero_point_node = ctx.make_const(
utils.make_name("zero_point"),
np.array(zero_point, dtype=np_q_type))
else:
scale_node = ctx.make_const(utils.make_name("scale"),
np.array(scale[0], dtype=np.float32))
zero_point_node = ctx.make_const(
utils.make_name("zero_point"),
np.array(zero_point[0], dtype=np_q_type))
ctx.replace_inputs(
node,
[node.input[0], scale_node.output[0], zero_point_node.output[0]])
del node.attr["scale"]
del node.attr["zero_point"]
del node.attr["quantized_dimension"]
def version_1(cls, ctx, node, **kwargs):
# T output = ZerosLike(T x)
# when params "dtype" used, tf will call another op "Fill" instead, so Cast is not needed here.
input_dtype = ctx.get_dtype(node.input[0])
node_name = utils.make_name("zero")
const_zero = ctx.make_const(node_name, np.array(0).astype(utils.map_onnx_to_numpy_type(input_dtype)))
shapes = node.output_shapes
dtypes = node.output_dtypes
ctx.remove_node(node.name)
ctx.make_node(op_type="Mul", inputs=[node.input[0], const_zero.output[0]],
name=node.name, outputs=node.output, shapes=shapes, dtypes=dtypes)
def version_1(cls, ctx, node, **kwargs):
node.type = "Mul"
node_shape = ctx.get_shape(node.input[0])
node_dtype = ctx.get_dtype(node.input[0])
node_np_dtype = utils.map_onnx_to_numpy_type(node_dtype)
const_two = ctx.make_const(utils.make_name("const_two"), np.array([2]).astype(node_np_dtype)).output[0]
node.input.append(const_two)
const_one = ctx.make_const(utils.make_name("const_one"), np.ones(node_shape, dtype=node_np_dtype)).output[0]
op_name = utils.make_name(node.name)
ctx.insert_new_node_on_output("Add", node.output[0], inputs=[node.output[0], const_one], name=op_name)
def any_version(cls, opset, ctx, node, **kwargs):
"""
Computes the modules of a complex.
If the matrix dtype is not complex64 or complex128,
it assumes the first dimension means real part (0)
and imaginary part (1, :, :...).
"""
supported_dtypes = [
onnx_pb.TensorProto.FLOAT,
onnx_pb.TensorProto.FLOAT16,
onnx_pb.TensorProto.DOUBLE,
onnx_pb.TensorProto.COMPLEX64,
onnx_pb.TensorProto.COMPLEX128,
]
onnx_dtype = ctx.get_dtype(node.input[0])
utils.make_sure(onnx_dtype in supported_dtypes, "Unsupported input type.")
shape = ctx.get_shape(node.input[0])
np_dtype = utils.map_onnx_to_numpy_type(onnx_dtype)
utils.make_sure(shape[0] == 2, "ComplexAbs expected the first dimension to be 2 but shape is %r", shape)
ind0 = ctx.make_const(name=utils.make_name('cst0'), np_val=np.array([0], dtype=np.int64))
ind1 = ctx.make_const(name=utils.make_name('cst1'), np_val=np.array([1], dtype=np.int64))
p2 = ctx.make_const(name=utils.make_name('p2'), np_val=np.array([2], dtype=np_dtype))
real_part = ctx.make_node(
'Gather', inputs=[node.input[0], ind0.name], attr=dict(axis=0),
name=utils.make_name('Real_' + node.name))
imag_part = ctx.make_node(
'Gather', inputs=[node.input[0], ind1.name], attr=dict(axis=0),
name=utils.make_name('Imag_' + node.name))
real_part2 = ctx.make_node(
'Pow', inputs=[real_part.output[0], p2.name],
name=utils.make_name(real_part.name + 'p2p'))
imag_part2 = ctx.make_node(
'Pow', inputs=[imag_part.output[0], p2.name],
name=utils.make_name(imag_part.name + 'p2p'))
ctx.remove_node(node.name)
add = ctx.make_node(
"Add", inputs=[real_part2.output[0], imag_part2.output[0]],
name=utils.make_name('ComplexAbs_' + node.name))
squeezed = GraphBuilder(ctx).make_squeeze(
{'data': add.output[0], 'axes': [0]}, name=utils.make_name('ComplexAbs' + node.name), return_node=True)
last_node = ctx.make_node(
"Sqrt", inputs=squeezed.output[:1],
name=utils.make_name('ComplexAbs' + node.name),
shapes=[shape[1:]], dtypes=[onnx_dtype])
ctx.replace_all_inputs(node.output[0], last_node.output[0]) # ops=ctx.get_nodes()
def get_weights_from_const_node(g, node):
temp = node
val = None
# this would help ignore Identity in non-const_folded graph.
while temp.type == 'Identity':
temp = temp.inputs[0]
if temp and temp.type == 'Const':
val = temp.get_tensor_value(as_list=False)
dtype = utils.map_onnx_to_numpy_type(g.get_dtype(temp.output[0]))
val = val.astype(dtype)
logger.debug("found weights %s", temp.name)
else:
logger.debug("weight node seems not to be Const, skip, node name is %s", temp.name)
return None
return val
def sparse_softmax_cross_entropy_with_logits_op_by_gathernd(ctx, node, **kwargs):
# make subgraph to implement one_hot, idea comes from onehot_op
indices_name = node.input[1]
indices_shape = ctx.get_shape(indices_name)
if len(indices_shape) != 1:
# TODO: this works for rank=1 but tensorflow supports more than this.
# Same principle should work but we need to implement our own eye.
raise ValueError("onehot op: only rank1 is supported")
logit_name = node.input[0]
logit_dtype = ctx.get_dtype(logit_name)
logit_shape = ctx.get_shape(logit_name)
utils.make_sure(logit_dtype, "Dtype of {} is None".format(logit_name))
indices_dtype = ctx.get_dtype(indices_name)
if indices_dtype != TensorProto.INT64:
indices_cast = ctx.make_node("Cast", [indices_name], attr={"to": TensorProto.INT64})
indices_name = indices_cast.output[0]
indices_size = ctx.make_node("Size", [indices_name])
indices_unsqueeze = ctx.make_node("Unsqueeze", [indices_name], attr={"axes": [1]})
zero_const = ctx.make_const(utils.make_name("zero"), np.array(0, dtype=np.int64))
one_const = ctx.make_const(utils.make_name("one"), np.array(1, dtype=np.int64))
id_name = utils.make_name("sparse_softmax_id")
id_output = utils.port_name(id_name)
controlflow.make_range(ctx, zero_const.output[0], indices_size.output[0], one_const.output[0],
id_output, id_name, shape=[-1], dtype=TensorProto.INT64)
id_unsqueeze = ctx.make_node("Unsqueeze", [id_output], attr={"axes": [1]})
indices_with_id = ctx.make_node("Concat",
[id_unsqueeze.output[0], indices_unsqueeze.output[0]],
attr={"axis": 1})
log_softmax = ctx.make_node(op_type="LogSoftmax",
inputs=[logit_name], dtypes=[logit_dtype], shapes=[logit_shape])
gathernd_name = utils.make_name("sparse_softmax_gathernd")
gathernd_output = utils.port_name(gathernd_name)
tensor.make_gathernd(ctx, log_softmax.output[0], indices_with_id.output[0], gathernd_output,
gathernd_name, logit_dtype, [logit_shape], [logit_dtype])
const_name = utils.make_name("const_negative_one")
const_negative_one = ctx.make_const(const_name, np.array(-1).astype(utils.map_onnx_to_numpy_type(logit_dtype)))
mul2 = ctx.make_node(op_type="Mul", inputs=[const_negative_one.output[0], gathernd_output])
shapes = node.output_shapes
dtypes = node.output_dtypes
ctx.remove_node(node.name)
ctx.make_node(op_type="Squeeze",
inputs=[mul2.output[0]], outputs=[node.output[0]],
attr={"axes": [1]}, shapes=[shapes[0]], dtypes=[dtypes[0]])
def _workaround_fill_ch_init_node(self, initializer_input_id, rnn_props):
node = self.g.get_node_by_output(initializer_input_id)
if node.type != "Fill":
return None
fill_val = node.inputs[1].get_tensor_value()
fill_val_dtype = utils.map_onnx_to_numpy_type(
self.g.get_dtype(node.input[1]))
# this must be int64, since Concat's input data type must be consistent.
num_direction_node = self.g.make_const(utils.make_name("Const"),
np.array([1], dtype=np.float32))
h_node = self.g.make_const(
utils.make_name("Const"),
np.array([rnn_props.hidden_size], dtype=np.float32))
b_node = rnn_props.batch_size_node
# Concat in OPSET7 does not support int64.
tile_shape = self.g.make_node(
"Concat",
[num_direction_node.output[0], b_node.output[0], h_node.output[0]],
attr={"axis": 0})
# Tile's repeats must be INT64
attr = {"to": onnx_pb.TensorProto.INT64}
tile_shape_int64 = self.g.make_node("Cast", [tile_shape.output[0]],
attr)
const_node = self.g.make_const(
utils.make_name("Const"),
np.array([[[fill_val]]], dtype=fill_val_dtype))
tile_node = self.g.make_node(
"Tile", [const_node.output[0], tile_shape_int64.output[0]])
self.all_nodes.extend([
tile_shape, tile_shape_int64, tile_node, num_direction_node,
h_node, const_node
])
return tile_node
def any_version(cls, opset, ctx, node, **kwargs):
node_inputs = node.input.copy()
num_segments_specified = False
num_segs_const = None
if node.type.endswith("WithNumSegments") or node.type.startswith(
"Unsorted"):
num_segments_specified = True
num_segments = node_inputs.pop()
node.type = node.type.replace("WithNumSegments", "")
node.type = node.type.replace("Unsorted", "")
if node.inputs[-1].is_const():
num_segs_const = node.inputs[-1].get_tensor_value(as_list=True)
if node.type.startswith("Sparse"):
data_inp, indices_inp, segment_inp = node_inputs
gather_node = ctx.make_node("Gather", [data_inp, indices_inp],
attr={'axis': 0})
data_inp = gather_node.output[0]
node.type = node.type.replace("Sparse", "")
else:
data_inp, segment_inp = node_inputs
# Data has shape [n, a, b, ..., c]
data_shape = ctx.get_shape(data_inp)
data_rank = len(data_shape) if data_shape is not None else None
data_dtype = ctx.get_dtype(data_inp)
seg_rank = ctx.get_rank(segment_inp)
if ctx.get_dtype(segment_inp) != onnx_pb.TensorProto.INT64:
segment_inp = ctx.make_node("Cast", [segment_inp],
attr={
"to": onnx_pb.TensorProto.INT64
}).output[0]
utils.make_sure(
seg_rank == 1,
"Segment ops only supported for segments of rank 1, not %s",
seg_rank)
data_np_dtype = utils.map_onnx_to_numpy_type(data_dtype)
seg_np_dtype = utils.map_onnx_to_numpy_type(ctx.get_dtype(segment_inp))
if num_segments_specified and ctx.get_dtype(
segment_inp) != ctx.get_dtype(num_segments):
num_segments = ctx.make_node("Cast", [num_segments],
attr={
"to": ctx.get_dtype(segment_inp)
}).output[0]
data_is_float = np.dtype(data_np_dtype).kind == 'f'
data_is_int = np.dtype(data_np_dtype).kind == 'i'
utils.make_sure(data_is_float or data_is_int,
"dtype for Segment ops must be float or int")
if node.type in ["SegmentSum", "SegmentMean", "SegmentSqrtN"]:
onnx_op = "ReduceSum"
identity_value = np.array(0, dtype=data_np_dtype)
elif node.type == "SegmentProd":
onnx_op = "ReduceProd"
identity_value = np.array(1, dtype=data_np_dtype)
elif node.type == "SegmentMax":
onnx_op = "ReduceMax"
if data_is_float:
identity_value = np.array('-inf', dtype=data_np_dtype)
else:
identity_value = np.iinfo(data_np_dtype).min
elif node.type == "SegmentMin":
onnx_op = "ReduceMin"
if data_is_float:
identity_value = np.array('inf', dtype=data_np_dtype)
else:
identity_value = np.iinfo(data_np_dtype).max
if not num_segments_specified:
max_segment = ctx.make_node("ReduceMax", [segment_inp],
attr={
'axes': [0],
'keepdims': 0
})
one_const = ctx.make_const(utils.make_name("const_one"),
np.array(1, dtype=seg_np_dtype))
num_segments = ctx.make_node(
"Add", [max_segment.output[0], one_const.output[0]]).output[0]
num_segments_unsq = GraphBuilder(ctx).make_unsqueeze({
'data': num_segments,
'axes': [0]
})
seg_shape = ctx.make_node("Shape", [segment_inp]).output[0]
seg_shape_sq = GraphBuilder(ctx).make_squeeze({
"data": seg_shape,
"axes": [0]
})
segs_sorted, indices = ctx.make_node("TopK", [segment_inp, seg_shape],
attr={
'axis': 0,
'largest': False,
'sorted': True
},
output_count=2,
op_name_scope=node.name).output
seg_unique_node = ctx.make_node("Unique", [segs_sorted],
attr={
'axis': 0,
'sorted': True
},
output_count=4,
op_name_scope=node.name)
seg_unique_node.output[1] = ""
seg_values, _, inv_indices, seg_cnts_sorted = seg_unique_node.output
max_cnt = ctx.make_node("ReduceMax", [seg_cnts_sorted],
attr={
'axes': [0],
'keepdims': True
}).output[0]
if node.type in ["SegmentMean", "SegmentSqrtN"]:
zero_tensor = helper.make_tensor("value",
onnx_pb.TensorProto.INT64,
dims=[1],
vals=[0])
if num_segs_const is not None:
zeros = ctx.make_const(utils.make_name("zeros"),
np.zeros([num_segs_const],
np.int64)).output[0]
else:
zeros = ctx.make_node("ConstantOfShape", [num_segments_unsq],
attr={
'value': zero_tensor
}).output[0]
seg_cnts = ctx.make_node("ScatterElements",
[zeros, seg_values, seg_cnts_sorted],
attr={
'axis': 0
}).output[0]
seg_cnts_float = ctx.make_node("Cast", [seg_cnts],
attr={
'to': onnx_pb.TensorProto.FLOAT
}).output[0]
if node.type == "SegmentMean":
scaling_amt = seg_cnts_float
elif node.type == "SegmentSqrtN":
scaling_amt = ctx.make_node("Sqrt", [seg_cnts_float]).output[0]
else:
scaling_amt = None
if scaling_amt is not None and num_segments_specified:
# If empty segments are possible, we must avoid division by zero
const_one_float = ctx.make_const(
utils.make_name("const_one_float"),
np.array(1, dtype=np.float32))
scaling_amt = ctx.make_node(
"Max", [scaling_amt, const_one_float.output[0]]).output[0]
zero_const_int64 = ctx.make_const(utils.make_name("const_zero"),
np.array(0,
dtype=np.int64)).output[0]
one_const_int64 = ctx.make_const(utils.make_name("const_one"),
np.array(1, dtype=np.int64)).output[0]
seg_range = ctx.make_node(
"Range",
[zero_const_int64, seg_shape_sq, one_const_int64]).output[0]
id_to_cnt = ctx.make_node("Gather",
[seg_cnts_sorted, inv_indices]).output[0]
range_mod = ctx.make_node("Mod", [seg_range, id_to_cnt]).output[0]
idx_grid_shape = ctx.make_node("Concat", [num_segments_unsq, max_cnt],
{
'axis': 0
}).output[0]
neg_one_tensor = helper.make_tensor("value",
onnx_pb.TensorProto.INT64,
dims=[1],
vals=[-1])
idx_grid = ctx.make_node("ConstantOfShape", [idx_grid_shape], {
'value': neg_one_tensor
}).output[0]
segs_sorted_unsq = GraphBuilder(ctx).make_unsqueeze({
'data': segs_sorted,
'axes': [-1]
})
range_mod_unsq = GraphBuilder(ctx).make_unsqueeze({
'data': range_mod,
'axes': [-1]
})
scatter_indices = ctx.make_node("Concat",
[segs_sorted_unsq, range_mod_unsq],
attr={
'axis': 1
}).output[0]
scatted_grid = ctx.make_node(
"ScatterND", [idx_grid, scatter_indices, indices]).output[0]
data_shape = ctx.make_node("Shape", [data_inp]).output[0]
max_int64 = int(utils.get_max_value(np.int64))
identity_shape = GraphBuilder(ctx).make_slice({
'data': data_shape,
'starts': [1],
'ends': [max_int64],
'axes': [0]
})
id_tensor = helper.make_tensor("value",
ctx.get_dtype(data_inp),
dims=[1],
vals=[identity_value])
identity = ctx.make_node("ConstantOfShape", [identity_shape], {
'value': id_tensor
}).output[0]
id_unsq = GraphBuilder(ctx).make_unsqueeze({
'data': identity,
'axes': [0]
})
data_with_id = ctx.make_node("Concat", [data_inp, id_unsq],
attr={
'axis': 0
}).output[0]
data_grid = ctx.make_node("Gather",
[data_with_id, scatted_grid]).output[0]
if onnx_op == "ReduceSum":
reduction_result = GraphBuilder(ctx).make_reduce_sum(
{
'data': data_grid,
'axes': [1],
"keepdims": False
},
op_name_scope=node.name)
else:
reduction_result = ctx.make_node(onnx_op, [data_grid],
attr={
'axes': [1],
'keepdims': False
},
op_name_scope=node.name).output[0]
if scaling_amt is not None:
if data_rank is None:
# Left pad scale to match data rank
data_slice_rank = ctx.make_node("Shape",
[identity_shape]).output[0]
one_tensor = helper.make_tensor("value",
onnx_pb.TensorProto.INT64,
dims=[1],
vals=[1])
ones_of_shape = ctx.make_node("ConstantOfShape",
[data_slice_rank], {
'value': one_tensor
}).output[0]
zero_unsq = ctx.make_const(utils.make_name('const_zero'),
np.array([0], np.int64)).output[0]
scaling_shape = ctx.make_node("Concat",
[zero_unsq, ones_of_shape],
attr={
'axis': 0
}).output[0]
scaling_amt = ctx.make_node(
"Reshape", [scaling_amt, scaling_shape]).output[0]
elif data_rank != 1:
scaling_amt = GraphBuilder(ctx).make_unsqueeze({
'data':
scaling_amt,
'axes':
list(range(1, data_rank))
})
reduction_result = ctx.make_node(
"Div", [reduction_result, scaling_amt]).output[0]
ctx.replace_all_inputs(node.output[0], reduction_result)
ctx.remove_node(node.name)
def compute_const_folding_using_tf(g, const_node_values, graph_outputs):
"""Find nodes with constant inputs and compute their values using TF"""
if const_node_values is None:
const_node_values = {}
graph_outputs = set(graph_outputs)
from tf2onnx.tf_loader import tf_session, tf_placeholder # pylint: disable=import-outside-toplevel
ops = g.get_operations()
outputs_to_values = {}
outputs_to_dtypes = {}
outputs_to_shapes = {}
shape_node_outputs = {}
def is_small_shape(x):
return np.product(x) <= 1000
def is_huge_shape(x):
return np.product(x) >= 1000000
for node in ops:
# Load values of constants. Use const_node_values if possible
if node.type in ["Const", "ConstV2"]:
tensor = node.node_def.attr["value"].tensor
if node.name in const_node_values:
tensor.tensor_content = const_node_values[node.name]
outputs_to_values[node.outputs[0].name] = get_tf_tensor_data(
tensor)
outputs_to_dtypes[node.outputs[0].name] = node.outputs[0].dtype
for out in node.outputs:
outputs_to_shapes[out.name] = get_tf_tensor_shape(out)
for node in ops:
if node.type == "Shape":
shape = outputs_to_shapes.get(node.inputs[0].name)
if shape is not None:
shape_node_outputs[node.outputs[0].name] = shape
unneeded_outputs = set()
progress = True
while progress:
progress = False
for node in ops:
# Find ops with constant inputs and compute their values
input_names = [i.name for i in node.inputs]
output_names = [i.name for i in node.outputs]
if node.type == 'StridedSlice' and input_names[0] in shape_node_outputs \
and output_names[0] not in outputs_to_values:
shape = shape_node_outputs[input_names[0]]
i = get_index_from_strided_slice_of_shape(
node, outputs_to_values)
if i is not None and 0 <= i < len(
shape) and shape[i] is not None:
np_dtype = map_onnx_to_numpy_type(
map_tf_dtype(node.outputs[0].dtype))
outputs_to_values[output_names[0]] = np.array(
shape[i], dtype=np_dtype)
outputs_to_dtypes[
node.outputs[0].name] = node.outputs[0].dtype
progress = True
can_fold = node.type not in [
'Enter', 'Placeholder', 'PlaceholderWithDefault'
]
can_fold = can_fold and len(input_names) > 0 and all(
inp in outputs_to_values for inp in input_names)
# We can only fold nodes with a single output
can_fold = can_fold and len(
output_names) == 1 and output_names[0] not in outputs_to_values
# Skip if value already computed, used, and discarded
can_fold = can_fold and output_names[
0] not in unneeded_outputs and output_names[
0] not in graph_outputs
if can_fold:
# Make a mini graph containing just the node to fold
g2 = tf.Graph()
with g2.as_default():
for inp in input_names:
tf_placeholder(outputs_to_dtypes[inp],
name=inp.split(':')[0])
mini_graph_def = g2.as_graph_def()
mini_graph_def.node.append(node.node_def)
g3 = tf.Graph()
with g3.as_default():
feed_dict = {}
inp_shapes = []
for inp in input_names:
inp_np = outputs_to_values[inp]
feed_dict[inp] = inp_np
inp_shapes.append(inp_np.shape)
try:
with tf_session() as sess:
tf.import_graph_def(mini_graph_def, name='')
results = sess.run(output_names,
feed_dict=feed_dict)
if is_huge_shape(results[0].shape) and all(
is_small_shape(inp) for inp in inp_shapes):
logger.debug(
"Skipping folding of node %s since result shape %s is much larger "
"than input shapes %s", node.name,
results[0].shape, inp_shapes)
else:
outputs_to_values[output_names[0]] = results[0]
outputs_to_dtypes[
output_names[0]] = node.outputs[0].dtype
progress = True
except Exception: # pylint: disable=broad-except
logger.debug("Could not fold node %s", node.name)
unneeded_outputs.update(outputs_to_values.keys())
for node in ops:
# Mark values we need to keep
input_names = [i.name for i in node.inputs]
output_names = [i.name for i in node.outputs]
if len(output_names) == 1 and output_names[0] in outputs_to_values:
continue
for i in input_names:
if i in unneeded_outputs:
unneeded_outputs.remove(i)
for node in unneeded_outputs:
# Remove unneeded values to prevent memory usage explosion
if node in outputs_to_values:
del outputs_to_values[node]
del outputs_to_dtypes[node]
for node in ops:
# We don't need the constants any more
if node.type in ["Const", "ConstV2"
] and node.outputs[0].name in outputs_to_values:
del outputs_to_values[node.outputs[0].name]
del outputs_to_dtypes[node.outputs[0].name]
logger.info("Computed %d values for constant folding",
len(outputs_to_values))
return outputs_to_values, outputs_to_dtypes
def version_9(cls, ctx, node, **kwargs):
node_inputs = node.input
num_segments_specified = False
if node.type.endswith("WithNumSegments") or node.type.startswith(
"Unsorted"):
num_segments_specified = True
num_segments = node_inputs.pop()
node.type = node.type.replace("WithNumSegments", "")
node.type = node.type.replace("Unsorted", "")
if node.type.startswith("Sparse"):
data_inp, indices_inp, segment_inp = node_inputs
gather_node = ctx.make_node("Gather", [data_inp, indices_inp],
attr={'axis': 0})
data_inp = gather_node.output[0]
node.type = node.type.replace("Sparse", "")
else:
data_inp, segment_inp = node_inputs
# Data has shape [n, a, b, ..., c]
data_shape = ctx.get_shape(data_inp)
data_rank = len(data_shape) if data_shape is not None else None
data_dtype = ctx.get_dtype(data_inp)
data_np_dtype = utils.map_onnx_to_numpy_type(data_dtype)
seg_np_dtype = utils.map_onnx_to_numpy_type(ctx.get_dtype(segment_inp))
if num_segments_specified and ctx.get_dtype(
segment_inp) != ctx.get_dtype(num_segments):
num_segments = ctx.make_node("Cast", [num_segments],
attr={
"to": ctx.get_dtype(segment_inp)
}).output[0]
data_is_float = np.dtype(data_np_dtype).kind == 'f'
data_is_int = np.dtype(data_np_dtype).kind == 'i'
utils.make_sure(data_is_float or data_is_int,
"dtype for Segment ops must be float or int")
if node.type in ["SegmentSum", "SegmentMean", "SegmentSqrtN"]:
onnx_op = "ReduceSum"
identity_value = np.array(0, dtype=data_np_dtype)
elif node.type == "SegmentProd":
onnx_op = "ReduceProd"
identity_value = np.array(1, dtype=data_np_dtype)
elif node.type == "SegmentMax":
onnx_op = "ReduceMax"
if data_is_float:
identity_value = np.array('-inf', dtype=data_np_dtype)
else:
identity_value = np.iinfo(data_np_dtype).min
elif node.type == "SegmentMin":
onnx_op = "ReduceMin"
if data_is_float:
identity_value = np.array('inf', dtype=data_np_dtype)
else:
identity_value = np.iinfo(data_np_dtype).max
if not num_segments_specified:
max_segment = ctx.make_node("ReduceMax", [segment_inp],
attr={
'axes': [0],
'keepdims': 0
})
one_const = ctx.make_const(utils.make_name("const_one"),
np.array(1, dtype=seg_np_dtype))
num_segments = ctx.make_node(
"Add", [max_segment.output[0], one_const.output[0]]).output[0]
# ORT doesn't support bool for OneHot so we use float32 and cast to bool
onehot_values = ctx.make_const(utils.make_name("onehot_values"),
np.array([0, 1], dtype=np.float32))
# one_hot_node has shape [s, n] (s is # segments)
one_hot_node = ctx.make_node(
"OneHot", [segment_inp, num_segments, onehot_values.output[0]],
attr={'axis': 0})
if node.type == "SegmentMean":
scaling_node_output = GraphBuilder(ctx).make_reduce_sum({
"data":
one_hot_node.output[0],
"axes": [1],
"keepdims":
0,
"noop_with_empty_axes":
1
})
elif node.type == "SegmentSqrtN":
seg_cnts_node_output = GraphBuilder(ctx).make_reduce_sum({
"data":
one_hot_node.output[0],
"axes": [1],
"keepdims":
0,
"noop_with_empty_axes":
1
})
scaling_node_output = ctx.make_node(
"Sqrt", [seg_cnts_node_output]).output[0]
else:
scaling_node_output = None
if scaling_node_output is not None and num_segments_specified:
# If empty segments are possible, we must avoid division by zero
const_one_float = ctx.make_const(
utils.make_name("const_one_float"),
np.array(1, dtype=np.float32))
scaling_node_output = ctx.make_node(
"Max",
[scaling_node_output, const_one_float.output[0]]).output[0]
if onnx_op == "ReduceSum":
# If the op is a summation, we can use MatMul instead of Where, which is faster
# Data shape is [n, a, b, ..., c]
data_shape_node = ctx.make_node("Shape", [data_inp])
new_shape = ctx.make_const(utils.make_name("reshape_const"),
np.array([0, -1], dtype=np.int64))
# Reshape the data from [n, a, b, ..., c] to [n, P]
data_reshape = ctx.make_node("Reshape",
[data_inp, new_shape.output[0]])
one_hot_cast = one_hot_node
if data_dtype != onnx_pb.TensorProto.FLOAT:
one_hot_cast = ctx.make_node("Cast", [one_hot_node.output[0]],
attr={'to': data_dtype})
# Shapes [s, n] * [n, P] => [s, P]
product = ctx.make_node(
"MatMul", [one_hot_cast.output[0], data_reshape.output[0]],
op_name_scope=node.name)
if scaling_node_output is not None:
scaling_node_unsqueeze = ctx.make_node("Unsqueeze",
[scaling_node_output],
attr={'axes': [1]})
product = ctx.make_node(
"Div",
[product.output[0], scaling_node_unsqueeze.output[0]])
# Create new shape [0, a, b, ..., c]
max_int64 = int(utils.get_max_value(np.int64))
new_shape_slice = GraphBuilder(ctx).make_slice({
"data":
data_shape_node.output[0],
"ends": [max_int64],
"starts": [1],
"axes": [0]
})
zero_const = ctx.make_const(utils.make_name("zero_const"),
np.array([0], dtype=np.int64))
new_shape = ctx.make_node("Concat",
[zero_const.output[0], new_shape_slice],
attr={'axis': 0})
shapes = node.output_shapes
dtypes = node.output_dtypes
ctx.remove_node(node.name)
# Reshape result from [s, P] to [s, a, b, ..., c]
ctx.make_node("Reshape", [product.output[0], new_shape.output[0]],
name=node.name,
outputs=node.output,
shapes=shapes,
dtypes=dtypes)
return
identity_const = ctx.make_const(utils.make_name("const_identity"),
identity_value)
one_hot_bool = ctx.make_node("Cast", [one_hot_node.output[0]],
attr={"to": onnx_pb.TensorProto.BOOL})
one_hot_unsqueeze = one_hot_bool
# Make one_hot_unsqueeze have shape [s, n, 1, 1, ..., 1]
if data_rank is None:
# Unsqueeze requires known rank, but we can use Reshape if rank is unknown
shape_node = ctx.make_node("Shape", [data_inp])
rank_node = ctx.make_node("Shape", [shape_node.output[0]])
one_const_int64 = ctx.make_const(utils.make_name("const_one"),
np.array([1], dtype=np.int64))
num_unsqueeze_dims = ctx.make_node(
"Sub", [rank_node.output[0], one_const_int64.output[0]])
one_tensor = helper.make_tensor("value",
onnx_pb.TensorProto.INT64,
dims=[1],
vals=[1])
unsqueeze_dims = ctx.make_node(
"ConstantOfShape",
inputs=[num_unsqueeze_dims.output[0]],
attr={"value": one_tensor})
# Zero indicates a dimension should be unchanged
double_zero_const = ctx.make_const(
utils.make_name("double_zero"), np.array([0, 0],
dtype=np.int64))
expanded_shape = ctx.make_node(
"Concat",
[double_zero_const.output[0], unsqueeze_dims.output[0]],
attr={'axis': 0})
one_hot_unsqueeze = ctx.make_node(
"Reshape", [one_hot_bool.output[0], expanded_shape.output[0]])
elif data_rank > 1:
new_dims = list(range(2, 2 + data_rank - 1))
one_hot_unsqueeze = ctx.make_node("Unsqueeze",
[one_hot_bool.output[0]],
attr={'axes': new_dims})
# Shape of data: [n, a, b, ..., c]
# Shape of one_hot: [s, n, 1, 1, ..., 1]
# Broadcast left-pads shape with 1s, so result is shape: [s, n, a, b, ..., c]
where_node = ctx.make_node(
"Where",
[one_hot_unsqueeze.output[0], data_inp, identity_const.output[0]])
shapes = node.output_shapes
dtypes = node.output_dtypes
ctx.remove_node(node.name)
# After reduction over axis 1, shape is: [s, a, b, ..., c]
ctx.make_node(onnx_op, [where_node.output[0]],
attr={
'axes': [1],
'keepdims': 0
},
name=node.name,
outputs=node.output,
shapes=shapes,
dtypes=dtypes)
def version_1(cls, ctx, node, **kwargs):
"""
Inspired from `Python implementation of RFFT
<https://jakevdp.github.io/blog/2013/08/28/understanding-the-fft/>`_.
Complex version:
::
import numpy as np
def _DFT_cst(N, fft_length):
n = np.arange(N)
k = n.reshape((N, 1)).astype(np.float64)
M = np.exp(-2j * np.pi * k * n / N)
return M[:fft_length // 2 + 1]
def DFT(x, fft_length=None):
if len(x.shape) == 1:
x = x.reshape((-1, 1))
else:
x = x.T
if fft_length is None:
fft_length = x.shape[0]
cst = _DFT_cst(x.shape[0], fft_length)
return np.dot(cst, x).T
Real version, first axis is (real, imag) part:
::
import numpy as np
def _DFT_real_cst(N, fft_length):
n = np.arange(N)
k = n.reshape((N, 1)).astype(np.float64)
M = np.exp(-2j * np.pi * k * n / N)
M = M[:fft_length // 2 + 1]
both = np.empty((2,) + M.shape)
both[0, :, :] = np.real(M)
both[1, :, :] = np.imag(M)
return both
def DFT_real(x, fft_length=None):
if len(x.shape) == 1:
x = x.reshape((-1, 1))
else:
x = x.T
if fft_length is None:
fft_length = x.shape[0]
cst = _DFT_real_cst(x.shape[0], fft_length)
res = np.dot(cst, x)
return np.transpose(res, (0, 2, 1))
"""
supported_dtypes = [
onnx_pb.TensorProto.FLOAT,
onnx_pb.TensorProto.FLOAT16,
onnx_pb.TensorProto.DOUBLE,
onnx_pb.TensorProto.COMPLEX64,
onnx_pb.TensorProto.COMPLEX128,
]
consumers = ctx.find_output_consumers(node.output[0])
consumer_types = set(op.type for op in consumers)
utils.make_sure(
consumer_types == {'ComplexAbs'},
"Current implementation of RFFT only allows ComplexAbs as consumer not %r",
consumer_types)
onnx_dtype = ctx.get_dtype(node.input[0])
utils.make_sure(onnx_dtype in supported_dtypes, "Unsupported input type.")
shape = ctx.get_shape(node.input[0])
np_dtype = utils.map_onnx_to_numpy_type(onnx_dtype)
shape_n = shape[-1]
utils.make_sure(len(node.input) == 2, "Two inputs expected not %r", len(node.input))
# This input should be a constant.
fft_length_name = node.input[1]
node_fft_length = ctx.get_node_by_output(fft_length_name, search_in_parent_graphs=True)
utils.make_sure(node_fft_length.type == 'Const',
"fft_length should be a constant, the other case is not implemented yet.")
value = node_fft_length.get_attr("value")
value_array = to_array(value.t)
utils.make_sure(value_array.shape == (1,), "Unexpected shape for fft_length (%r)", value_array.shape)
fft_length = value_array[0]
# TODO: handle this parameter when onnx.helper.make_node is fixed.
# Tcomplex = node.get_attr("Tcomplex")
if np_dtype == np.float16:
res_onnx_dtype = utils.map_numpy_to_onnx_dtype(np.float16)
np_dtype = np.float16
elif np_dtype in (np.float32, np.complex64):
res_onnx_dtype = utils.map_numpy_to_onnx_dtype(np.float32)
np_dtype = np.float32
else:
res_onnx_dtype = utils.map_numpy_to_onnx_dtype(np.float64)
np_dtype = np.float64
real_imag_part = make_dft_constant(shape_n, np_dtype, fft_length)
onx_real_imag_part = ctx.make_const(
name=utils.make_name('cst_rfft_%d' % shape_n), np_val=real_imag_part)
shapei = list(np.arange(len(shape)))
perm = shapei[:-2] + [shapei[-1], shapei[-2]]
trx = ctx.make_node(
"Transpose", inputs=[node.input[0]], attr=dict(perm=perm),
name=utils.make_name(node.name + 'tr'))
ctx.remove_node(node.name)
mult = ctx.make_node(
"MatMul", inputs=[onx_real_imag_part.name, trx.output[0]],
name=utils.make_name('CPLX_' + node.name + 'rfft'))
new_shape = [2] + list(shape)
shapei = list(np.arange(len(new_shape)))
perm = shapei[:-2] + [shapei[-1], shapei[-2]]
last_node = ctx.make_node(
"Transpose", inputs=[mult.output[0]], attr=dict(perm=perm),
name=utils.make_name('CPLX_' + node.name + 'rfft'),
shapes=[new_shape], dtypes=[res_onnx_dtype])
ctx.replace_all_inputs(node.output[0], last_node.output[0]) # ops=ctx.get_nodes()
def version_1(cls, ctx, node, **kwargs):
"""
Args:
x: A `Tensor`. Must be one of the following types: `float32`.
The input to the LSTM cell, shape (batch_size, num_inputs).
cs_prev: A `Tensor`. Must have the same type as `x`.
Value of the cell state at previous time step.
h_prev: A `Tensor`. Must have the same type as `x`.
Output of the previous cell at previous time step.
w: A `Tensor`. Must have the same type as `x`. The weight matrix.
wci: A `Tensor`. Must have the same type as `x`.
The weight matrix for input gate peephole connection.
wcf: A `Tensor`. Must have the same type as `x`.
The weight matrix for forget gate peephole connection.
wco: A `Tensor`. Must have the same type as `x`.
The weight matrix for output gate peephole connection.
b: A `Tensor`. Must have the same type as `x`. The bias vector.
forget_bias: An optional `float`. Defaults to `1`. The forget gate bias.
cell_clip: An optional `float`. Defaults to `-1` (no clipping).
Value to clip the 'cs' value to. Disable by setting to negative value.
use_peephole: An optional `bool`. Defaults to `False`.
Whether to use peephole weights.
name: A name for the operation (optional).
Returns:
A tuple of `Tensor` objects (i, cs, f, o, ci, co, h).
i: A `Tensor`. Has the same type as `x`. The input gate.
cs: A `Tensor`. Has the same type as `x`. The cell state before the tanh.
f: A `Tensor`. Has the same type as `x`. The forget gate.
o: A `Tensor`. Has the same type as `x`. The output gate.
ci: A `Tensor`. Has the same type as `x`. The cell input.
co: A `Tensor`. Has the same type as `x`. The cell after the tanh.
h: A `Tensor`. Has the same type as `x`. The output h vector.
```python
xh = [x, h_prev]
[i, ci, f, o] = xh * w + b
f = f + forget_bias
if not use_peephole:
wci = wcf = wco = 0
i = sigmoid(cs_prev .* wci + i)
f = sigmoid(cs_prev .* wcf + f)
ci = tanh(ci)
cs = ci .* i + cs_prev .* f
cs = clip(cs, cell_clip)
o = sigmoid(cs * wco + o)
co = tanh(cs)
h = co .* o
```
"""
nodes = []
x, cs_prev, h_prev, w, wci, wcf, wco, b = node.input
forget_bias = float(node.get_attr("forget_bias").f)
cell_clip = float(node.get_attr("cell_clip").f)
use_peephole = bool(node.get_attr("use_peephole").i)
def make_sigmoid(i, w, b):
i_w_node = ctx.make_node("Mul", [i, w])
i_w_b_node = ctx.make_node("Add", [i_w_node.output[0], b])
output_node = ctx.make_node("Sigmoid", [i_w_b_node.output[0]])
nodes.extend([i_w_node, i_w_b_node, output_node])
return output_node.output[0]
# xh = [x, h]
xh_node = ctx.make_node("Concat", [x, h_prev], attr={"axis": 1})
# i, ci, f, o = xh * w + b
xh_w_node = ctx.make_node("MatMul", [xh_node.output[0], w])
w_shape = ctx.get_shape(w)
if len(w_shape) != 2 or w_shape[1] % 4 != 0:
raise RuntimeError(
"shape of W of LSTMBlockCell {} should be times of 4".format(
node.name))
merged_output_node = ctx.make_node("Add", [xh_w_node.output[0], b])
w_last_dim = int(w_shape[1] / 4)
split = [w_last_dim] * 4
split_output_node = ctx.make_node("Split",
[merged_output_node.output[0]],
attr={
"axis": 1,
"split": split
},
output_count=4)
i, ci, f, o = split_output_node.output
# f = f + forget_bias
forget_bias_const = ctx.make_const(
utils.make_name("{}__forget_bias".format(node.name)),
np.array(forget_bias, dtype=np.float32))
f_node = ctx.make_node("Add", [f, forget_bias_const.output[0]])
if not use_peephole:
zeros_const = ctx.make_const(
utils.make_name("{}__zeros_const".format(node.name)),
np.zeros([w_last_dim], dtype=np.float32))
nodes.append(zeros_const)
wci = zeros_const.output[0]
wcf = zeros_const.output[0]
wco = zeros_const.output[0]
# i = sigmoid(cs_prev .* wci + i)
i = make_sigmoid(cs_prev, wci, i)
# f = sigmoid(cs_prev .* wcf + f)
f = make_sigmoid(cs_prev, wcf, f_node.output[0])
# ci = Tanh(ci)
ci_node = ctx.make_node("Tanh", [ci])
# cs = ci .* i + f .* cs_prev
ci_i_node = ctx.make_node("Mul", [ci_node.output[0], i])
cs_prev_f_node = ctx.make_node("Mul", [cs_prev, f])
cs_node = ctx.make_node(
"Add", [ci_i_node.output[0], cs_prev_f_node.output[0]])
cs = cs_node.output[0]
# cs = clip(cs)
if cell_clip > 0:
if ctx.opset < 11:
cs_clip_node = ctx.make_node("Clip", [cs],
attr={
"max": cell_clip,
"min": -cell_clip
})
nodes.append(cs_clip_node)
cs = cs_clip_node.output[0]
else:
dtype = utils.map_onnx_to_numpy_type(ctx.get_dtype(cs))
name_min = utils.make_name("{}_min".format(node.name))
name_max = utils.make_name("{}_max".format(node.name))
min_const = ctx.make_const(name_min,
np.array(-cell_clip, dtype=dtype))
max_const = ctx.make_const(name_max,
np.array(cell_clip, dtype=dtype))
cs_clip_node = ctx.make_node(
'Clip', [cs, min_const.output[0], max_const.output[0]])
nodes.append(cs_clip_node)
cs = cs_clip_node.output[0]
# o = cs * wco + o
o = make_sigmoid(cs, wco, o)
# co = Tanh(cs)
co_node = ctx.make_node("Tanh", [cs])
# h = co .* o
h_node = ctx.make_node("Mul", [co_node.output[0], o])
def replace_output(old_output, new_output):
ctx.replace_all_inputs(old_output,
new_output) # ops=ctx.get_nodes()
ctx.copy_dtype(old_output, new_output)
ctx.copy_shape(old_output, new_output)
replace_output(node.output[0], i)
replace_output(node.output[1], cs)
replace_output(node.output[2], f)
replace_output(node.output[3], o)
replace_output(node.output[4], ci_node.output[0])
replace_output(node.output[5], co_node.output[0])
replace_output(node.output[6], h_node.output[0])
def version_9(cls, ctx, node, **kwargs):
data_inp = node.input[0]
segment_inp = node.input[1]
data_shape = ctx.get_shape(data_inp)
data_rank = len(data_shape) if data_shape is not None else None
data_np_dtype = utils.map_onnx_to_numpy_type(ctx.get_dtype(data_inp))
seg_np_dtype = utils.map_onnx_to_numpy_type(ctx.get_dtype(segment_inp))
data_is_float = np.dtype(data_np_dtype).kind == 'f'
data_is_int = np.dtype(data_np_dtype).kind == 'i'
utils.make_sure(data_is_float or data_is_int,
"dtype for Segment ops must be float or int")
if node.type == "SegmentSum":
onnx_op = "ReduceSum"
identity_value = np.array(0, dtype=data_np_dtype)
elif node.type == "SegmentProd":
onnx_op = "ReduceProd"
identity_value = np.array(1, dtype=data_np_dtype)
elif node.type == "SegmentMax":
onnx_op = "ReduceMax"
if data_is_float:
identity_value = np.array('-inf', dtype=data_np_dtype)
else:
identity_value = np.iinfo(data_np_dtype).min
elif node.type == "SegmentMin":
onnx_op = "ReduceMin"
if data_is_float:
identity_value = np.array('inf', dtype=data_np_dtype)
else:
identity_value = np.iinfo(data_np_dtype).max
max_segment = ctx.make_node("ReduceMax", [segment_inp],
attr={
'axes': [0],
'keepdims': 0
})
one_const = ctx.make_const(utils.make_name("const_one"),
np.array(1, dtype=seg_np_dtype))
identity_const = ctx.make_const(utils.make_name("const_identity"),
identity_value)
num_segments = ctx.make_node(
"Add", [max_segment.output[0], one_const.output[0]])
# ORT doesn't support bool for OneHot so we use float32 and cast to bool
onehot_values = ctx.make_const(utils.make_name("onehot_values"),
np.array([0, 1], dtype=np.float32))
one_hot_node = ctx.make_node(
"OneHot",
[segment_inp, num_segments.output[0], onehot_values.output[0]],
attr={'axis': 0})
one_hot_bool = ctx.make_node("Cast", [one_hot_node.output[0]],
attr={"to": onnx_pb.TensorProto.BOOL})
one_hot_unsqueeze = one_hot_bool
if data_rank is None:
# Unsqueeze requires known rank, but we can use Reshape if rank is unknown
shape_node = ctx.make_node("Shape", [data_inp])
rank_node = ctx.make_node("Shape", [shape_node.output[0]])
one_const_int64 = ctx.make_const(utils.make_name("const_one"),
np.array([1], dtype=np.int64))
num_unsqueeze_dims = ctx.make_node(
"Sub", [rank_node.output[0], one_const_int64.output[0]])
one_tensor = helper.make_tensor("value",
onnx_pb.TensorProto.INT64,
dims=[1],
vals=[1])
unsqueeze_dims = ctx.make_node(
"ConstantOfShape",
inputs=[num_unsqueeze_dims.output[0]],
attr={"value": one_tensor})
double_zero_const = ctx.make_const(
utils.make_name("double_zero"), np.array([0, 0],
dtype=np.int64))
expanded_shape = ctx.make_node(
"Concat",
[double_zero_const.output[0], unsqueeze_dims.output[0]],
attr={'axis': 0})
one_hot_unsqueeze = ctx.make_node(
"Reshape", [one_hot_bool.output[0], expanded_shape.output[0]])
elif data_rank > 1:
new_dims = list(range(2, 2 + data_rank - 1))
one_hot_unsqueeze = ctx.make_node("Unsqueeze",
[one_hot_bool.output[0]],
attr={'axes': new_dims})
mul_node = ctx.make_node(
"Where",
[one_hot_unsqueeze.output[0], data_inp, identity_const.output[0]])
shapes = node.output_shapes
dtypes = node.output_dtypes
ctx.remove_node(node.name)
ctx.make_node(onnx_op, [mul_node.output[0]],
attr={
'axes': [1],
'keepdims': 0
},
name=node.name,
outputs=node.output,
shapes=shapes,
dtypes=dtypes)
def version_9(cls, ctx, node, **kwargs):
"""
Obtained with a linear regression.
::
def atan2(y, x):
sx = numpy.sign(x)
sy = numpy.sign(y)
pi_part = (sy + sx * (sy ** 2 - 1)) * (sx - 1) * (-numpy.pi/2)
atan_part = numpy.arctan(y / (x + (1 - sx ** 2))) * sx ** 2
return atan_part + pi_part
"""
onnx_dtype = ctx.get_dtype(node.input[0])
shape = ctx.get_shape(node.input[0])
np_dtype = utils.map_onnx_to_numpy_type(onnx_dtype)
# sign part
sign_x_node = ctx.make_node(
"Sign", inputs=node.input[1:],
name=utils.make_name(node.name + 'signx'))
sign_y_node = ctx.make_node(
"Sign", inputs=node.input[:1],
name=utils.make_name(node.name + 'signy'))
sx_node = ctx.make_node(
"Cast", sign_x_node.output[:1], attr={"to": onnx_dtype},
name=utils.make_name(node.name + 'csignx'))
sy_node = ctx.make_node(
"Cast", sign_y_node.output[:1], attr={"to": onnx_dtype},
name=utils.make_name(node.name + 'csigny'))
# cst
one_node = ctx.make_const(
utils.make_name("{}_one".format(node.name)),
np.array([1], dtype=np_dtype))
pib2_node = ctx.make_const(
utils.make_name("{}_pi".format(node.name)),
np.array(- np.pi / 2, dtype=np_dtype))
# pi_part = (sy + sx * (sy ** 2 - 1)) * (sx - 1) * (-numpy.pi/2)
sxm1_node = ctx.make_node(
"Sub", [sx_node.output[0], one_node.output[0]],
name=utils.make_name(node.name + 'sxm1'))
sy2_node = ctx.make_node(
"Mul", [sy_node.output[0], sy_node.output[0]],
name=utils.make_name(node.name + 'sy2'))
sy2m1_node = ctx.make_node(
"Sub", [sy2_node.output[0], one_node.output[0]],
name=utils.make_name(node.name + 'sy2m1'))
sxsy2m1_node = ctx.make_node(
"Mul", [sx_node.output[0], sy2m1_node.output[0]],
name=utils.make_name(node.name + 'sxsy2m1'))
sysxsy2m1_node = ctx.make_node(
"Add", [sy_node.output[0], sxsy2m1_node.output[0]],
name=utils.make_name(node.name + 'sysxsy2m1'))
m1_node = ctx.make_node(
"Mul", [sysxsy2m1_node.output[0], sxm1_node.output[0]],
name=utils.make_name(node.name + 'm1'))
pi_part = ctx.make_node(
"Mul", [m1_node.output[0], pib2_node.output[0]],
name=utils.make_name(node.name + 'pip'))
# atan
sx2_node = ctx.make_node(
"Mul", [sx_node.output[0], sx_node.output[0]],
name=utils.make_name(node.name + 'sx2'))
sx2m1_node = ctx.make_node(
"Sub", [sx2_node.output[0], one_node.output[0]],
name=utils.make_name(node.name + 'sx2m1'))
xsx2m1_node = ctx.make_node(
"Add", [node.input[1], sx2m1_node.output[0]],
name=utils.make_name(node.name + 'xsx2m1'))
div_node = ctx.make_node(
"Div", inputs=[node.input[0], xsx2m1_node.output[0]],
name=utils.make_name(node.name + 'div'))
atan0_node = ctx.make_node(
"Atan", inputs=[div_node.output[0]],
name=utils.make_name(node.name + 'atan0'))
atan_node = ctx.make_node(
"Mul", inputs=[sx2_node.output[0], atan0_node.output[0]],
name=utils.make_name(node.name + 'atan'))
# final
ctx.remove_node(node.name)
last_node = ctx.make_node(
"Add", inputs=[atan_node.output[0], pi_part.output[0]],
op_name_scope=node.name + 'all',
shapes=[shape], dtypes=[onnx_dtype])
ctx.replace_all_inputs(ctx.get_nodes(), node.output[0], last_node.output[0])
def wire_while_body(parent_g, g, loop_node, body_input_to_state_var,
cond_input_to_state_var, output_shapes, output_dtypes,
scope, parent, cond_graph, tf_while_inputs,
scan_output_names, ragged_scan_output_names):
"""Wire subgraph graph into main."""
remove_parents = []
to_remove = []
# tensorflow function inputs that are state_vars come from outer context and
# we need to remove them from the inputs by making the placeholder an identity
for n in g.inputs:
if n.output[0] in body_input_to_state_var:
n.type = "Identity"
g.replace_inputs(n, [body_input_to_state_var[n.output[0]]])
# onnx will pass in cond as argument
cond_node = g.make_node("Placeholder", [],
name=utils.make_name("cond"),
output_count=1,
dtypes=[onnx_pb.TensorProto.BOOL],
shapes=[[]])
# in onnx the body inputs are: index, cond, [loop_vars]
func_inputs = [
i for i in g.input_names[2:] if i not in body_input_to_state_var
]
func_inputs = [g.input_names[0], cond_node.output[0]] + func_inputs
g.set_dtype(func_inputs[0], onnx_pb.TensorProto.INT64)
g.inputs = [g.get_node_by_output(inp) for inp in func_inputs]
for p, c in zip(loop_node.input, func_inputs):
g.copy_shape(p, c)
for i, node in enumerate(g.inputs):
if node.output[0] not in func_inputs:
remove_parents.append(node.output[0])
# this is a tensor array write - make it an identity
scan_outputs = []
ragged_scan_outputs_cnt = 0
names_to_scan_outputs = {}
for node in g.get_nodes():
if node.type == "TensorListSetItem":
if node.inputs[2].type == "RaggedTensorToVariant":
node.type = "SequenceInsert"
row_content = node.inputs[2].input[0]
g.replace_inputs(node, [node.input[0], row_content])
g.set_shape(node.output[0], g.get_shape(node.input[1]))
g.set_dtype(node.output[0],
utils.SeqType(g.get_dtype(node.input[1])))
dense_shape = g.make_node("Shape", [row_content]).output[0]
zero_const = g.make_const(utils.make_name("zero_const"),
np.array(0, np.int64)).output[0]
row_length = g.make_node("Gather",
[dense_shape, zero_const]).output[0]
row_length_id = g.make_node("Identity", [row_length])
scan_outputs.append(row_length_id.output[0])
names_to_scan_outputs[ragged_scan_output_names[
ragged_scan_outputs_cnt]] = row_length_id.output[0]
ragged_scan_outputs_cnt += 1
continue
remove_parents.append(node.input[0])
node.type = "Identity"
g.set_shape(node.output[0], g.get_shape(node.input[2]))
g.set_dtype(node.output[0], g.get_dtype(node.input[2]))
g.replace_inputs(node, [node.input[2]])
scan_outputs.append(node.output[0])
if len(scan_outputs) != len(scan_output_names):
raise ValueError(
"While loop couldn't find scan output index for nodes")
for output in scan_outputs:
if output in names_to_scan_outputs.values():
continue
last_output = output
consumers = g.find_output_consumers(last_output)
while consumers:
node = consumers[0]
if node.type != "Identity":
raise ValueError(
"While loop couldn't find scan output index for node " +
node.name)
last_output = node.output[0]
consumers = g.find_output_consumers(last_output)
if last_output not in scan_output_names:
raise ValueError(
"While loop couldn't find scan output index for node " +
node.name)
names_to_scan_outputs[last_output] = output
# Reorder scan outputs
scan_outputs = [names_to_scan_outputs[name] for name in scan_output_names]
# Use shapes from subgraph if loop node shapes for scan outputs are missing
for i in range(-len(scan_output_names), 0):
if loop_node.output_shapes[i] is None:
shape = g.get_shape(scan_outputs[i])
if shape is not None:
parent_g.set_shape(loop_node.output[i], [-1] + shape)
# remove all nodes feeding to TensorListSetItem's reserved tensor
while remove_parents:
output_name = remove_parents[0]
del remove_parents[0]
node = g.get_node_by_output(output_name)
if node:
if output_name not in func_inputs:
if node.input:
remove_parents.extend(node.input)
g.remove_node(node.name)
for node in to_remove:
g.remove_node(node.name)
cond_binding = parameter_binding(cond_graph,
func_inputs[:1] + g.outputs[2:],
cond_input_to_state_var)
cond_outputs = inline_subgraph(g, cond_graph, "cond__", cond_binding)
g.outputs = [cond_outputs[0]] + g.outputs[2:] + scan_outputs
# onnx does not have a variant type so we try to fish for the dtype in a prior TensorListSetItem.
for o in g.outputs:
if g.get_dtype(o) == onnx_pb.TensorProto.UNDEFINED:
curr_o = o
while g.get_node_by_output(curr_o).type == "Identity":
curr_o = g.get_node_by_output(curr_o).input[0]
g.copy_dtype(curr_o, o)
for node in g.ragged_variant_list_reads:
# Requires opset 11
gather = node.inputs[0]
inp = gather.inputs[0]
while inp.type == "Identity":
inp = inp.inputs[0]
err_msg1 = "Could not find corresponding RaggedTensorToVariant for node %s" % node.name
err_msg2 = "Input to RaggedTensorToVariant for loop has batched_input=False for node %s" % inp.name
err_msg3 = "RAGGED_RANK != 1 for RaggedTensorToVariant node %s" % node.name
utils.make_sure(inp.type == "RaggedTensorToVariant", err_msg1)
utils.make_sure(inp.get_attr_value("batched_input"), err_msg2)
utils.make_sure(inp.get_attr_value("RAGGED_RANK") == 1, err_msg3)
idx = gather.input[1]
idx_unsq = GraphBuilder(g).make_unsqueeze({'data': idx, 'axes': [0]})
np_dtype = utils.map_onnx_to_numpy_type(g.get_dtype(idx_unsq))
const_one = g.make_const(utils.make_name("const_1"),
np.array(1, np_dtype)).output[0]
idx_plus_1 = g.make_node("Add", [idx_unsq, const_one]).output[0]
splits, values = inp.input
start = g.make_node("Gather", [splits, idx_unsq]).output[0]
end = g.make_node("Gather", [splits, idx_plus_1]).output[0]
np_dtype2 = utils.map_onnx_to_numpy_type(g.get_dtype(splits))
axes = g.make_const(utils.make_name("const_zero"),
np.array([0], np_dtype2)).output[0]
sliced_vals = g.make_node("Slice",
[values, start, end, axes]).output[0]
g.replace_all_inputs(node.output[0], sliced_vals)
return g
def any_version(cls, const_length, opset, ctx, node, **kwargs):
"""
Inspired from `Python implementation of RFFT
<https://jakevdp.github.io/blog/2013/08/28/understanding-the-fft/>`_.
Complex version:
::
import numpy as np
def _DFT_cst(N, fft_length):
n = np.arange(N)
k = n.reshape((N, 1)).astype(np.float64)
M = np.exp(-2j * np.pi * k * n / N)
return M[:fft_length // 2 + 1]
def DFT(x, fft_length=None):
if len(x.shape) == 1:
x = x.reshape((-1, 1))
else:
x = x.T
if fft_length is None:
fft_length = x.shape[0]
cst = _DFT_cst(x.shape[0], fft_length)
return np.dot(cst, x).T
Real version, first axis is (real, imag) part:
::
import numpy as np
def _DFT_real_cst(N, fft_length):
n = np.arange(N)
k = n.reshape((N, 1)).astype(np.float64)
M = np.exp(-2j * np.pi * k * n / N)
M = M[:fft_length // 2 + 1]
both = np.empty((2,) + M.shape)
both[0, :, :] = np.real(M)
both[1, :, :] = np.imag(M)
return both
def DFT_real(x, fft_length=None):
if len(x.shape) == 1:
x = x.reshape((-1, 1))
else:
x = x.T
if fft_length is None:
fft_length = x.shape[0]
cst = _DFT_real_cst(x.shape[0], fft_length)
res = np.dot(cst, x)
return np.transpose(res, (0, 2, 1))
"""
supported_dtypes = [
onnx_pb.TensorProto.FLOAT,
onnx_pb.TensorProto.FLOAT16,
onnx_pb.TensorProto.DOUBLE,
onnx_pb.TensorProto.COMPLEX64,
onnx_pb.TensorProto.COMPLEX128,
]
consumers = ctx.find_output_consumers(node.output[0])
consumer_types = set(op.type for op in consumers)
utils.make_sure(
consumer_types == {'ComplexAbs'},
"Current implementation of RFFT or FFT only allows ComplexAbs as consumer not %r",
consumer_types)
input_name = node.input[0]
onnx_dtype = ctx.get_dtype(input_name)
utils.make_sure(onnx_dtype in supported_dtypes, "Unsupported input type.")
shape = ctx.get_shape(node.input[0])
shape_n = shape[-1]
if onnx_dtype in (onnx_pb.TensorProto.COMPLEX64, onnx_pb.TensorProto.COMPLEX128):
parent = ctx.get_node_by_output_in_current_graph(node.input[0])
utils.make_sure(
parent.type == 'Cast' and parent.get_attr_value('to') == onnx_dtype,
"Current implementation of FFT or RFFT assumes the input is real or complex produced "
"by a node Cast just before this one.")
input_name = parent.input[0]
onnx_dtype = ctx.get_dtype(input_name)
np_dtype = utils.map_onnx_to_numpy_type(onnx_dtype)
if np_dtype == np.float16:
res_onnx_dtype = utils.map_numpy_to_onnx_dtype(np.float16)
np_dtype = np.float16
elif np_dtype in (np.float32, np.complex64):
res_onnx_dtype = utils.map_numpy_to_onnx_dtype(np.float32)
np_dtype = np.float32
else:
res_onnx_dtype = utils.map_numpy_to_onnx_dtype(np.float64)
np_dtype = np.float64
if const_length:
# RFFT: length of FFT is known, some computation
# (see function make_dft_constant)
# can be done at conversion time and stored as constant
utils.make_sure(len(node.input) == 2, "Two inputs expected not %r", len(node.input))
# This input should be a constant.
fft_length_name = node.input[1]
node_fft_length = ctx.get_node_by_output(fft_length_name, search_in_parent_graphs=True)
utils.make_sure(node_fft_length.type == 'Const',
"fft_length should be a constant, the other case is not implemented yet.")
value = node_fft_length.get_attr("value")
value_array = to_array(value.t)
utils.make_sure(value_array.shape == (1,), "Unexpected shape for fft_length (%r)", value_array.shape)
fft_length = value_array[0]
# TODO: handle this parameter when onnx.helper.make_node is fixed.
# Tcomplex = node.get_attr("Tcomplex")
real_imag_part = make_dft_constant(shape_n, np_dtype, fft_length)
onx_real_imag_part = ctx.make_const(
name=utils.make_name('cst_rfft_%d' % shape_n), np_val=real_imag_part)
onx_real_imag_part_name = onx_real_imag_part.name
else:
# FFT: length of FFT is unknown, the matrix
# created by function make_dft_constant must be
# done in ONNX.
dyn_shape_all = ctx.make_node("Shape", inputs=[input_name],
name=utils.make_name('CPLX_' + node.name + 'shape'))
m1_cst = ctx.make_const(name=utils.make_name('CPLX_m1'), np_val=np.array([-1], dtype=np.int64))
dyn_shape = ctx.make_node('Gather', inputs=[dyn_shape_all.output[0], m1_cst.name])
one_tensor = helper.make_tensor("value", res_onnx_dtype, dims=[1], vals=[1])
cst_1 = ctx.make_node("ConstantOfShape", inputs=[dyn_shape.output[0]], attr={"value": one_tensor})
just_0 = ctx.make_const(name=utils.make_name('CPLX1'), np_val=np.array([0], dtype=np.int64))
rng1 = ctx.make_node("CumSum", inputs=[cst_1.output[0], just_0.name],
name=utils.make_name('CPLX_' + node.name + 'range'))
p1_cst = ctx.make_const(name=utils.make_name('CPLX_p1'), np_val=np.array([1], dtype=np_dtype))
rng = ctx.make_node("Sub", inputs=[rng1.output[0], p1_cst.name],
name=utils.make_name('CPLX_' + node.name + 'range'))
resh_cst = ctx.make_const(name=utils.make_name('CPLX_reshape'), np_val=np.array([1, -1], dtype=np.int64))
rng_tr1 = ctx.make_node("Reshape", inputs=[rng.output[0], resh_cst.name],
name=utils.make_name('CPLX_' + node.name + 'range'))
resh_cst = ctx.make_const(name=utils.make_name('CPLX_reshape'), np_val=np.array([-1, 1], dtype=np.int64))
rng_tr2 = ctx.make_node("Reshape", inputs=[rng.output[0], resh_cst.name],
name=utils.make_name('CPLX_' + node.name + 'range'))
rng_mat = ctx.make_node('MatMul', inputs=[rng_tr2.output[0], rng_tr1.output[0]],
name=utils.make_name('CPLX_' + node.name + 'range2'))
pi_cst = ctx.make_const(name=utils.make_name('CPLX_pi'), np_val=np.array([np.pi * 2], dtype=np_dtype))
angle_pi = ctx.make_node("Mul", inputs=[rng_mat.output[0], pi_cst.name],
name=utils.make_name('CPLX_' + node.name + 'angle_pi'))
shape_cast = ctx.make_node('Cast', inputs=[dyn_shape.output[0]], attr={'to': res_onnx_dtype})
angle_pibn = ctx.make_node("Div", inputs=[angle_pi.output[0], shape_cast.output[0]],
name=utils.make_name('CPLX_' + node.name + 'angle'))
if opset >= 13:
angle = ctx.make_node("Unsqueeze", inputs=[angle_pibn.output[0], just_0.name],
name=utils.make_name('CPLX_' + node.name + 'angles'))
else:
angle = ctx.make_node("Unsqueeze", inputs=[angle_pibn.output[0]],
name=utils.make_name('CPLX_' + node.name + 'angles'),
attr={'axes': [0]})
rng_cos = ctx.make_node("Cos", inputs=[angle.output[0]],
name=utils.make_name('CPLX_' + node.name + 'cos'))
rng_sin = ctx.make_node("Sin", inputs=[angle.output[0]],
name=utils.make_name('CPLX_' + node.name + 'sin'))
onx_real_imag_part = ctx.make_node("Concat", inputs=[rng_cos.output[0], rng_sin.output[0]],
name=utils.make_name('CPLX_' + node.name + '_cst_fft'),
attr={'axis': 0})
onx_real_imag_part_name = onx_real_imag_part.output[0]
shapei = list(np.arange(len(shape)))
perm = shapei[:-2] + [shapei[-1], shapei[-2]]
trx = ctx.make_node(
"Transpose", inputs=[input_name], attr=dict(perm=perm),
name=utils.make_name(node.name + 'tr'))
ctx.remove_node(node.name)
mult = ctx.make_node(
"MatMul", inputs=[onx_real_imag_part_name, trx.output[0]],
name=utils.make_name('CPLX_' + node.name + 'rfft'))
new_shape = [2] + list(shape)
shapei = list(np.arange(len(new_shape)))
perm = shapei[:-2] + [shapei[-1], shapei[-2]]
last_node = ctx.make_node(
"Transpose", inputs=[mult.output[0]], attr=dict(perm=perm),
name=utils.make_name('CPLX_' + node.name + 'rfft'),
shapes=[new_shape], dtypes=[res_onnx_dtype])
ctx.replace_all_inputs(node.output[0], last_node.output[0]) # ops=ctx.get_nodes()
def read_tfjs_graph(nodes,
weights,
func=None,
graph_inputs=None,
graph_outputs=None,
shape_override=None,
ignore_default=None,
use_default=None):
"""Creates an onnx graph from the provided tfjs nodes"""
if shape_override is None:
shape_override = {}
onnx_nodes = []
output_shapes = {}
tf_dtypes = {}
op_info = {}
graph_name = 'tfjs_model'
func_name = None
def update_shapes(new_shapes):
if isinstance(new_shapes, dict):
new_shapes = new_shapes.items()
for k, v in new_shapes:
output_shapes[k] = shape_override.get(k, v)
if func is not None:
tf_dtypes, fn_input_shapes, graph_inputs, graph_outputs, func_name = read_tfjs_function(
func)
update_shapes(fn_input_shapes)
graph_name = func_name
for inp in graph_inputs:
onnx_nodes.append(
helper.make_node("Placeholder", [], outputs=[inp], name=inp))
if graph_inputs is None:
placeholder_ops = [
"Placeholder", "PlaceholderWithDefault", "PlaceholderV2"
]
graph_inputs = [
n['name'] + ':0' for n in nodes if n['op'] in placeholder_ops
]
for node in nodes:
if node['op'] == "NextIteration":
# NextIteration nodes can violate the topological sort with cyclic dependencies, so we do them first.
node_name = node['name']
output_name = node_name + ':0'
output_shapes[output_name] = None
tf_dtypes[output_name] = read_tfjs_attr(node['attr']['T'],
tf_dtypes=True)
op_info[node_name] = (node['op'], {
'dtype': tf_dtypes[output_name]
}, [tf_dtypes[output_name]])
for node in nodes:
op_type = node['op']
node_name = node['name']
if op_type == "Const":
np_arr = weights[node_name]
out_name = node_name + ':0'
tf_dtype = read_tfjs_attr(node['attr']['dtype'], tf_dtypes=True)
onnx_dtype = tf_utils.map_tf_dtype(tf_dtype)
# The dtype of a Const in tfjs can differ from that of the weight used to get its value
np_dtype = utils.map_onnx_to_numpy_type(onnx_dtype)
onnx_tensor = numpy_helper.from_array(np_arr.astype(np_dtype),
out_name)
onnx_node = helper.make_node("Const", [],
outputs=[out_name],
name=node_name,
value=onnx_tensor)
onnx_nodes.append(onnx_node)
output_shapes[out_name] = shape_override.get(
out_name, list(np_arr.shape))
tf_dtypes[out_name] = tf_dtype
op_info[node_name] = (op_type, {'dtype': tf_dtypes[out_name]}, [])
continue
tf_attr = {}
onnx_attr = {}
fix_string_attr(node)
node_def = tfjs_node_to_tf_node_def(node)
for k, v in node.get('attr', {}).items():
tf_attr[k] = read_tfjs_attr(v, tf_dtypes=True)
if k in tf_utils.TF_IGNORED_NODE_ATTRS:
continue
if k == 'DstT':
k = 'to'
onnx_attr[k] = read_tfjs_attr(v)
if op_type == "FusedDepthwiseConv2dNative":
# This op isn't in tensorflow but can be converted to a TF op
op_type = "_FusedDepthwiseConv2dNative"
err_msg = "explicit_paddings for supported for _FusedDepthwiseConv2dNative"
utils.make_sure(len(tf_attr['explicit_paddings']) == 0, err_msg)
del tf_attr['explicit_paddings']
del onnx_attr['explicit_paddings']
del node_def.attr['explicit_paddings']
node_def.op = op_type
input_names = [
inp for inp in node.get('input', []) if not inp.startswith('^')
]
input_names = [
resolve_output(inp, op_info, func_name) for inp in input_names
]
inp_dtypes = [tf_dtypes[inp] for inp in input_names]
inp_shapes = [output_shapes[inp] for inp in input_names]
inp_consts = [weights.get(inp.split(':')[0]) for inp in input_names]
out_dtypes = get_output_dtypes(op_type, tf_attr, inp_dtypes)
out_shapes = get_output_shapes(node_def, inp_dtypes, inp_shapes,
inp_consts)
op_info[node_name] = (op_type, tf_attr, inp_dtypes)
output_names = [
node_name + ":" + str(i) for i in range(len(out_dtypes))
]
tf_dtypes.update(zip(output_names, out_dtypes))
update_shapes(zip(output_names, out_shapes))
if op_type == "PlaceholderWithDefault":
remove = False
if ignore_default and node_name in ignore_default:
op_type = 'Placeholder'
input_names = []
elif use_default and node_name in use_default:
remove = True
elif node_name.endswith('keras_learning_phase'):
logger.warning(
"Removing optional input %s that appears to be a keras learning phase parameter. "
"Use --ignore_default to force this into an input.",
node_name)
remove = True
if remove:
op_type = 'Identity'
graph_inputs = [
inp for inp in graph_inputs if inp != node_name + ":0"
]
onnx_node = helper.make_node(op_type,
input_names,
output_names,
name=node_name,
**onnx_attr)
onnx_nodes.append(onnx_node)
dtypes = {k: tf_utils.map_tf_dtype(v) for k, v in tf_dtypes.items()}
if graph_outputs is None:
output_to_node = {
out: node.name
for node in onnx_nodes for out in node.output
}
node_to_outputs = {node.name: list(node.output) for node in onnx_nodes}
used_nodes = set(output_to_node[out] for node in onnx_nodes
for out in node.input)
unused_nodes = [
node for node in onnx_nodes if node.name not in used_nodes
]
graph_outputs = [
out for node in unused_nodes for out in node_to_outputs[node.name]
]
graph_outputs_mapped = [
resolve_output(out, op_info, func_name) for out in graph_outputs
]
g = Graph(onnx_nodes,
output_shapes,
dtypes,
input_names=graph_inputs,
output_names=graph_outputs_mapped,
is_subgraph=func is not None,
graph_name=graph_name)
g.rename_tensors(dict(zip(graph_outputs_mapped, graph_outputs)))
return g
def version_9(cls, ctx, node, **kwargs):
data_inp = node.input[0]
segment_inp = node.input[1]
data_shape = ctx.get_shape(data_inp)
utils.make_sure(data_shape is not None,
"Segment ops require input rank to be known")
data_rank = len(data_shape)
data_np_dtype = utils.map_onnx_to_numpy_type(ctx.get_dtype(data_inp))
seg_np_dtype = utils.map_onnx_to_numpy_type(ctx.get_dtype(segment_inp))
data_is_float = np.dtype(data_np_dtype).kind == 'f'
data_is_int = np.dtype(data_np_dtype).kind == 'i'
utils.make_sure(data_is_float or data_is_int,
"dtype for Segment ops must be float or int")
if node.type == "SegmentSum":
onnx_op = "ReduceSum"
identity_value = np.array(0, dtype=data_np_dtype)
elif node.type == "SegmentProd":
onnx_op = "ReduceProd"
identity_value = np.array(1, dtype=data_np_dtype)
elif node.type == "SegmentMax":
onnx_op = "ReduceMax"
if data_is_float:
identity_value = np.array('-inf', dtype=data_np_dtype)
else:
identity_value = np.iinfo(data_np_dtype).min
elif node.type == "SegmentMin":
onnx_op = "ReduceMin"
if data_is_float:
identity_value = np.array('inf', dtype=data_np_dtype)
else:
identity_value = np.iinfo(data_np_dtype).max
max_segment = ctx.make_node("ReduceMax", [segment_inp],
attr={
'axes': [0],
'keepdims': 0
})
one_const = ctx.make_const(utils.make_name("const_one"),
np.array(1, dtype=seg_np_dtype))
identity_const = ctx.make_const(utils.make_name("const_identity"),
identity_value)
num_segments = ctx.make_node(
"Add", [max_segment.output[0], one_const.output[0]])
# ORT doesn't support bool for OneHot so we use float32 and cast to bool
onehot_values = ctx.make_const(utils.make_name("onehot_values"),
np.array([0, 1], dtype=np.float32))
one_hot_node = ctx.make_node(
"OneHot",
[segment_inp, num_segments.output[0], onehot_values.output[0]],
attr={'axis': 0})
one_hot_bool = ctx.make_node("Cast", [one_hot_node.output[0]],
attr={"to": onnx_pb.TensorProto.BOOL})
one_hot_unsqueeze = one_hot_bool
if data_rank > 1:
new_dims = list(range(2, 2 + data_rank - 1))
one_hot_unsqueeze = ctx.make_node("Unsqueeze",
[one_hot_bool.output[0]],
attr={'axes': new_dims})
mul_node = ctx.make_node(
"Where",
[one_hot_unsqueeze.output[0], data_inp, identity_const.output[0]])
shapes = node.output_shapes
dtypes = node.output_dtypes
ctx.remove_node(node.name)
ctx.make_node(onnx_op, [mul_node.output[0]],
attr={
'axes': [1],
'keepdims': 0
},
name=node.name,
outputs=node.output,
shapes=shapes,
dtypes=dtypes)
