TPUEstimator 不适用于 use_tpu=False答案

【问题标题】：TPUEstimator does not work with use_tpu=FalseTPUEstimator 不适用于 use_tpu=False
【发布时间】：2019-01-25 00:00:10
【问题描述】：

我正在尝试首先在 CPU 上本地使用 TPUEstimator 运行模型，以通过在估计器初始化上设置 use_tpu=False 来验证它是否有效。运行 train 时出现此错误。

InternalError: failed to synchronously memcpy host-to-device: host 0x7fcc7e4d4000 to device 0x1deffc002 size 4096: Failed precondition: Unable to enqueue when not opened, queue: [0000:00:04.0 PE0 C0 MC0 TN0 Queue HBM_WRITE]. State is: CLOSED [[Node: optimizer/gradients/neural_network/fully_connected_2/BiasAdd_grad/BiasAddGrad_G14 = _Recv[client_terminated=false, recv_device="/job:worker/replica:0/task:0/device:TPU:0", send_device="/job:worker/replica:0/task:0/device:CPU:0", send_device_incarnation=-7832507818616568453, tensor_name="edge_42_op...iasAddGrad", tensor_type=DT_FLOAT, _device="/job:worker/replica:0/task:0/device:TPU:0"]()]]

它看起来仍在尝试使用 TPU，正如它所说的 recv_device="/job:worker/replica:0/task:0/device:TPU:0"。为什么use_tpu设置为False时尝试使用TPU？

【问题讨论】：

标签： google-cloud-platform google-cloud-tpu

【解决方案1】：

您使用的是什么优化器？如果您使用tf.contrib.tpu.CrossShardOptimizer 并且use_tpu 设置为False，则可能会发生此类错误。优化器正在尝试跨 TPU 内核分片工作，但不能因为您在 CPU 上运行。

通常的做法是有一个命令行标志来设置是否正在使用 TPU。此标志用于切换 CrossShardOptimizer 和 use_tpu 之类的内容。例如，在 MNIST 参考模型中：

if FLAGS.use_tpu:
  optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)

https://github.com/tensorflow/models/blob/ad3526a98e7d5e9e57c029b8857ef7b15c903ca2/official/mnist/mnist_tpu.py#L102

【讨论】：