тензор маски с 0 в соответствии со значением по некоторому определенному индексу в самом внутреннем измерении

Question

Допустим, у меня есть тензор A формы [x, y, z].

Для объяснения, давайте предположим, что 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]]]

Как правильно использовать тензор потока для выполнения этой операции? Я попробовал несколько подходов и был ограничен тем, что работа с тензорами, которые содержат тензор переменного размера, болезненна в тензорном потоке.

Единственное решение, которое я могу придумать, - это использовать map_fn для перебора тензора (до измерения -2). Но использование map_fn сложно и ухудшит производительность, потому что

1. Если у меня есть тензоры более высокого ранга (скажем, 4+), необходимо использовать несколько map_fn внутренностей map_fn.

2. map_fn не может работать на графическом процессоре и может снизить производительность, особенно в случае большого набора данных.

Кто-нибудь может пролить свет на это?

Answer 1

Вот, наслаждайтесь (проверено):

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

---

Наиболее эффективная однострочная версия с идеями Ответ Олдрима и комментарий , хотя и гораздо менее понятный для чтения:

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

Answer 2

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