Способы реализации многоядерных блоков BN с синхронизирующими средствами и vars

Я хотел бы знать возможные способы реализации уровней пакетной нормализации с синхронизацией статистики партии при обучении с несколькими GPU.

Caffe Возможно, есть варианты, которые могли бы сделать, например ссылка. Но для слоя BN я понимаю, что он все еще синхронизирует только выходы слоев, а не средства и vars. Возможно, MPI может синхронизировать средства и vars, но я думаю, что MPI немного сложно реализовать.

Факел Я видел несколько комментариев здесь и здесь, которые показывают, что run_mean и run_var могут быть синхронизированы, но я думаю, что пакетное среднее и пакетное var не могут или не могут быть синхронизированы.

Tensorflow Обычно это то же самое, что и кофе и факел. Реализация BN означает this. Я знаю, что tensorflow может распространять операцию на любое устройство, указанное tf.device(). Но вычисление средств и vars находится в середине слоя BN, поэтому, если я соберу средства и vars в CPU, мой код будет таким:

cpu_gather = []
label_batches = []
for i in range(num_gpu):
    with tf.device('/gpu:%d' % i):
        with tf.variable_scope('block1', reuse=i > 0):
            image_batch, label_batch = cifar_input.build_input('cifar10', train_data_path, batch_size, 'train')
            label_batches.append(label_batch)

            x = _conv('weights', image_batch, 3, 3, 16, _stride_arr(1))
            block1_gather.append(x)

with tf.device('/cpu:0'):
    print block1_gather[0].get_shape()
    x1 = tf.concat(block1_gather, 0)
    # print x1.get_shape()
    mean, variance = tf.nn.moments(x1, [0, 1, 2], name='moments')

for i in range(num_gpu):
    with tf.device('/gpu:%d' % i):
        with tf.variable_scope('block2', reuse=i > 0):
            shape = cpu_gather[i].get_shape().as_list()
            assert len(shape) in [2, 4]
            n_out = shape[-1]
            beta, gamma, moving_mean, moving_var = get_bn_variables(n_out, True, True)

            x = tf.nn.batch_normalization(
                cpu_gather[i], mean, variance, beta, gamma, 0.00001)

            x = _relu(x)

Это просто для одного слоя BN. Для сбора статистики в CPU мне нужно разбить код. Если у меня будет более 100 слоев BN, это будет громоздким.

Я не эксперт в этих библиотеках, поэтому, возможно, есть некоторые недоразумения, не стесняйтесь указывать на мои ошибки.

Я не очень-то люблю скорость обучения. Я занимаюсь сегментацией изображений, которая потребляет много памяти GPU, а BN требует разумного размера партии (например, больше 16) для стабильной статистики. Поэтому использование multi-GPU неизбежно. На мой взгляд, тензорный поток может быть лучшим выбором, но я не могу решить проблему с нарушением кода. Решение с другими библиотеками также будет приветствоваться.

8
28 марта '17 в 0:42
источник поделиться
1 ответ

Я не уверен, полностью ли я понимаю ваш вопрос, но при условии, что вы правильно настроили область переменной, в коллекции tf.GraphKeys.UPDATE_OPS должны автоматически появляться операции обновления для batch_norm для каждой из ваших башен. Если все update_ops применяются синхронно, они будут неявно усреднены сервером параметров, все, что вам нужно сделать, это убедиться, что обновления применяются до того, как вы усреднёте и примените градиенты. (Если я правильно понимаю ваши намерения).

Из-за переменной видимости каждый набор операций обновления будет обновлять одни и те же переменные, поэтому для синхронизации операций по обновлению все, что вам нужно сделать, это забросить ваш градиентный расчет на полный набор операций обновления. Вы должны также инкапсулировать все ваши уровни пакетной нормы в один name_scope, чтобы избежать захвата любых внешних операций в UPDATE_OPS. Скелет кода ниже:

update_ops = []
for i, device in enumerate(devices):
  with tf.variable_scope('foo', reuse=bool(i > 0)):
    with tf.name_scope('tower_%d' % i) as name_scope:
      with tf.device(device):
        # Put as many batch_norm layers as you want here
      update_ops.extend(tf.get_collection(tf.GraphKeys.UPDATE_OPS,
                                          name_scope))
# make gradient calculation ops here
with tf.device(averaging_device):
  with tf.control_dependencies(update_ops):
    # average and apply gradients.

Если вы хотите попробовать это на каком-то существующем коде, попробуйте просто удалить строку if i == 0 здесь: https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10_estimator/cifar10_main.py#L115

Вы увидите некоторое замедление (мы обычно используем только одну башню для вычисления статистики серийной нормы по этой причине), но она должна делать то, что вы хотите.

2
06 янв. '18 в 3:15
источник

Посмотрите другие вопросы по меткам или Задайте вопрос