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)