84 lines
3.2 KiB
Python
84 lines
3.2 KiB
Python
#!/usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
|
|
# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
|
|
"""A module for helper tensorflow ops.
|
|
|
|
This is originally implemented in TensorFlow Object Detection API.
|
|
"""
|
|
|
|
import tensorflow as tf
|
|
|
|
from mask_rcnn.object_detection import shape_utils
|
|
|
|
|
|
def indices_to_dense_vector(indices,
|
|
size,
|
|
indices_value=1.,
|
|
default_value=0,
|
|
dtype=tf.float32):
|
|
"""Creates dense vector with indices set to specific value and rest to zeros.
|
|
|
|
This function exists because it is unclear if it is safe to use
|
|
tf.sparse_to_dense(indices, [size], 1, validate_indices=False)
|
|
with indices which are not ordered.
|
|
This function accepts a dynamic size (e.g. tf.shape(tensor)[0])
|
|
|
|
Args:
|
|
indices: 1d Tensor with integer indices which are to be set to
|
|
indices_values.
|
|
size: scalar with size (integer) of output Tensor.
|
|
indices_value: values of elements specified by indices in the output vector
|
|
default_value: values of other elements in the output vector.
|
|
dtype: data type.
|
|
|
|
Returns:
|
|
dense 1D Tensor of shape [size] with indices set to indices_values and the
|
|
rest set to default_value.
|
|
"""
|
|
size = tf.cast(size, dtype=tf.int32)
|
|
zeros = tf.ones([size], dtype=dtype) * default_value
|
|
values = tf.ones_like(indices, dtype=dtype) * indices_value
|
|
|
|
return tf.dynamic_stitch([tf.range(size), tf.cast(indices, dtype=tf.int32)],
|
|
[zeros, values])
|
|
|
|
|
|
def matmul_gather_on_zeroth_axis(params, indices, scope=None):
|
|
"""Matrix multiplication based implementation of tf.gather on zeroth axis.
|
|
|
|
TODO(rathodv, jonathanhuang): enable sparse matmul option.
|
|
|
|
Args:
|
|
params: A float32 Tensor. The tensor from which to gather values.
|
|
Must be at least rank 1.
|
|
indices: A Tensor. Must be one of the following types: int32, int64.
|
|
Must be in range [0, params.shape[0])
|
|
scope: A name for the operation (optional).
|
|
|
|
Returns:
|
|
A Tensor. Has the same type as params. Values from params gathered
|
|
from indices given by indices, with shape indices.shape + params.shape[1:].
|
|
"""
|
|
|
|
params_shape = shape_utils.combined_static_and_dynamic_shape(params)
|
|
indices_shape = shape_utils.combined_static_and_dynamic_shape(indices)
|
|
params2d = tf.reshape(params, [params_shape[0], -1])
|
|
indicator_matrix = tf.one_hot(indices, params_shape[0])
|
|
gathered_result_flattened = tf.matmul(indicator_matrix, params2d)
|
|
return tf.reshape(gathered_result_flattened,
|
|
tf.stack(indices_shape + params_shape[1:]))
|