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

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

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

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

В 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 реплики, но я не знаю, разделяется ли первая реплика на обоих графических процессорах в соответствии с назначением, или каждая реплика принудительно используется только для одного графического процессора. Любое понимание будет приятно.

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

0 ответов

Другие вопросы по тегам