Question

我正在Tensorflow中实现特殊的损失功能。这是一个特殊函数的numpy样式代码，该函数选择前q个元素并掩盖每一行和每一列中的其他元素。请注意，A是n*n矩阵，而q是小于n的整数。

def thresh(A, q):
    A_ = A.copy()
    n = A_.shape[1]
    for i in range(n):
        A_[i, :][A_[i, :].argsort()[0:n - q]] = 0
        A_[:, i][A_[:, i].argsort()[0:n - q]] = 0
    return A_

现在的问题是我有一个Tensorflow张量A，其形状为(n,n)，我想实现与numpy相同的逻辑。但是，我不能使用索引直接将值分配给张量A。 anyont有一些解决方案吗？

Answer 1

TLDR;

我们可以创建一个函数来屏蔽除顶部k元素以外的所有元素，如下所示：

def mask_all_but_top_k(X, k):
  n = X.shape[1]
  top_k_indices = tf.math.top_k(X, k).indices
  mask = tf.reduce_sum(tf.one_hot(top_k_indices, n), axis=1)
  return mask * X

不幸的是，tf.map.top_k不允许我们指定尺寸，但是我们当然可以通过首先转置X然后将结果与tf.transpose()

来明智地复制此列。

说明

我们可以通过创建一个由1和0组成的掩码然后逐元素相乘来达到目标。

例如，考虑到n=4, k=2的情况，我们具有以下矩阵：

array([[0.67757607, 0.74070597, 0.89508283, 0.11858773],
       [0.7661159 , 0.8737055 , 0.73599136, 0.1552105 ],
       [0.7093129 , 0.44203556, 0.48861897, 0.83231044],
       [0.24682868, 0.36648738, 0.92984104, 0.9881872 ]], dtype=float32)

然后我们可以使用tf.math.top_k函数来获取矩阵每一行中前2个值的索引：

top_k_indices = tf.math.top_k(X, 2).indices

现在，我们将使用一些技巧来首先one_hot进行编码：

tf.one_hot(top_k_indices, 4)
array([[[0., 0., 1., 0.],
        [0., 1., 0., 0.]],

       [[0., 1., 0., 0.],
        [1., 0., 0., 0.]],

       [[0., 0., 0., 1.],
        [1., 0., 0., 0.]],

       [[0., 0., 0., 1.],
        [0., 0., 1., 0.]]], dtype=float32)>

然后在倒数第二个维度上reduce_sum来创建我们的遮罩：

tf.reduce_sum(tf.one_hot(top_k_indices, 4), axis=1)
array([[0., 1., 1., 0.],
       [1., 1., 0., 0.],
       [1., 0., 0., 1.],
       [0., 0., 1., 1.]], dtype=float32)>

现在我们可以进行Hadamard（逐元素）乘法以获得所需的结果：

array([[0.        , 0.74070597, 0.89508283, 0.        ],
       [0.7661159 , 0.8737055 , 0.        , 0.        ],
       [0.7093129 , 0.        , 0.        , 0.83231044],
       [0.        , 0.        , 0.92984104, 0.9881872 ]], dtype=float32)>

将所有这些放在一起，我们可以创建一个函数，该函数将按行方式屏蔽除顶部k元素之外的所有元素，如下所示：

def mask_all_but_top_k(X, k):
  n = X.shape[1]
  top_k_indices = tf.math.top_k(X, k).indices
  mask = tf.reduce_sum(tf.one_hot(top_k_indices, n), axis=1)
  return mask * X

如何在Tensorflow张量中的每一行和每一列中选择前q个元素？

1 个答案:

TLDR;

说明