在 Tensorflow 中使用 2D 占位符遮盖 3D 占位符张量

Masking a 3D Placeholder Tensor with a 2D Placeholder in Tensorflow

我有一个动态大小为 A(批量大小、序列大小、5)的占位符和另一个具有零值和一个值以及动态大小 B(批量大小、序列大小)的占位符。我想使用第二个 2D 占位符来屏蔽第一个占位符,就好像张量 B[0][0]=0 的值那么 A[0][0][0:5] 将被设置为零,如果它是等于 B[0][0]=1 那么 A[0][0][0:5] 不会改变。

palceholder A:(shape=(2,2,5))    
   [[[ 1,  2,  3,  1,  4],
    [ 2,  3,  5,  2,  4]],
   [[ 2,  7,  5,  8,  1],
    [ 4,  5,  1,  3,  9]]] 

palceholder B:(shape=(2,2))
  [[ 1,  0],
   [ 0,  1]]

Tensor C= Mask(A,B)    
   [[[ 1,  2,  3,  1,  4],
    [ 0,  0,  0,  0,  0]],
   [[ 0,  0,  0,  0,  0],
    [ 4,  5,  1,  3,  9]]]

我试过 tf.boolean_mask 但它不适用于动态尺寸蒙版。

怎么样

import numpy as np
import tensorflow as tf

a = tf.placeholder(tf.int32, [None, None, 5])
b = tf.placeholder(tf.int32, [None, None])
c = tf.multiply(a, tf.expand_dims(b, -1))

with tf.Session() as sess:
    in_a = np.reshape(np.arange(20, dtype=np.int32), [2, 2, 5])
    in_b = np.eye(2)
    print("A: {}".format(in_a))
    print("B: {}".format(in_b))
    out_c = sess.run(c, feed_dict={a: in_a, b: in_b})
    print("C: {}".format(out_c))

打印

A: [[[ 0  1  2  3  4]
  [ 5  6  7  8  9]]

 [[10 11 12 13 14]
  [15 16 17 18 19]]]
B: [[1. 0.]
 [0. 1.]]
C: [[[ 0  1  2  3  4]
  [ 0  0  0  0  0]]

 [[ 0  0  0  0  0]
  [15 16 17 18 19]]]