Tensorflow中四维张量的tf.nn.top_k索引的二进制掩码？

Question

我有一个形状的四维张量（10,32,32,128）。我想为所有前N个元素生成一个二进制掩码。

arr = tf.random_normal(shape=(10, 32, 32, 128)) 
values, indices = tf.nn.top_k(arr, N=64) 

我的问题是如何使用indices通过tf.nn.top_k

Answer 1

如果有人要的答案返回给获得相同形状的二进制掩码为arr：在这里不言而喻。

K = 64 
arr = tf.random_normal(shape=(10, 32, 32, 128)) 
values, indices = tf.nn.top_k(arr, k=K, sorted=False) 

temp_indices = tf.meshgrid(*[tf.range(d) for d in (tf.unstack(
     tf.shape(arr)[:(arr.get_shape().ndims - 1)]) + [K])], indexing='ij') 
temp_indices = tf.stack(temp_indices[:-1] + [indices], axis=-1) 
full_indices = tf.reshape(temp_indices, [-1, arr.get_shape().ndims]) 
values = tf.reshape(values, [-1]) 

mask_st = tf.SparseTensor(indices=tf.cast(
     full_indices, dtype=tf.int64), values=tf.ones_like(values), dense_shape=arr.shape) 
mask = tf.sparse_tensor_to_dense(tf.sparse_reorder(mask_st))