Реализация обучения смешанной точности - PullRequest
Новая
       14

Реализация обучения смешанной точности

0 голосов
/ 05 августа 2020

Итак, с выпуском 1.16 я попытался оптимизировать часть своего кода для обеспечения смешанной точности и посмотреть, увижу ли я измеримое ускорение. rtx 2060), но я не измеряю скорость, поэтому я хотел бы показать их здесь, чтобы сначала проверить правильность моей реализации. Моя реализация основана на примере кода https://pytorch.org/blog/accelerating-training-on-nvidia-gpus-with-pytorch-automatic-mixed-precision/

Функция поезда БЕЗ смешанной точности: 

def train(model, device, train_loader, criterion, optimizer, scheduler, epoch, iter_meter, writer):
    model.train()
    data_len = len(train_loader.dataset)
    train_start_time = time.time()
    for batch_idx, _data in enumerate(train_loader):
        spectrograms, labels, input_lengths, label_lengths = _data
        spectrograms, labels = spectrograms.to(device), labels.to(device)

        optimizer.zero_grad()

        output = model(spectrograms)  # (batch, time, n_class)
        output = F.log_softmax(output, dim=2)
        output = output.transpose(0, 1)  # (time, batch, n_class)

        loss = criterion(output, labels, input_lengths, label_lengths)
        loss.backward()

        writer.add_scalar("Loss/train", loss.item(), iter_meter.get())
        writer.add_scalar("learning_rate", scheduler.get_last_lr()[0], iter_meter.get())
        optimizer.step()
        scheduler.step()
        iter_meter.step()
        
    return loss.item()

И функция поезда со смешанной точностью 

def train(model, device, train_loader, criterion, optimizer, scheduler, epoch, iter_meter, scaler, writer):
    model.train()
    data_len = len(train_loader.dataset)
    train_start_time = time.time()
    for batch_idx, _data in enumerate(train_loader):

        spectrograms, labels, input_lengths, label_lengths = _data
        spectrograms, labels = spectrograms.to(device), labels.to(device)

        optimizer.zero_grad()
        with torch.cuda.amp.autocast():
            output = model(spectrograms)  # (batch, time, n_class)
            output = F.log_softmax(output, dim=2)
            output = output.transpose(0, 1)  # (time, batch, n_class)

            loss = criterion(output, labels, input_lengths, label_lengths)

        # Mixed precision
        scaler.scale(loss).backward()  # loss.backward()

        scaler.step(optimizer)  # optimizer.step()
        scheduler.step() #Should I also put this steps in the scaler?
        iter_meter.step()

        # Updates the scale for next iteration
        scaler.update()

        writer.add_scalar("Loss/train", loss.item(), iter_meter.get())
        writer.add_scalar("learning_rate", scheduler.get_last_lr()[0], iter_meter.get())
    return loss.item()

Основное отличие состоит в том, что я передаю входные данные модели и вычисляю потери внутри with torch.cuda.amp.autocast(): и в обратном порядке с помощью скейлера.

Время для обучения с использованием обеих функций занимает одинаковое время, мои мысли о возможных причинах:

  • Automati c обучение смешанной точности из pytorch (делая обе функции кода равными)
  • Я что-то пропустил в своей реализации
  • Моя модель не может получить выгоду от ускорения от обучения смешанной точности

На случай, если кто-то захочет проверить исходную обучаемую модель, у меня она есть на Github

Есть ли у вас опытные глаза, которые могут на него взглянуть?

...