Серверы параметров Tensorflow на SageMaker

Question

Я пытаюсь понять, как работают серверы параметров (PS) для распределенного обучения в Tensorflow на Amazon SageMaker.

Чтобы сделать вещи более конкретными, я могу запустить пример из AWS с помощью PS: 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. В примере от 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, но этот флаг установлен в False для PS). Каким-то образом Tensorflow или контейнер SageMaker могут решить, где следует хранить переменные для каждого слоя. Однако я не вижу в коде контейнера ничего, что могло бы сделать что-либо со стратегией распространения.

Я могу запустить этот код и обучить модель, используя 1 и 2 экземпляра. Когда я это делаю, я вижу сокращение времени выполнения почти на 50%, что говорит о распределенном обучении.

Мои вопросы:

1. Как Tensorflow определяет распределение переменных в PS? В примере кода нет явной ссылки на устройства. Как-то раздача происходит автоматически.
2. Можно ли увидеть, какие параметры были назначены каждому PS? Или посмотреть, как выглядит связь между PS? Если да, то как?

Answer 1

Мои вопросы:

Как Tensorflow определяет распределение переменных в PS? В примере кода нет явной ссылки на устройства. Как-то раздача происходит автоматически.

Изображение TensorFlow, предоставленное SageMaker, содержит код для настройки TF_CONFIG и запуска сервера параметров для многопользовательского обучения. См. Код [здесь][1] Настройка такова, что для каждого узла в кластере настроен PS и рабочий поток.

Он не использует никакой DistributionStrategy, поэтому используется стратегия по умолчанию. См. [Здесь][2].

Если вы хотите использовать другую DistributionStrategy или другой TF_CONFIG, вам необходимо отключить parameter_server при запуске учебного задания SageMaker и настройте все в своем учебном сценарии.

Можно ли увидеть, какие параметры были назначены каждому PS? Или посмотреть, как выглядит связь между PS? Если да, то как?

Вы сможете получить некоторую информацию из выходного журнала, который можно найти в CloudWatch. Ссылка доступна на странице консоли Training Job. [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