Sposoby implementacji warstw BN wielu GPU ze środkami synchronizującymi i zmiennymi

Question

Chciałbym poznać możliwe sposoby implementacji warstw normalizacji wsadowej z synchronizowaniem statystyk wsadowych podczas treningu z wieloma GPU.

Caffe Być może jest kilka wariantów caffe, które mogą zrobić, jak link. Ale jeśli chodzi o warstwę BN, rozumiem, że nadal synchronizuje ona tylko wyjścia warstw, a nie środki i warianty. Może MPI może zsynchronizować środki i vary, ale myślę, że MPI jest trochę trudny do zrealizowania.

Torch Widziałem kilka uwag here i here, które pokazują running_mean i running_var mogą być synchronizowane, ale myślę, że partia myśli i partii var nie mogą lub są trudne do synchronizacji.

Tensorflow Zwykle jest to to samo, co caffe and torch. Implementacja BN dotyczy this. Wiem, że tensorflow może dystrybuować operację do dowolnego urządzenia określonego przez tf.device(). Ale obliczanie środków i vars jest w środku warstwy BN, więc jeśli zebrać środki i vars w CPU, mój kod będzie tak:

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) 

To jest tylko dla jednej BN warstwy. Do zbierania statystyk w procesorze, muszę złamać kod. Jeśli mam więcej niż 100 warstw BN, będzie to uciążliwe.

Nie jestem ekspertem w tych bibliotekach, więc może są pewne nieporozumienia, proszę wskazać moje błędy.

Nie zależy mi zbytnio na prędkości treningu. Robię segmentację obrazu, która zużywa dużo pamięci GPU, a BN potrzebuje rozsądnej wielkości partii (na przykład większej niż 16) dla stabilnych statystyk. Więc korzystanie z wielu GPU jest nieuniknione. Moim zdaniem, tensorflow może być najlepszym wyborem, ale nie mogę rozwiązać problemu z łamaniem kodu. Rozwiązanie z innymi bibliotekami również będzie mile widziane.

Answer 1

Nie jestem pewien, czy w pełni rozumiem twoje pytanie, ale pod warunkiem, że prawidłowo skonfigurowałeś swój zakres zmiennych, kolekcja tf.GraphKeys.UPDATE_OPS powinna automatycznie mieć aktualizację ops dla batch_norm dla każdej z twoich wież. Jeśli wszystkie parametry update_ops zostaną zastosowane synchronicznie, zostaną niejawnie uśrednione przez serwer parametrów, wszystko co musisz zrobić, to upewnić się, że aktualizacje są stosowane przed średnią i zastosować gradienty. (Jeśli dobrze rozumiem twoje zamiary).

Ze względu na zakres zmienny, każdy zestaw aktualizacji ops aktualizuje te same zmienne, więc aby zsynchronizować aktualizację, wszystko, co musisz zrobić, to ustawić obliczenia gradientowe na pełnym zestawie aktualizacji. Powinieneś również hermetyzować wszystkie warstwy norm wsadowych w jednym numerze name_scope, aby uniknąć przechwycenia jakichkolwiek dodatkowych operacji w UPDATE_OPS. Kod szkielet poniżej:

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. 

Jeśli chcesz spróbować to na jakimś istniejącym kodem, spróbuj po prostu usuwając linię if i == 0 tutaj: https://github.com/tensorflow/models/blob/master/tutorials/image/cifar10_estimator/cifar10_main.py#L115

Jedziesz zobaczyć niektóre spowolnienie (zazwyczaj używamy tylko jedną wieżę do obliczania statystyk norm wsadowych z tego powodu), ale powinno robić to, co chcesz.