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]]]

我想“屏蔽”这个张量，

如果最内层维度的索引1处的元素等于 2 ，则周围的张量不应更改，否则它们都会更改为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执行此操作的正确方法是什么？我已经尝试了几种方法，并受到以下事实的限制：使用包含可变大小张量的张量在张量流中是痛苦的。

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

如果我拥有更高等级的张量（例如4+），则需要使用多个map_fn内部map_fn。
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]]]))

Answer 2

在这里，享受（测试）：

import tensorflow as tf

a = tf.constant( [[[1,2,3],[2,2,3],[4,4,4],[1,1,1]], [[2,2,2],[2,2,2],[2,2,2],[3,3,3]]] )
a2 = a[ :, :, 1 ]
b = tf.where( tf.equal( a2, 2 ), tf.ones_like( a2 ), tf.zeros_like( a2 ) )[ :, :, None ]
c = tf.tile( b, [ 1, 1, a.get_shape()[ 2 ].value ] )
d = a * c

with tf.Session() as sess:
    print( sess.run ( d ) )

输出：

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

[[2 2 2]
    [2 2 2]
    [2 2 2]
    [0 0 0]]]

效率最高的单行版本，包含来自Aldream's answer和comment的提示，虽然阅读不太清楚：

import tensorflow as tf

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

b = a * tf.cast( tf.equal( a[ ..., 1 ], 2 ), a.dtype )[ ..., None ]

with tf.Session() as sess:
    print( sess.run ( b ) )

输出：

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

[[2 2 2]
    [2 2 2]
    [2 2 2]
    [0 0 0]]]

根据最内层维度的某个特定索引处的值，使用0表示张量张量

2 个答案: