python - 用形状为 (x, y) 的二维 bool 掩码掩蔽形状为 (x, y, z) 的 3 维张量

Question

我有以下张量，形状为 (2, 6, 2) 的输出:

[[[0.4 0.2]
  [0.7 0.5]
  [0.4 0.1]
  [0.5 0.4]
  [0.9 0.7]
  [0.2 0.9]]

 [[0.6 0.6]
  [0.3 0.5]
  [0.7 0.2]
  [0.8 0.1]
  [0.3 0.5]
  [0.4 0.7]]]

并具有以下 bool 掩码张量，形状为 (2, 6) 的掩码:

mask = tf.sequence_mask(lengths=[3, 4] maxlen=6)

[[ True  True  True False False False]
 [ True  True  True  True False False]]

我如何使用 mask(或其调整)来应用 masked_output = tf.boolean_mask(output, masks) 导致以下结果:

[[[0.4 0.2]
  [0.7 0.5]
  [0.4 0.1]
  [0.0 0.0]
  [0.0 0.0]
  [0.0 0.0]]

 [[0.6 0.6]
  [0.3 0.5]
  [0.7 0.2]
  [0.8 0.1]
  [0.0 0.0]
  [0.0 0.0]]]

编辑

以下更改，但它仍然看起来是一种困惑的方式。任何其他建议表示赞赏。

mask = tf.sequence_mask(lengths=[[3, 3], [4, 4]] maxlen=6)
mask = tf.transpose(mask, [0, 2, 1])

Answer 1

可能不是最有效的方法，但它有效

flat_mask = tf.reshape(mask, shape=(-1,))
expanded_flat_mask = tf.tile(flat_mask, [2])
expanded_mask = tf.reshape(expanded_flat_mask, shape=(2, 6, 2))
masked_output = tf.boolean_mask(output, expanded_mask)

或者另一种方式

masked_output = output * tf.cast(tf.expand_dims(mask, axis=-1), output.dtype)