在TensorFlow中将张量的k个最大元素设置为零

时间:2016-06-02 13:55:58

标签: tensorflow

我想在 h 的每一行中找到 k 最大元素,并将零值设置为这些最大元素。

我可以使用top_k函数选择每行最高值的索引,如:

<input type="file />"

但我无法使用top_k返回的索引来更新张量。

我该怎么做?提前谢谢......

3 个答案:

答案 0 :(得分:10)

这有点棘手,也许有更好的解决方案。 tf.scatter_update()在这里不起作用,因为它只能修改第一维的张量部分(例如,不是第一行和第二列中的元素)。

您必须从values获取indicestf.nn.top_k()以创建稀疏Tensor并将其减去初始Tensor x

x = tf.constant([[6., 2., 0.], [0., 4., 5.]])  # of type tf.float32

k = 2
values, indices = tf.nn.top_k(x, k, sorted=False)  # indices will be [[0, 1], [1, 2]], values will be [[6., 2.], [4., 5.]]

# We need to create full indices like [[0, 0], [0, 1], [1, 2], [1, 1]]
my_range = tf.expand_dims(tf.range(0, indices.get_shape()[0]), 1)  # will be [[0], [1]]
my_range_repeated = tf.tile(my_range, [1, k])  # will be [[0, 0], [1, 1]]

# change shapes to [N, k, 1] and [N, k, 1], to concatenate into [N, k, 2]
full_indices = tf.concat([tf.expand_dims(my_range_repeated, 2), tf.expand_dims(indices, 2)], axis=2)
full_indices = tf.reshape(full_indices, [-1, 2])

to_substract = tf.sparse_to_dense(full_indices, x.get_shape(), tf.reshape(values, [-1]), default_value=0.)

res = x - to_substract  # res should be all 0.

答案 1 :(得分:2)

我遇到了相反的问题,想要一个支持渐变的操作。 top_k不支持梯度传播,因此一种好的方法是在c ++中实现该函数。

找到

top_k c ++代码here

你的操作内核看起来像这样:

template <typename T>
class MakeSparseOp : public OpKernel {
   public:
    explicit MakeSparseOp(OpKernelConstruction *context) : OpKernel(context) {}

    void Compute(OpKernelContext *context) override {
        // Grab the input tensors
        const auto &k_in = context->input(1);

        OP_REQUIRES(context, TensorShapeUtils::IsScalar(k_in.shape()),
                    errors::InvalidArgument("k must be scalar, got shape ",
                                            k_in.shape().DebugString()));

        int k = k_in.scalar<int32>()();
        OP_REQUIRES(context, k >= 0,
                    errors::InvalidArgument("Need k >= 0, got ", k));

        const Tensor &x_in = context->input(0);
        OP_REQUIRES(context, x_in.dims() >= 1,
                    errors::InvalidArgument("input must be >= 1-D, got shape ",
                                            x_in.shape().DebugString()));
        OP_REQUIRES(
            context, x_in.dim_size(x_in.dims() - 1) >= k,
            errors::InvalidArgument("input must have at least k columns"));



        // Flattening the input tensor
        const auto &x = x_in.flat_inner_dims<T>();

        const auto num_rows = x.dimension(0);
        const auto num_cols = x.dimension(1);

        TensorShape output_shape = x_in.shape();

        // Create an output tensor
        Tensor *x_out = nullptr;
        OP_REQUIRES_OK(context,
                       context->allocate_output(0, output_shape, &x_out));

        /*
         * Get the top k values along the first dimension for input
         */

        auto x_sparse = x_out->flat_inner_dims<T>();

        if (k == 0) return;  // Nothing to do

        // Using TopN to get the k max element
        gtl::TopN<std::pair<T, int32>> filter(k);

        x_sparse = x; // Copy all elements

        for (int r = 0; r < num_rows; r++) {
            // Processing a row at a time
            for (int32 c = 0; c < num_cols; c++) {
                // The second element is the negated index, so that lower-index
                // elements
                // are considered larger than higher-index elements in case of
                // ties.
                filter.push(std::make_pair(x(r, c), -c));

            }

            for (auto top_k_it = filter.unsorted_begin();
                 top_k_it != filter.unsorted_end(); ++top_k_it) {
                x_sparse(r, -top_k_it->second) = 0; // Set max k to zero
            }

            filter.Reset();
        }
    }
};

我对相关问题的实施是here

答案 2 :(得分:0)

最近在tensorflow中可用scatter_nd_update函数,以下是来自Oliver的答案的修改版本。

k = 2
val_to_replace_with = -333
x = tf.Variable([[6., 2., 0.], [0., 4., 5.]])  # of type tf.float32


values, indices = tf.nn.top_k(x, k, sorted=False)  # indices will be [[0, 1], [1, 2]], values will be [[6., 2.], [4., 5.]]
# We need to create full indices like [[0, 0], [0, 1], [1, 2], [1, 1]]
my_range = tf.expand_dims(tf.range(0, tf.shape(indices)[0]), 1)  # will be [[0], [1]]
my_range_repeated = tf.tile(my_range, [1, k])  # will be [[0, 0], [1, 1]]
# change shapes to [N, k, 1] and [N, k, 1], to concatenate into [N, k, 2]
full_indices = tf.concat([tf.expand_dims(my_range_repeated, -1), tf.expand_dims(indices, -1)], axis=2)
full_indices = tf.reshape(full_indices, [-1, 2])


# only significant modification -----------------------------------------------------------------
updates = val_to_replace_with + tf.zeros([tf.size(indices)], dtype=tf.float32)
c = tf.scatter_nd_update(x, full_indices, updates)
# only significant modification -----------------------------------------------------------------


with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    print(sess.run(c))