diff --git a/README.md b/README.md index 6f6aaf0d..0a0d5bd9 100644 --- a/README.md +++ b/README.md @@ -70,7 +70,7 @@ Video speed is adjusted approximately 50 times slower than actual speed. $ docker run --rm -it \ -v `pwd`:/workdir \ -w /workdir \ -ghcr.io/pinto0309/onnx2tf:1.0.23 +ghcr.io/pinto0309/onnx2tf:1.0.24 or diff --git a/onnx2tf/__init__.py b/onnx2tf/__init__.py index 75a92add..1107ed75 100644 --- a/onnx2tf/__init__.py +++ b/onnx2tf/__init__.py @@ -1,3 +1,3 @@ from onnx2tf.onnx2tf import convert, main -__version__ = '1.0.23' +__version__ = '1.0.24' diff --git a/onnx2tf/ops/Slice.py b/onnx2tf/ops/Slice.py index 6e12ea5f..98b4d670 100644 --- a/onnx2tf/ops/Slice.py +++ b/onnx2tf/ops/Slice.py @@ -63,6 +63,8 @@ def make_node( ) input_tensor = tf_layers_dict[input_tensor.name]['tf_node'] \ if isinstance(input_tensor, gs.Variable) else input_tensor + input_tensor_shape = input_tensor.shape + input_tensor_rank = len(input_tensor_shape) starts = None if len(graph_node.inputs) >= 2: @@ -123,61 +125,104 @@ def make_node( } # Generation of TF OP - slice_len = len(starts) - - if slice_len == 1: - # Shape output as int64 since the spec implicitly allows int64 - full_sizes = None - if None not in input_tensor.shape: - full_sizes = np.asarray(input_tensor.shape, dtype=np.int64) - else: - full_sizes = [s if s is not None else -1 for s in input_tensor.shape] - - updated_full_sizes = [0] * len(input_tensor.get_shape()) - updated_full_begin = [0] * len(input_tensor.get_shape()) - updated_starts = [0] * slice_len - updated_ends = [0] * slice_len - - if axes is None: - axes = [i for i in range(input_tensor.shape.rank)] - - for axis in range(input_tensor.shape.rank): - if axis not in axes: - # Update the sizes for axes that are not in the axes attribute - # No need to change the default of 0 in begins - updated_full_sizes[axis] = full_sizes[axis] - else: - # Update the begins and sizes for each axis in the axes attribute - for i in range(slice_len): - if axis == axes[i]: - updated_starts[i] = full_sizes[axis] + starts[i] if starts[i] < 0 else starts[i] - updated_ends[i] = full_sizes[axis] + ends[i] if ends[i] < 0 else ends[i] - if full_sizes[axis] is not None: - updated_ends[i] = min(full_sizes[axis], updated_ends[i]) - updated_starts[i] = min(full_sizes[axis], updated_starts[i]) - - updated_full_begin[axis] = updated_starts[i] - updated_full_sizes[axis] = updated_ends[i] - - starts = updated_full_begin - ends = updated_full_sizes - - tf_layers_dict[graph_node_output.name]['tf_node'] = \ - tf.strided_slice( - input_=input_tensor, - begin=starts, - end=ends, - name=graph_node.name, + + # first of all, get the input tensor shape + is_axes_negative = tf.less(axes, tf.zeros_like(axes)) + axes = tf.where(is_axes_negative, axes + tf.cast(input_tensor_rank, axes.dtype), axes) + + # expand a dimension of 1 at the end + sparse_indices = tf.cast(tf.expand_dims(axes, -1), tf.int64) + + # build the indexed dimension sizes as sparse_shape + sparse_shape = tf.gather_nd(params=input_tensor_shape, indices=sparse_indices) + sparse_shape = tf.cast(sparse_shape, tf.int64) + + # take care of starts, ends that are larger than the dim size. + starts_min = tf.minimum(starts, sparse_shape) + ends_min = tf.minimum(ends, sparse_shape) + + # need to densify everything for the inputs to slice + # the output shape is the input_tensor rank + output_shape = tf.reshape(input_tensor_rank, [1]) + output_shape = tf.cast(output_shape, tf.int64) + + def tf_SparseTensor(sparse_indices, pos, sparse_shape, start_or_end): + dense = None + # take care of starts, ends that are negative + if start_or_end == 'start': + is_starts_negative = tf.less(pos, tf.zeros_like(pos)) + final = tf.where(is_starts_negative, pos + sparse_shape, pos) + sparse_tensor = tf.sparse.reorder( + sp_input=tf.sparse.SparseTensor( + indices=sparse_indices, + values=final, + dense_shape=sparse_shape, + ) ) - else: - tf_layers_dict[graph_node_output.name]['tf_node'] = \ - tf.strided_slice( - input_=input_tensor, - begin=starts, - end=ends, - strides=steps, - name=graph_node.name, + dense = tf.sparse.to_dense(sparse_tensor) + elif start_or_end == 'end': + is_ends_negative = tf.less(pos, tf.zeros_like(pos)) + final = tf.where(is_ends_negative, pos + sparse_shape, pos) + sparse_tensor = tf.sparse.reorder( + sp_input=tf.sparse.SparseTensor( + indices=sparse_indices, + values=final, + dense_shape=sparse_shape, + ) + ) + dense_ends = tf.sparse.to_dense( + sparse_tensor, + default_value=tf.constant(-1, dtype=dense_begins.dtype) + ) + dense = tf.where( + tf.equal(dense_ends, tf.constant(-1, dtype=dense_begins.dtype)), + input_tensor_shape, + dense_ends, + ) + return dense + + # create dense tensor, pad 0 as default begins + dense_begins = Lambda( + tf_SparseTensor, + arguments={ + 'pos': starts_min, + 'sparse_shape': output_shape, + 'start_or_end': 'start', + } + )(sparse_indices) + # create dense tensor, pad -1 for next step + dense_ends = Lambda( + tf_SparseTensor, + arguments={ + 'pos': ends_min, + 'sparse_shape': output_shape, + 'start_or_end': 'end' + } + )(sparse_indices) + + # create dense tensor for steps if not already so + if steps is not None: + dense_steps = tf.sparse.reorder( + sp_input=tf.sparse.SparseTensor( + sparse_indices, + steps, + output_shape, ) + ) + dense_steps = tf.sparse.to_dense( + dense_steps, + default_value=tf.constant(1, dtype=steps.dtype) + ) + else: + dense_steps = tf.ones(input_tensor_shape, ends.dtype) + + tf_layers_dict[graph_node_output.name]['tf_node'] = \ + tf.strided_slice( + input_=input_tensor, + begin=dense_begins, + end=dense_ends, + strides=tf.cast(dense_steps, dtype=dense_begins.dtype), + ) # Generation of Debug Info tf_layers_dict[graph_node_output.name]['tf_node_info'] = \