SageMaker 上的 Tensorflow 參數服務器

Question

我試圖了解參數服務器 (PS) 如何在 Amazon SageMaker 上的 Tensorflow 中進行分布式訓練。

為了使事情更具體，我可以使用 PS 從 AWS 運行示例： https : //github.com/aws-samples/amazon-sagemaker-script-mode/blob/master/tf-distribution-options/tf-分布式訓練.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 上運行訓練作業時，我從未見過這樣做。 在 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%，這表明正在進行分布式訓練。

我的問題是：

1. Tensorflow 如何決定 PS 上變量的分布？ 在示例代碼中，沒有對設備的顯式引用。 不知何故，分配是自動完成的。
2. 是否可以查看每個 PS 分配了哪些參數？ 或者看看PS之間的通信是什么樣的？ 如果是這樣，如何？

Answer 1

我的問題是：

Tensorflow 如何決定 PS 上變量的分布？ 在示例代碼中，沒有對設備的顯式引用。 不知何故，分配是自動完成的。

SageMaker 提供的 TensorFlow 圖像具有設置 TF_CONFIG 和啟動參數服務器以進行多工作訓練的代碼。 參見代碼[此處][1] 設置是為集群中的每個節點配置一個 PS 和一個工作線程。

它沒有使用任何 DistributionStrategy，因此使用默認策略。 見[這里][2]。

如果您想使用不同的 DistributionStrategy 或不同的 TF_CONFIG，則需要在啟動 SageMaker 訓練作業時禁用parameter_server選項並在訓練腳本中設置所有內容。

是否可以查看每個 PS 分配了哪些參數？ 或者看看PS之間的通信是什么樣的？ 如果是這樣，如何？

您應該能夠從 CloudWatch 中的輸出日志中獲取一些信息。 該鏈接在培訓作業控制台頁面上可用。 [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