TPUEstimator 不适用于 use_tpu=False

TPUEstimator does not work with use_tpu=False

我正在尝试 运行 在 CPU 上本地使用 TPUEstimator 的模型,首先通过在估计器初始化时设置 use_tpu=False 来验证它是否有效。当 运行ning 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?

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

上 运行

通常的做法是使用命令行标志来设置是否正在使用 TPU。此标志用于切换 CrossShardOptimizeruse_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