python-3.x - TPU分类器InvalidArgumentError : No OpKernel was registered to support Op 'CrossReplicaSum' with these attrs

Question

我尝试使用Estimator API实现基于TPUEstimator的Tensorflow模型失败。在训练过程中出现错误:

InvalidArgumentError (see above for traceback): No OpKernel was registered to support Op 'CrossReplicaSum' with these attrs.  Registered devices: [CPU], Registered kernels: <no registered kernels>
[[Node: CrossReplicaSum_5 = CrossReplicaSum[T=DT_FLOAT](gradients/dense_2/BiasAdd_grad/tuple/control_dependency_1)]]

在一开始也有一个警告，尽管我不确定它是否相关:

WARNING:tensorflow:CrossShardOptimizer should be used within a tpu_shard_context, but got unset number_of_shards. Assuming 1.

这是模型函数的相关部分:

def model_fn(features, labels, mode, params):
"""A simple NN with two hidden layers of 10 nodes each."""
input_layer = tf.feature_column.input_layer(features, params['feature_columns'])

dense1 = tf.layers.dense(inputs=input_layer, units=10, activation=tf.nn.relu, kernel_initializer=tf.glorot_uniform_initializer())
dense2 = tf.layers.dense(inputs=dense1, units=10, activation=tf.nn.relu, kernel_initializer=tf.glorot_uniform_initializer())
logits = tf.layers.dense(inputs=dense2, units=4)

reshaped_logits = tf.reshape(logits, [-1, 1, 4])

onehot_labels = tf.one_hot(indices=tf.cast(labels, tf.int32), depth=4)

loss = tf.losses.softmax_cross_entropy(onehot_labels=onehot_labels, logits=reshaped_logits)

if mode == tf.estimator.ModeKeys.TRAIN:

    optimizer = tf.contrib.tpu.CrossShardOptimizer(tf.train.AdagradOptimizer(learning_rate=0.05))

    train_op = optimizer.minimize(
        loss=loss,
        global_step=tf.train.get_global_step())

我正在尝试通过将TPUEstimator标志设置为--use_tpu来使用False进行本地CPU执行。实例化了TPUEstimator并因此调用了train:

estimator_classifier = tf.contrib.tpu.TPUEstimator(
        model_fn=model_fn, 
        model_dir="/tmp/estimator_classifier_logs",
        config=tf.contrib.tpu.RunConfig(
            session_config=tf.ConfigProto(
                allow_soft_placement=True, log_device_placement=True),
            tpu_config=tf.contrib.tpu.TPUConfig()
        ),
        train_batch_size=DEFAULT_BATCH_SIZE,
        use_tpu=False,
        params={
            'feature_columns': feature_columns  
        }
    )

    tensors_to_log = {"probabilities": "softmax_tensor"}
    logging_hook = tf.train.LoggingTensorHook(tensors=tensors_to_log, every_n_iter=50)

    estimator_classifier.train(
        input_fn=data_factory.make_tpu_train_input_fn(train_x, train_y, DEFAULT_BATCH_SIZE),
        steps=DEFAULT_STEPS,
        hooks=[logging_hook]
    )

此错误的含义是什么，如何解决该错误？

Answer 1

上下文尚不清楚。

您是在Cloud TPU环境中还是在具有TPU硬件的环境中运行您的工作？

如果否，则可以预期。 TPUEstimator设计为主要用于Cloud TPU环境，在该环境中，后端工作器将所有内核正确链接到Tensorflow服务器。 CrossReplicaSum是为设备TPU(而非CPU)注册的内核的一部分。

如果是，您是否正确设置了主地址。根据日志显示，您的tensorflow session 主机似乎没有其中的TPU设备。如果您正在Cloud TPU中运行作业，则可以执行

with tf.Session('<replace_with_your_worker_address>') as sess:
    print(sess.list_devices())

您应该至少看到类似"/<some_thing_varies_in_your_env>/device:TPU:0"的设备。