TF2.x Eager 模式现在不能支持 ParameterServerStrategy 了吗？

Question

TF 版本：最新主控，b083cea

下面是一个使用 TF2.0 eager 模式的简单示例，使用 MirroredStrategy 运行成功，但使用 ParameterServerStrategy 出错。

TF2.0 eager 模式现在可以支持 ParameterServerStrategy 了吗？到目前为止我没有找到一个成功的例子:(

from __future__ import absolute_import, division, print_function, unicode_literals
import tensorflow as tf
import tensorflow_datasets as tfds
import os, json

datasets, info = tfds.load(name='mnist', with_info=True, as_supervised=True)
mnist_train, mnist_test = datasets['train'], datasets['test']

os.environ['TF_CONFIG'] = json.dumps({
    "cluster": {
        "worker": ["localhost:12345"],
        "ps": ["localhost:12346"]
    },
    "task": {"type": "worker", "index": 0}
})

strategy = tf.distribute.experimental.ParameterServerStrategy()
#strategy = tf.distribute.MirroredStrategy()

print('Number of devices: {}'.format(strategy.num_replicas_in_sync))
num_train_examples = info.splits['train'].num_examples
num_test_examples = info.splits['test'].num_examples
BUFFER_SIZE = 10000
BATCH_SIZE_PER_REPLICA = 64
BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync

def scale(image, label):
  image = tf.cast(image, tf.float32)
  image /= 255
  return image, label

train_dataset = mnist_train.map(scale).shuffle(BUFFER_SIZE).batch(BATCH_SIZE)
eval_dataset = mnist_test.map(scale).batch(BATCH_SIZE)

with strategy.scope():
  model = tf.keras.Sequential([
      tf.keras.layers.Conv2D(32, 3, activation='relu', input_shape=(28, 28, 1)),
      tf.keras.layers.MaxPooling2D(),
      tf.keras.layers.Flatten(),
      tf.keras.layers.Dense(64, activation='relu'),
      tf.keras.layers.Dense(10, activation='softmax')
  ])

  model.compile(loss='sparse_categorical_crossentropy',
                optimizer=tf.keras.optimizers.Adam(),
                metrics=['accuracy'])

checkpoint_dir = './training_checkpoints'
# Name of the checkpoint files
checkpoint_prefix = os.path.join(checkpoint_dir, "ckpt_{epoch}")

# Function for decaying the learning rate.
# You can define any decay function you need.
def decay(epoch):
  if epoch < 3:
    return 1e-3
  elif epoch >= 3 and epoch < 7:
    return 1e-4
  else:
    return 1e-5

# Callback for printing the LR at the end of each epoch.
class PrintLR(tf.keras.callbacks.Callback):
  def on_epoch_end(self, epoch, logs=None):
    print('\nLearning rate for epoch {} is {}'.format(epoch + 1,
                                                      model.optimizer.lr.numpy()))
callbacks = [
    tf.keras.callbacks.TensorBoard(log_dir='./logs'),
    tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_prefix,
                                       save_weights_only=True),
    tf.keras.callbacks.LearningRateScheduler(decay),
    PrintLR()
]

model.fit(train_dataset, epochs=12, callbacks=callbacks)
model.load_weights(tf.train.latest_checkpoint(checkpoint_dir))
eval_loss, eval_acc = model.evaluate(eval_dataset)
print('Eval loss: {}, Eval Accuracy: {}'.format(eval_loss, eval_acc))

错误信息

tf.keras.layers.Dense(10, 激活='softmax') _method_wrapper 中的文件“/usr/local/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py”，第 456 行 结果 = 方法（自我，*args，**kwargs） 文件“/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/engine/sequential.py”，第 116 行，在 __init__ super(Sequential, self).__init__(name=name, autocast=False) 文件“/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py”，第 199 行，在 __init__ self._init_batch_counters() _method_wrapper 中的文件“/usr/local/lib/python3.7/site-packages/tensorflow/python/training/tracking/base.py”，第 456 行 结果 = 方法（自我，*args，**kwargs） _init_batch_counters 中的文件“/usr/local/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py”，第 206 行 self._train_counter = variables.Variable(0, dtype='int64', aggregation=agg) __call__ 中的文件“/usr/local/lib/python3.7/site-packages/tensorflow/python/ops/variables.py”，第 261 行 返回 cls._variable_v2_call(*args, **kwargs) _variable_v2_call 中的文件“/usr/local/lib/python3.7/site-packages/tensorflow/python/ops/variables.py”，第 255 行 形状=形状） 文件“/usr/local/lib/python3.7/site-packages/tensorflow/python/ops/variables.py”，第 66 行，在 getter 返回 capture_getter(captured_previous, **kwargs) 文件“/usr/local/lib/python3.7/site-packages/tensorflow/python/distribute/distribute_lib.py”，第 1769 行，在 creator_with_resource_vars return self._create_variable(next_creator, **kwargs) _create_variable 中的文件“/usr/local/lib/python3.7/site-packages/tensorflow/python/distribute/parameter_server_strategy.py”，第 455 行 使用 ops.device（self._variable_device）： 文件“/usr/local/lib/python3.7/site-packages/tensorflow/python/framework/ops.py”，第 5183 行，在设备中 “tf.device 不支持急切执行时的功能” RuntimeError: 启用急切执行时，tf.device 不支持函数。

Answer 1

在tf.distribute.experimental.ParameterServerStrategy 页面中，如下所述

注意：此策略仅适用于 Estimator API。创建 RunConfig 时，将此策略的实例传递给experimental_distribute 参数。然后应该将此 RunConfig 实例传递给调用 train_and_evaluate 的 Estimator 实例。

下面是一个使用tf.distribute.experimental.ParameterServerStrategy()的例子-

strategy = tf.distribute.experimental.ParameterServerStrategy()
run_config = tf.estimator.RunConfig(
    experimental_distribute.train_distribute=strategy)
estimator = tf.estimator.Estimator(config=run_config)
tf.estimator.train_and_evaluate(estimator,...)

另外如果你去页面Distributed training with TensorFlow，它解释了目前TF 2.0在哪些场景中支持什么，

希望这能回答您的问题。快乐学习。