Как адаптировать размер партии GPU во время тренировки? - PullRequest
Новая
       36

Как адаптировать размер партии GPU во время тренировки?

1 голос
/ 21 ноября 2019

Мне показалось удивительным, что я не смог найти в Интернете никаких ресурсов о том, как динамически адаптировать размер пакета GPU без остановки обучения.

Идея заключается в следующем:

1) Пройти обучениескрипт, который (почти) не зависит от используемого графического процессора. Размер партии будет динамически изменяться без вмешательства пользователя или необходимости настройки.

2) Возможность определения желаемого размера тренировочной партии, даже если она слишком велика для самой большой известной графической карты.

Например, допустим, я хочу обучить модель с использованием пакета размером 4096 изображений, каждое изображение 1024x1024. Скажем также, что у меня есть доступ к серверу с разными графическими процессорами NVidea, но я не знаю, какой из них будет мне назначен заранее. (Или что все хотят использовать самый большой графический процессор, и что я долго жду, когда наступит мой срок).

Я хочу, чтобы мой учебный скрипт нашел максимальный размер пакета (скажем, это 32 изображенияна пакет GPU) и обновляйте оптимизатор только после обработки всех 4096 изображений (один пакет обучения = 128 пакетов GPU).

1 Ответ

1 голос
/ 21 ноября 2019

Существуют разные способы решения этой проблемы. Но если указать графический процессор, который может выполнять эту работу, или использовать несколько графических процессоров, не вариант, то удобно динамически адаптировать размер пакета графического процессора.

Я подготовил этот репозиторий с иллюстративным обучающим примером в pytorch (он должен работать аналогично в TensorFlow)

В приведенном ниже коде try / исключения используется для попыткиРазличные размеры партий GPU без остановки обучения. Когда партия становится слишком большой, она уменьшается, и адаптация отключается. Пожалуйста, проверьте репозиторий для деталей реализации и возможных исправлений ошибок.

Также реализована методика, называемая Batch Spoofing, которая выполняет несколько прямых проходов перед выполнением обратного распространения. В PyTorch требуется только замена optimizer.zero_grad (). 

import torch
import torchvision
import torch.optim as optim
import torch.nn as nn

# Example of how to use it with Pytorch
if __name__ == "__main__":

    # #############################################################
    # 1) Initialize the dataset, model, optimizer and loss as usual.
    # Initialize a fake dataset

    trainset = torchvision.datasets.FakeData(size=1_000_000,
                                             image_size=(3, 224, 224),
                                             num_classes=1000)

    # initialize the model, loss and SGD-based optimizer
    resnet = torchvision.models.resnet152(pretrained=True,
                                          progress=True)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(resnet.parameters(), lr=0.01)

    continue_training = True  # criteria to stop the training

    # #############################################################
    # 2) Set parameters for the adaptive batch size
    adapt = True  # while this is true, the algorithm will perform batch adaptation
    gpu_batch_size = 2  # initial gpu batch_size, it can be super small
    train_batch_size = 2048  # the train batch size of desire

    # Modified training loop to allow for adaptive batch size
    while continue_training:

        # #############################################################
        # 3) Initialize dataloader and batch spoofing parameter
        # Dataloader has to be reinicialized for each new batch size.
        trainloader = torch.utils.data.DataLoader(trainset,
                                                  batch_size=int(gpu_batch_size),
                                                  shuffle=True)

        # Number of repetitions for batch spoofing
        repeat = max(1, int(train_batch_size / gpu_batch_size))

        try:  # This will make sure that training is not halted when the batch size is too large

            # #############################################################
            # 4) Epoch loop with batch spoofing
            optimizer.zero_grad()  # done before training because of batch spoofing.

            for i, (x, y) in enumerate(trainloader):

                y_pred = resnet(x)
                loss = criterion(y_pred, y)
                loss.backward()

                # batch spoofing
                if not i % repeat:
                    optimizer.step()
                    optimizer.zero_grad()

                # #############################################################
                # 5) Adapt batch size while no RuntimeError is rased.
                # Increase batch size and get out of the loop
                if adapt:
                    gpu_batch_size *= 2
                    break

                # Stopping criteria for training
                if i > 100:
                    continue_training = False

        # #############################################################
        # 6) After the largest batch size is found, the training progresses with the fixed batch size.
        # CUDA out of memory is a RuntimeError, the moment we will get to it when our batch size is too large.
        except RuntimeError as run_error:
            gpu_batch_size /= 2  # resize the batch size for the biggest that works in memory
            adapt = False  # turn off the batch adaptation

            # Number of repetitions for batch spoofing
            repeat = max(1, int(train_batch_size / gpu_batch_size))

            # Manual check if the RuntimeError was caused by the CUDA or something else.
            print(f"---\nRuntimeError: \n{run_error}\n---\n Is it a cuda error?")

Если у вас есть код, который может работать аналогичным образом в Tensorflow, Caffe или других, пожалуйста, поделитесь!

...