python - 根据最内维度某个特定索引处的值将张量掩码为 0

Question

假设我有一个形状为 [x, y, z] 的张量 A。

为了便于解释，我们假设 A 的形状为 [2,4,3]:

[[[1,2,3],[2,2,3],[4,4,4],[1,1,1]], [[2,2,2],[2,2,2],[2,2,2],[3,3,3]]]

我想“掩盖”这个张量，使得，

if elements at index 1 at the inner most dimension equal to 2, then the surrounding tensors should not change, otherwise they are all change to 0.

在此示例中，张量应变为

[[[1,2,3],[2,2,3],[0,0,0],[0,0,0]], [[2,2,2],[2,2,2],[2,2,2],[0,0,0]]]

在 tensorflow 中执行此操作的正确方法是什么？我尝试了几种方法，但受到以下事实的限制:在 tensorflow 中使用包含可变大小张量的张量是痛苦的。

我能想到的唯一解决方案是使用map_fn来迭代张量(直到-2维度)。但使用 map_fn 很棘手，并且会损害性能，因为

1. 如果我有更高阶的张量(比如 4+)，则需要在 map_fn 内部使用多个 map_fn。

2. map_fn 无法在 GPU 上运行，可能会损害性能，尤其是在数据集较大的情况下。

有人能解释一下吗？

Answer 1

import tensorflow as tf

x = tf.constant([[[1,2,3],[2,2,3],[4,4,4],[1,1,1]], [[2,2,2],[2,2,2],[2,2,2],[2,3,3]]])

x_idx1 = x[..., 1]
mask = tf.cast(tf.equal(x_idx1, 2), tf.int32)
mask = tf.expand_dims(tf.cast(mask, x.dtype), -1)
masked_x = x * mask

with tf.Session() as sess:
    print(sess.run(masked_x))
    # [[[1 2 3], [2 2 3], [0 0 0], [0 0 0]]
    #  [[2 2 2], [2 2 2], [2 2 2], [0 0 0]]]))