Операции на месте факела для экономии памяти (softmax) - PullRequest
Новая
       28

Операции на месте факела для экономии памяти (softmax)

0 голосов
/ 11 декабря 2018

Некоторые операции в факеле выполняются на месте.Сокращенные операторы, такие как + =, например.

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

Я сейчас работаюс языковой обработкой.Модель создает длинную последовательность распределений вероятностей по большому словарному запасу.Этот конечный тензор вывода отвечает за около 60% выделенной памяти.Это огромная проблема, так как мне нужно вычислить softmax по нему, и это удваивает требуемую память.

Вот пример проблемы.Меня не интересует тензор t, только его softmax: 

import numpy as np
import torch
import torch.nn.functional as F

t = torch.tensor(np.zeros((30000,30000))).cuda()  #allocates 6.71 GB of GPU
softmax = F.softmax(t, 1)  #out of memory error
del t  #too late, program crashed

Даже следующее не работает: 

F.softmax(torch.tensor(np.zeros((30000,30000))).cuda(), 1)

Ответы [ 2 ]

0 голосов
/ 12 декабря 2018

Я создал версию softmax на месте: 

import numpy as np
import torch
import torch.nn.functional as F

# in-place version
t = torch.tensor(np.ones((100,200)))
torch.exp(t, out=t)
summed = torch.sum(t, dim=1, keepdim=True)
t /= summed

# original version
t2 = torch.tensor(np.ones((100,200)))
softmax = F.softmax(t2, 1)

assert torch.allclose(t, softmax)

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

  • многим функциям, таким как torch.exp, может быть предоставлен необязательный параметр out.
  • назначения t[idx] = something на месте
  • стенографические операторы /=, *=, +=, -= на месте

Это требует тщательной отладки и может быть не интуитивно понятным: 

t = t / summed  #not in-place
t /= summed  #in-place

Я читал, что операции на месте могут вызвать проблемы с градиентами.Я проведу еще несколько тестов с этим кодом.

0 голосов
/ 12 декабря 2018

Это невозможно с PyTorch на текущий момент.Вы можете попробовать развернуть свое собственное ядро ​​графического процессора , но я вижу проблемы (если не стену) впереди, что, вероятно, является причиной, почему эта операция вообще не доступна.

Softmax может быть легко применен параллельно, за исключением нормализации, которая требует сокращения.Сокращения нетривиальны, и они могут быть параллельными или (имея в виду, что использование атомных элементов равнозначно параллельной, но не параллельной операции).Это означает, что ваша операция на месте в любом случае должна была бы либо распределиться под капотом, что противоречит цели, либо работать ужасно медленно.Пожалуйста, сочтите мой ответ несколько спекулятивным, я не эксперт по GPGPU, но я хочу сказать, что, по крайней мере, трудно решить проблему дешево, быстро и правильно.насчет софтмакса, мне кажется, что вы делаете вывод.И, возможно, для вашего приложения целесообразно переместить логиты в ЦП и запустить там softmax, пока ваш графический процессор уже обрабатывает следующий пакет?

...