Блог

В процессе отслеживания ошибки OOM GPU я сделал следующие контрольные точки в своем коде Pytorch (работает на Google Colab P100):

 learning_rate = 0.001
num_epochs = 50

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

print('check 1')
!nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

model = MyModel()

print('check 2')
!nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

model = model.to(device)

print('check 3')
!nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

print('check 4')
!nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

for epoch in range(num_epochs):
    train_running_loss = 0.0
    train_accuracy = 0.0

    model = model.train()

    print('check 5')
    !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

    ## training step
    for i, (name, output_array, input) in enumerate(trainloader):
        
        output_array = output_array.to(device)
        input = input.to(device)
        comb = torch.zeros(1,1,100,1632).to(device)

        print('check 6')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        ## forward   backprop   loss
        output = model(input, comb)

        print('check 7')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        loss = my_loss(output, output_array)

        print('check 8')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        optimizer.zero_grad()

        print('check 9')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        loss.backward()

        print('check 10')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        ## update model params
        optimizer.step()

        print('check 11')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        train_running_loss  = loss.detach().item()

        print('check 12')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        temp = get_accuracy(output, output_array)

        print('check 13')
        !nvidia-smi | grep MiB | awk '{print $9 $10 $11}'

        train_accuracy  = temp     
 

со следующим выводом:

 check 1
2MiB/16160MiB
check 2
2MiB/16160MiB
check 3
3769MiB/16160MiB
check 4
3769MiB/16160MiB
check 5
3769MiB/16160MiB
check 6
3847MiB/16160MiB
check 7
6725MiB/16160MiB
check 8
6725MiB/16160MiB
check 9
6725MiB/16160MiB
check 10
9761MiB/16160MiB
check 11
16053MiB/16160MiB
check 12
16053MiB/16160MiB
check 13
16053MiB/16160MiB
check 6
16053MiB/16160MiB
check 7
16071MiB/16160MiB
check 8
16071MiB/16160MiB
check 9
16071MiB/16160MiB
check 10
16071MiB/16160MiB
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-11-f566d09448f9> in <module>()
     65 
     66         ## update model params
---> 67         optimizer.step()
     68 
     69         print('check 11')

3 frames
/usr/local/lib/python3.7/dist-packages/torch/optim/optimizer.py in wrapper(*args, **kwargs)
     86                 profile_name = "Optimizer.step#{}.step".format(obj.__class__.__name__)
     87                 with torch.autograd.profiler.record_function(profile_name):
---> 88                     return func(*args, **kwargs)
     89             return wrapper
     90 

/usr/local/lib/python3.7/dist-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
     26         def decorate_context(*args, **kwargs):
     27             with self.__class__():
---> 28                 return func(*args, **kwargs)
     29         return cast(F, decorate_context)
     30 

/usr/local/lib/python3.7/dist-packages/torch/optim/adam.py in step(self, closure)
    116                    lr=group['lr'],
    117                    weight_decay=group['weight_decay'],
--> 118                    eps=group['eps'])
    119         return loss

/usr/local/lib/python3.7/dist-packages/torch/optim/_functional.py in adam(params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps, amsgrad, beta1, beta2, lr, weight_decay, eps)
     92             denom = (max_exp_avg_sqs[i].sqrt() / math.sqrt(bias_correction2)).add_(eps)
     93         else:
---> 94             denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)
     95 
     96         step_size = lr / bias_correction1

RuntimeError: CUDA out of memory. Tried to allocate 2.32 GiB (GPU 0; 15.78 GiB total capacity; 11.91 GiB already allocated; 182.75 MiB free; 14.26 GiB reserved in total by PyTorch)
 

Для меня имеет смысл, что model = model.to(device) создает 3,7 Г памяти.

Но почему запуск модели output = model(input, comb) создает еще один объем памяти 3G?
А затем loss.backward() создает еще один объем памяти 3G?
А затем optimizer.step() создает еще 6,3 Г памяти?

Я был бы признателен, если бы кто-нибудь объяснил, как работает модель распределения памяти GPU PyTorch в этом примере.

• Вывод

  По умолчанию вывод в вашей модели будет выделять память для хранения активаций каждого слоя (активация, как на входах промежуточного слоя). Это необходимо для обратного распространения, когда эти тензоры используются для вычисления градиентов. Простым, но эффективным примером является функция, определяемая f: x -> x² . Здесь df/dx = 2x , то есть для того , чтобы вычислять df/dx , вы должны хранить x в памяти.

  Если вы используете torch.no_grad() контекстный менеджер, вы позволите PyTorch не сохранять эти значения, тем самым экономя память. Это особенно полезно при оценке или тестировании вашей модели, т. е. при выполнении обратного распространения. Конечно, вы не сможете использовать это во время тренировок!

• Обратное распространение

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

• Обновление параметров модели

  Поскольку вы используете оптимизатор с отслеживанием состояния (Adam), для сохранения некоторых параметров требуется дополнительная память. Прочитайте соответствующий пост на форуме PyTorch об этом. Если вы попытаетесь использовать оптимизатор без состояния (например, SGD), у вас не должно быть никаких накладных расходов на память step при вызове.

Все три шага могут иметь потребности в памяти. Таким образом, объем памяти, выделенной на вашем устройстве, будет эффективно зависеть от трех элементов:

1. Размер вашей нейронной сети: чем больше модель, тем больше активаций слоев и градиентов будет сохранено в памяти.
2. Независимо от того, находитесь ли вы в torch.no_grad контексте: в этом случае в памяти должно быть только состояние вашей модели (никаких активаций или градиентов не требуется).
3. Тип используемого оптимизатора: с сохранением состояния (сохраняет некоторые текущие оценки во время обновления параметров) или без состояния (не требует этого).

требуется ли вам сделать это обратно