эффективное обучение трансформаторам без принуждения учителя

#python #pytorch #time-series #transformer

Вопрос:

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

 # src shape is [n_batch, src_length, d_model] --> [N, S, D]
# tgt and gt (groundtruth) shape is [n_batch, tgt_length, d_model] --> [N, T, D]

l = 0.
for src, tgt, gt in loader:
    src, tgt, gt = src.to(device), tgt.to(device), gt.to(device)
    optimizer.zero_grad()

    # first tgt element is the last src one, followed by a zero-filled tensor
    tgt = torch.cat((src[:, -1:, :], torch.zeros(tgt.shape[0], tgt.shape[1] - 1, tgt.shape[2], device=device)), dim=1)

    out = torch.zeros_like(tgt, device=device)
    memory = model.encode(src)
    for j in range(gt.shape[1]):
        out = model.decode(tgt, memory)
        l  = loss(out[:, j, :], gt[:, j, :])
        tgt = tgt.detach()
        tgt = torch.cat((tgt[:, :j   1, :], out[:, j:j 1, :], tgt[:, j   1:gt.shape[1] - 1, :]), dim=1) 
    l.backward()
 

Я начинаю инициализацию цели с последнего элемента исходной последовательности, за которым следует столько нулей, сколько необходимо для соответствия правильной форме.
Затем, внутри внутреннего for, я добавляю новое прогнозируемое значение для предыдущей цели (снова добавляя необходимые нули), чтобы создать цель для следующей итерации.

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

Теперь приведенный выше код отлично работает для чрезвычайно небольших партий (десятки образцов), несмотря на то, что время для одной эпохи обучения увеличивается в 10 раз. обучение той же модели с использованием принудительного обучения. Проблема в том, что если я попытаюсь использовать «обычный» размер пакета, скажем, 32, 64 образца, код выйдет из строя во время обучения, потому что у CUDA закончится память. Форма для каждого образца моего набора данных равна [12, 5], которая становится [12, 128] после встраивания в модель.

Мой графический процессор локальной машины имеет только 4 ГБ оперативной памяти, но то же самое поведение происходит при использовании 12 Гб графических процессоров Colab. Я почти уверен, что есть лучший способ реализовать обучение трансформаторам без принуждения учителей.

 RuntimeError: CUDA out of memory. Tried to allocate 2.00 MiB (GPU 0; 11.17 GiB total capacity; 10.38 GiB already allocated; 1.81 MiB free; 10.67 GiB reserved in total by PyTorch)
 

Ошибка выше обычно проверяется при вызове model.encode(src) или model.decode(tgt, память). С другой стороны, использование учителя, заставляющего код, работает нормально.

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