时间戳

首页 > 解决方案 > 如何在 Tensorflow 张量的每一行和每一列中选择前 q 个元素？

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

问题描述

我正在 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_

现在的问题是我有一个A形状为的 Tensorflow 张量(n,n)，我想实现与 numpy 相同的逻辑。但是，我不能使用索引直接为张量赋值A。有没有关于它的一些解决方案？

标签： pythontensorflowdeep-learning

解决方案

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 (element-wise) 乘法来得到想要的结果：

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

推荐阅读