Переписывают ли стратегии распределения тензорного потока назначения устройств? - PullRequest
Новая
       90

Переписывают ли стратегии распределения тензорного потока назначения устройств?

0 голосов
/ 07 марта 2020

Я использую оценщик API. Я настроил обучающий конвейер, но я хочу реализовать различные стратегии распределения, когда доступно несколько графических процессоров. Я хочу иметь возможность выбирать между параллелизм данных и параллелизм модели перед началом обучения. То, что я сделал до сих пор, выглядит примерно так: 

    if args.nodist: # True if no data parallelism distribution selected
        strategy = None
    else:
        strategy = tf.distribute.MirroredStrategy()

Затем я обычно передаю стратегию tf.estimator.RunConfig(train_distribute=strategy)

In model_fn(), которая схематически определяется как 

    def model_fn(features, labels, mode, params):
        loss, train_op, = None, None
        eval_metric_ops, training_hooks, evaluation_hooks = None, None, None
        predictions_dict = None
        output = model(input_tensor=features['image'], params=params)
        with tf.device('/CPU:0'):
            with tf.name_scope('arg_max_output'):
                output_arg = tf.math.argmax(output, axis=-1)
        if mode in (estimator.ModeKeys.TRAIN, estimator.ModeKeys.EVAL):
            with tf.name_scope('Loss_Calculation'):
                with tf.device('/GPU:0'):
                    loss = custom_loss(predictions=output, labels=labels['label'])
            with tf.name_scope('Metrics_Calculations'):
                with tf.device('/GPU:0'):
                    metric_1 = tf..metrics.metric(labels=labels['label'], predictions=output)
                with tf.device('/GPU:1'):
                    metric_2 = tf..metrics.metric(labels=labels['label'], predictions=output)
            with tf.name_scope('Images'):
                with tf.name_scope('{}'.format(mode)):
                    with tf.device('/CPU:0'):
                        summary.image('1_Image', features['image'], max_outputs=1)
                        summary.image('2_Label', tf.expand_dims(tf.cast(label * 255, 
        type=tf.uint8), axis=-1), max_outputs=1)
        if mode == estimator.ModeKeys.TRAIN:
            with tf.device('/CPU:0'):
                with tf.name_scope('Metrics'):
                    summary.scalar('1_metric', metric_2[1])
                    summary.scalar('2_metric', metric_2[1])
                                       .
                                       .
                                       .
        return estimator.EstimatorSpec(mode,
                                       predictions=predictions_dict,
                                       loss=loss,
                                       train_op=train_op,
                                       eval_metric_ops=eval_metric_ops,
                                       training_hooks=training_hooks,
                                       evaluation_hooks=evaluation_hooks)

В вышеупомянутом model() есть что-то вроде: 

    def model(input_tensor, params):
        with variable_scope('Part_1'):
            with tf.device('/GPU:0'):
                output = layer_1(input_tensor=input_tensor)
                output = layer_2(input_tensor=output)
        with variable_scope('Part_2'):
            with tf.device('/GPU:1'):
                output = layer_3(input_tensor=output)
                output = layer_4(input_tensor=output)
        return output

Теперь возникает вопрос: если я включу MirroredStrategy, будет ли он конфликтовать с назначениями закодированного устройства или будет перезаписывать эти назначения для настройки Зеркальная стратегия распространения? Когда я не выбираю паралелизм данных, модель обычно разделяется на 2 графических процессора. На MirroredStrategy это правильно создает 2 реплики, но я не знаю, если первая реплика разделена на обоих графических процессорах, как это назначено, или если каждая реплика вынуждена только на один графический процессор. Любое понимание будет приятным.

Спасибо вам всем заранее!

...