SageMaker 上的 Tensorflow 参数服务器

Question

我想了解参数服务器 (PS's) 如何在 Amazon SageMaker 上的 Tensorflow 中进行分布式训练。

为了使事情更具体，我可以运行使用 PS 的 AWS 示例：https://github.com/aws-samples/amazon-sagemaker-script-mode/blob/master/tf-distribution-options/tf-distributed-training.ipynb

这是初始化 Tensorflow 估计器的代码块：

from sagemaker.tensorflow import TensorFlow

git_config = {'repo': 'https://github.com/aws-samples/amazon-sagemaker-script-mode', 'branch': 'master'}

ps_instance_type = 'ml.p3.2xlarge'
ps_instance_count = 2

model_dir = "/opt/ml/model"

distributions = {'parameter_server': {
                    'enabled': True}
                }
hyperparameters = {'epochs': 60, 'batch-size' : 256}

estimator_ps = TensorFlow(
                       git_config=git_config,
                       source_dir='tf-distribution-options/code',
                       entry_point='train_ps.py', 
                       base_job_name='ps-cifar10-tf',
                       role=role,
                       framework_version='1.13',
                       py_version='py3',
                       hyperparameters=hyperparameters,
                       train_instance_count=ps_instance_count, 
                       train_instance_type=ps_instance_type,
                       model_dir=model_dir,
                       tags = [{'Key' : 'Project', 'Value' : 'cifar10'},{'Key' : 'TensorBoard', 'Value' : 'dist'}],
                       distributions=distributions)

查看 Tensorflow 的文档，似乎可以使用设备范围将变量分配给特定的工作人员。但是，在 SageMaker 上运行ning 训练作业时，我从未见过这样做。在来自 AWS 的示例中，模型定义为：

https://github.com/aws-samples/amazon-sagemaker-script-mode/blob/master/tf-distribution-options/code/model_def.py

这是一个片段：

def get_model(learning_rate, weight_decay, optimizer, momentum, size, mpi=False, hvd=False):

    model = Sequential()
    model.add(Conv2D(32, (3, 3), padding='same', input_shape=(HEIGHT, WIDTH, DEPTH)))
    model.add(BatchNormalization())
    model.add(Activation('relu'))
    model.add(Conv2D(32, (3, 3)))

    ...

    model.add(Flatten())
    model.add(Dense(512))
    model.add(Activation('relu'))
    model.add(Dropout(0.5))
    model.add(Dense(NUM_CLASSES))
    model.add(Activation('softmax'))

    if mpi:
        size = hvd.size()

    if optimizer.lower() == 'sgd':
        ...

    if mpi:
        opt = hvd.DistributedOptimizer(opt)

    model.compile(loss='categorical_crossentropy',
                  optimizer=opt,
                  metrics=['accuracy'])

    return model

这里没有对分发策略的引用（MPI 除外，但 PS 的该标志设置为 False）。不知何故，Tensorflow 或 SageMaker 容器能够决定每个层的变量应该存储在哪里。但是，我没有在容器代码中看到任何与分发策略相关的内容。

我能够运行此代码并使用 1 和 2 个实例训练模型。当我这样做时，我看到运行时间减少了近 50%，这表明正在进行分布式训练。

我的问题是：

Tensorflow 如何决定变量在 PS 上的分布？在示例代码中，没有明确引用设备。以某种方式自动完成分发。
是否可以查看每个 PS 分配了哪些参数？或者看看 PS 之间的通信是什么样的？如果是，怎么做？

Answer 1

My questions are:

How does Tensorflow decide the distribution of variables on the PS's? In the example code, there is no explicit reference to devices. Somehow the distribution is done automatically.

SageMaker 提供的 TensorFlow 映像具有设置 TF_CONFIG 和启动参数服务器以进行多任务训练的代码。请参阅代码 [此处][1] 该设置是为集群中的每个节点配置一个 PS 和一个工作线程。

它没有使用任何 DistributionStrategy，因此使用默认策略。请参阅[此处][2]。

如果您想使用不同的 DistributionStrategy 或不同的 TF_CONFIG，您需要在启动 SageMaker 训练作业时禁用 parameter_server 选项，并在训练脚本中设置所有内容。

Is it possible to see which parameters have been assigned to each PS? Or to see what the communication between PS's looks like? If so, how?

您应该能够从可在 CloudWatch 中找到的输出日志中获取一些信息。 link 在培训作业控制台页面上可用。 [1]: https://github.com/aws/sagemaker-tensorflow-training-toolkit/blob/master/src/sagemaker_tensorflow_container/training.py#L37 [2]: https://www.tensorflow.org/guide/distributed_training#default_strategy

SageMaker 上的 Tensorflow 参数服务器

Tensorflow Parameter Servers on SageMaker

tensorflow

amazon-sagemaker

distributed-tensorflow