Question

我想根据位置列表（索引）修改3D张量中的特定向量：

#indices: 1D vector of positions, indices.shape: (k)

mask = np.zeros(k, n, m)
for i in range(k):
  mask[i][indices[i]] = 1

此蒙版将应用于另一个3D张量（相同形状），我想保留特定向量，并将其余部分归零。

在TensorFlow中构建此类蒙版的最佳方法是什么？我可以使用assign op进行循环，但我想找到一个更优雅的解决方案。也许使用tf.scatter_nd？

编辑：示例：

>>> mask_before
array([[[ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.]],

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

       [[ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.]]])

>>> indices
array([2, 1, 4])

>>> mask_after
array([[[ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.],
        [ 1.,  1.,  1.,  1.],
        [ 0.,  0.,  0.,  0.],
        [ 0.,  0.,  0.,  0.]],

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

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

Answer 1

一种矢量化方式是扩展维度，从而利用broadcasting -

np.tile((indices[:,None,None] == np.arange(n)[:,None]), m)

示例运行 -

In [755]: # Sample Setup
     ...: indices = np.array([2,3,1])
     ...: 
     ...: k = 3
     ...: n = 4
     ...: m = 2
     ...: mask = np.zeros((k, n, m),dtype=bool)
     ...: for i in range(k):
     ...:     mask[i][indices[i]] = 1

In [756]: out = np.tile((indices[:,None,None] == np.arange(n)[:,None]), m)

In [757]: np.allclose(out, mask)
Out[757]: True

要移至tensorflow，我们会在那里找到对方：

tf.expand_dims和tf.tile。

Answer 2

这里是Divakar答案的相应TF代码：

indices = tf.constant([2,1,4])
a1 = tf.expand_dims(indices, axis=1)
a1 = tf.expand_dims(a1, axis=1)
a2 = tf.range(5)
a2 = tf.expand_dims(a2, axis=1)
a3 = tf.equal(a1, a2)
mask = tf.tile(a3, [1,1,4])


>>> tf.cast(mask, dtype=tf.int8)
<tf.Tensor: id=55, shape=(3, 5, 4), dtype=int8, numpy=
array([[[0, 0, 0, 0],
    [0, 0, 0, 0],
    [1, 1, 1, 1],
    [0, 0, 0, 0],
    [0, 0, 0, 0]],

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

   [[0, 0, 0, 0],
    [0, 0, 0, 0],
    [0, 0, 0, 0],
    [0, 0, 0, 0],
    [1, 1, 1, 1]]], dtype=int8)>

如何屏蔽张量中每行的特定切片

2 个答案: