Me gustaría saber las posibles formas de implementar capas de normalización por lotes con la sincronización de estadísticas de lotes cuando se entrena con múltiples GPU.

Cafe Tal vez hay algunas variantes de cafeína que podrían hacer, comoenlazar. Pero para la capa BN, entiendo que todavía sincroniza solo las salidas de las capas, no los medios y variables. Quizás MPI puede sincronizar medios y variables, pero creo que MPI es un poco difícil de implementar.

Antorcha He visto algunos comentariosaquí yaquí, que muestran que running_mean y running_var se pueden sincronizar, pero creo que la media del lote y el lote var no pueden o son difíciles de sincronizar.

Tensorflow Normalmente, es lo mismo que caffe y antorcha. La implementación de BN se refiereesta. Sé que tensorflow puede distribuir una operación a cualquier dispositivo especificado portf.device(). Pero el cálculo de medias y variables está en el medio de la capa BN, por lo que si reúno las medias y variables en la CPU, mi código será así:

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)

Eso es solo para una capa BN. Para recopilar estadísticas en la CPU, tengo que descifrar el código. Si tengo más de 100 capas BN, será engorroso.

No soy experto en esas bibliotecas, así que tal vez haya algunos malentendidos, no dude en señalar mis errores.

No me importa mucho la velocidad de entrenamiento. Estoy haciendo una segmentación de imagen que consume mucha memoria de GPU y BN necesita un tamaño de lote razonable (por ejemplo, más grande que 16) para estadísticas estables. Por lo tanto, usar múltiples GPU es inevitable. En mi opinión, tensorflow podría ser la mejor opción, pero no puedo resolver el problema del código de ruptura. La solución con otras bibliotecas también será bienvenida.