2

我需要使用 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
4

1 回答 1

3

这是由于（在 TPU 上调用，请参见此处）的限制。tf.wheretf.boolean_mask

    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]，我不知道。

于 2019-04-25T06:42:53.230 回答