如何在 TPU 上实现布尔屏蔽

Question

我需要使用 TPUEstimator 实现布尔屏蔽操作。 tf.boolean_mask 未实现。有解决办法吗？

以下代码在 CPU 和 GPU 上非常适合我的目的：

  all_out = model.get_sequence_output()
  P = tf.boolean_mask(all_out, P_mask)

all_out 是一个形状为 [?, 128, 768] 的张量

P_mask 是形状 [?, 128] 并且第二维是 one-hot 编码以表示要提取的所需张量。

P 的期望形状是 [?,768]

当我使用 TPUEstimator 在 TPU 上运行此程序时，我收到以下错误消息：

Compilation failure: Detected unsupported operations when trying to
compile graph _functionalize_body_1[] on XLA_TPU_JIT: Where (No 
registered 'Where' OpKernel for XLA_TPU_JIT devices compatible with node
node boolean_mask/Where

Answer 1

这是由于tf.where 上的limitation（在TPU 上由tf.boolean_mask 调用，也请参见here。

    tf.where    Both x and y must be non-None. If both x and y are None, the operator would not have a static shape.

根本原因是这样做不会获得静态形状，因此截至今天，tpu 对此并不满意。

如果您的唯一目的是最终计算损失或总和，那么重写您的代码可能是可行的。

Rewrite this:
   reduce_sum(gather_nd(tf.where(cond),Y))
to this:
   reduce_sum(Y * tf.cast(cond))

但是，如果您确实需要屏蔽输出 [?, 768] 的动态形状，我不知道。