def chk_idx_out_of_bounds(cls, data, indices):
""" Check indices out of bounds for ScatterND and GatherND
In Tensorflow GPU version, if an out of bound index is found,
a 0 is stored in the corresponding output value for GatherND;
and the index is ignored for ScatterND/TensorScatterNDUpdate.
But ONNX spec state that it is an error if any index values
are out of bounds. Therefore the converter need to run this
function to verify all the indices are in bounds before send
it to Tensoflow. If out of bound is detected then the caller
of this function need to throw InvalidArgumentError exception.
"""
data_shape = tf_shape(data)
indices_shape = tf_shape(indices)
def _chk_idx_out_of_bounds(i, result):
indices_i = tf.transpose(indices)[i]
limit_i = tf.cast(data_shape, indices.dtype)[i]
cond1 = tf.greater_equal(indices_i, tf.negative(limit_i))
cond2 = tf.less(indices_i, limit_i)
result = tf.reduce_all(tf.logical_and(cond1, cond2))
return i + 1, result
_, result = tf.while_loop(
lambda i, result: tf.logical_and(tf.less(i, indices_shape[-1]),
result), _chk_idx_out_of_bounds,
[tf.zeros([], tf.int64), True])
return result
def process_neg_idx(cls, data, indices):
""" Convert all the negative indices to positive
GatherND and ScatterND/TensorScatterNDUpdate in Tensorflow
doesn't support negative indices. Therefore need to run this
function to convert all the negative indices to positive before
send it to Tensorflow.
"""
data_shape = tf_shape(data)
indices_shape = tf_shape(indices)
max_i = tf.cast(data_shape[:indices_shape[-1]], indices.dtype)
return tf.math.floormod(tf.add(indices, max_i), max_i)
def _common(cls, node, **kwargs):
tensor_dict = kwargs["tensor_dict"]
input = tensor_dict[node.inputs[0]]
x_shape = tf_shape(input)
attrs = copy.deepcopy(node.attrs)
axis = attrs.get("axis", 0)
axis = axis if axis >= 0 else len(x_shape) + axis
if "split" in node.attrs:
split = attrs["split"]
elif len(node.inputs) == 2: # since version 1
split = tensor_dict[node.inputs[1]]
else:
per_part = x_shape[axis] / len(node.outputs)
if input.shape.is_fully_defined():
if int(per_part) != per_part:
raise ValueError("Split can not be evenly divided.")
split = [int(per_part)] * len(node.outputs)
else:
split = [tf.cast(per_part, tf.int32)] * len(node.outputs)
attrs["num_or_size_splits"] = split
return list(
cls.make_tensor_from_onnx_node(node,
inputs=[input],
attrs=attrs,
**kwargs))
def _pad_input(self):
"""
Pad the input according to the parameters
"""
# check if we need to do any padding at all
if not self.ceil_mode and ((type(self.padding) is list and self.padding
== [0] * self.spatial_size * 2)
or self.padding == "VALID"):
self.pads = np.array([0] * self.spatial_size * 2)
return (self.input, self.pads)
in_spatial_shape = self.input_shape[1:self.spatial_size + 1]
pads = self._calc_pads(in_spatial_shape)
if self.is_known_shape and np.count_nonzero(pads) == 0:
self.pads = pads
return (self.input, pads)
tf_paddings = [[0, 0]]
for i in range(self.spatial_size):
tf_paddings += [[pads[i * 2], pads[i * 2 + 1]]]
tf_paddings += [[0, 0]]
self.input = tf.pad(self.input,
tf_paddings,
mode='CONSTANT',
constant_values=self.padding_constant)
# update input shape and pads values
self.input_shape = tf_shape(self.input)
self.pads = pads
def __init__(self,
input,
kernel_shape,
strides,
dilations,
padding="VALID",
ceil_mode=False,
count_include_pad=False,
pooling_type="MAX"):
self.input = tf.convert_to_tensor(input)
self.kernel_shape = kernel_shape
self.strides = strides
self.dilations = dilations
self.padding = padding
self.is_explicit_padding = type(padding) is list
self.ceil_mode = ceil_mode
self.count_include_pad = count_include_pad
self.pooling_type = pooling_type.upper()
self.is_known_shape = self.input.shape.is_fully_defined()
self.spatial_size = len(kernel_shape)
self.input_rank = self.spatial_size + 2
# if the rank is not defined, set it to the calculated input_rank
# rank should be known for ops like tf.gather_nd
if not input.shape.rank:
input.set_shape([None] * self.input_rank)
self.orig_input_shape = tf_shape(input)
self.input_shape = self.orig_input_shape
if pooling_type.startswith("MAX"):
self.padding_constant = input.dtype.min
else:
self.padding_constant = 0
def version_11(cls, node, **kwargs):
axis = node.attrs.get("axis", 0)
data = kwargs["tensor_dict"][node.inputs[0]]
indices = kwargs["tensor_dict"][node.inputs[1]]
updates = kwargs["tensor_dict"][node.inputs[2]]
# poocess negative axis
axis = axis if axis >= 0 else tf.add(tf.rank(data), axis)
# check are there any indices are out of bounds
result = cls.chk_idx_out_of_bounds_along_axis(data, axis, indices)
msg = 'ScatterElements indices are out of bounds, please double check the indices and retry.'
with tf.control_dependencies(
[tf.compat.v1.assert_equal(result, True, message=msg)]):
# process negative indices
indices = cls.process_neg_idx_along_axis(data, axis, indices)
# Calculate shape of the tensorflow version of indices tensor.
sparsified_dense_idx_shape = tf_shape(updates)
# Move on to convert ONNX indices to tensorflow indices in 2 steps:
#
# Step 1:
# What would the index tensors look like if updates are all
# dense? In other words, produce a coordinate tensor for updates:
#
# coordinate[i, j, k ...] = [i, j, k ...]
# where the shape of "coordinate" tensor is same as that of updates.
#
# Step 2:
# But the coordinate tensor needs some correction because coord
# vector at position axis is wrong (since we assumed update is dense,
# but it is not at the axis specified).
# So we update coordinate vector tensor elements at psotion=axis with
# the sparse coordinate indices.
idx_tensors_per_axis = tf.meshgrid(*list(
map(lambda x: tf.range(x, dtype=tf.dtypes.int64),
sparsified_dense_idx_shape)),
indexing='ij')
idx_tensors_per_axis[axis] = indices
dim_expanded_idx_tensors_per_axis = list(
map(lambda x: tf.expand_dims(x, axis=-1),
idx_tensors_per_axis))
coordinate = tf.concat(dim_expanded_idx_tensors_per_axis, axis=-1)
# Now the coordinate tensor is in the shape
# [updates.shape, updates.rank]
# we need it to flattened into the shape:
# [product(updates.shape), updates.rank]
indices = tf.reshape(coordinate, [-1, tf.rank(data)])
updates = tf.reshape(updates, [-1])
return [tf.tensor_scatter_nd_update(data, indices, updates)]
def version_10(cls, node, **kwargs):
inp = kwargs["tensor_dict"][node.inputs[0]]
dtype = inp.dtype
shape = tf_shape(inp)
zero = tf.zeros(shape, dtype)
dn = node.attrs.get("detect_negative", 1)
dp = node.attrs.get("detect_positive", 1)
# detecting only positive infinity, zero out elements < 0
if dn == 0:
inp = tf.maximum(zero, inp)
# detecting only negative infinity, zero out elements > 0
if dp == 0:
inp = tf.minimum(zero, inp)
return [cls.make_tensor_from_onnx_node(node, inputs=[inp], **kwargs)]
def chk_idx_out_of_bounds_along_axis(cls, data, axis, indices):
""" Check indices out of bounds for ScatterElement
In Tensorflow GPU version, if an out of bound index is found,
the index is ignored for ScatterND/TensorScatterNDUpdate.
But ONNX spec state that it is an error if any index values
are out of bounds. Therefore the converter need to run this
function to verify all the indices are in bounds along the
axis before send it to Tensoflow. If out of bound is detected
then the caller of this function need to throw
InvalidArgumentError exception.
"""
data_shape = tf.cast(tf_shape(data), indices.dtype)
limit = data_shape[axis]
cond1 = tf.greater_equal(indices, tf.negative(limit))
cond2 = tf.less(indices, limit)
return tf.logical_and(cond1, cond2)
def version_10(cls, node, **kwargs):
x = kwargs["tensor_dict"][node.inputs[0]]
x_shape = tf_shape(x)
scales = kwargs["tensor_dict"][node.inputs[1]]
n_in_scales_is_one = tf.equal(scales[0], 1)
c_in_scales_is_one = tf.logical_or(tf.equal(scales[1], 1),
tf.equal(scales[3], 1))
assert_n_c_in_scales_are_ones = tf.Assert(
tf.logical_and(n_in_scales_is_one, c_in_scales_is_one), [scales])
with tf.control_dependencies([assert_n_c_in_scales_are_ones]):
x_in_NCHW_format = tf.equal(scales[1], 1)
h_w_scale = tf.where(x_in_NCHW_format, scales[2:], scales[1:3])
h_w_shape = tf.where(x_in_NCHW_format, x_shape[2:], x_shape[1:3])
new_h_w_shape = tf.cast(
h_w_scale * tf.cast(h_w_shape, scales.dtype), tf.int32)
mode = node.attrs.get("mode", "nearest")
if mode.lower() == "linear":
mode = tf.image.ResizeMethod.BILINEAR
else:
mode = tf.image.ResizeMethod.NEAREST_NEIGHBOR
def process_NCHW_format(x):
x_t = tf.transpose(x, perm=[0, 2, 3, 1])
y = tf.image.resize(x_t, size=new_h_w_shape, method=mode)
y_t = tf.transpose(y, perm=[0, 3, 1, 2])
return y_t
def process_NHWC_format(x):
y = tf.image.resize(x, size=new_h_w_shape, method=mode)
return y
output = tf.cond(x_in_NCHW_format, lambda: process_NCHW_format(x),
lambda: process_NHWC_format(x))
return [output]
def _remove_dilations(self):
"""
This method removes the dilations by extracting the values from
the input for every sliding window according to the dilations,
strides and kernel size and generates output that can be used by
pooling operations with strides = kernel_shape to accomplish
dilated pooling
Example:
Input: [[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11],
[ 12, 13, 14, 15]]
Kernel: [2, 2]
Dilations: [2, 2]
Strides: [1, 1]
Will return:
[[ 0, 2, 1, 3],
[ 8, 10, 9, 11],
[ 4, 6, 5, 7],
[ 12, 14, 13, 15]]
After max_pool2d with kernel_shape = strides = [2, 2]
the result is:
[[ 10, 11],
[ 14, 15]]
"""
input_shape = tf_shape(self.input)
in_spatial_shape = input_shape[1:self.spatial_size + 1]
channels_count = input_shape[self.spatial_size + 1]
# Initialize gather_ind with the range of channels
# e.g. [0 1]
gather_ind = tf.range(channels_count, dtype=tf.int64)
# convert the vector to column vector
# in the following logic we use column vectors
gather_ind = tf.expand_dims(gather_ind, 1)
# initilize the output_shape with zeros
# self.output_shape will contain the shape of the
# output tensor after the loop below is executed
self.output_shape = [0] * (self.spatial_size + 2)
self.output_shape[0] = input_shape[0]
"""
Loop over the input spatial dimensions starting from the
last (most internal) going up to the first dimension
On every step of the loop calculate the output indices and
map them to the input indices using `_calc_input_ind`,
then "combine" with the already calculated indices from the
previous dimensions using cartesian product.
For the following example input:
Input: [[ 0, 1, 2, 3],
[ 4, 5, 6, 7],
[ 8, 9, 10, 11],
[ 12, 13, 14, 15]]
Kernel: [2, 2]
Dilations: [2, 2]
Strides: [1, 1]
these are the steps that will be executed:
1. Initilize gather_ind = [[0]] # we have only 1 channel
2. Loop step 0 (axis 1):
filter_size = 3
output_size = 4
dim_ind = [[0]
[2]
[1]
[3]]
gather_ind = [[0 0]
[2 0]
[1 0]
[3 0]]
3. Loop step 1 (axis 0):
filter_size = 3
output_size = 4
dim_ind = [[0]
[2]
[1]
[3]]
gather_ind = [[0 0 0]
[0 2 0]
[0 1 0]
[0 3 0]
[2 0 0]
[2 2 0]
[2 1 0]
[2 3 0]
[1 0 0]
[1 2 0]
[1 1 0]
[1 3 0]
[3 0 0]
[3 2 0]
[3 1 0]
[3 3 0]]
These are the indices used for gather_nd operation to collect
the values from the input data.
"""
for dim in range(self.spatial_size - 1, -1, -1):
filter_size = (self.kernel_shape[dim] - 1) * \
self.dilations[dim] + 1
output_size = ((
(in_spatial_shape[dim] - filter_size) // self.strides[dim]) +
1) * self.kernel_shape[dim]
self.output_shape[dim + 1] = output_size
# initialize the output dimension index with the range of the
# dimension output size (e.g. 4): [0, 1, 2, 3]
dim_ind = tf.range(output_size)
# calculate the matching indices in the input data
# [0, 1, 2, 3] will calculate to [0, 2, 1, 3]
# from the above example
dim_ind = self._calc_input_ind(dim_ind, self.kernel_shape[dim],
self.dilations[dim],
self.strides[dim])
# convert to column vector
dim_ind = tf.expand_dims(dim_ind, 1)
# "combine" current dimension indices with the previous dimensions
# using cartesian product
gather_ind = tf_product(dim_ind, gather_ind)
# The result from the above loop for 2D data will be:
# [[y1, x1, c], [y2, x2, c], ..., [yn, xm, c]] where n is the height,
# m is the width and c is the channel number.
# set the channels count in the output_shape
self.output_shape[self.spatial_size + 1] = channels_count
# expand the dimensions to match the input dimensions + 1
for x in range(self.spatial_size):
gather_ind = tf.expand_dims(gather_ind, 0)
# dublicate the indices for every batch
gather_ind = tf.tile(gather_ind,
[input_shape[0]] + [1] * (self.spatial_size + 1))
# extract the selected values from the input
output = tf.gather_nd(self.input, gather_ind, batch_dims=1)
# reshape the output to the correct shape calculated earlier
output = tf.reshape(output, self.output_shape)
return output
