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]，第二维为一维编码，表示要提取的所需张量。

所需的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]的动态形状，我不知道。

如何在TPU上实现布尔掩码

1 个答案: