Как использовать модуль Pytorch в качестве экстрактора только в прямом проходе? - PullRequest
Новая
       24

Как использовать модуль Pytorch в качестве экстрактора только в прямом проходе?

0 голосов
/ 29 октября 2018

Я нашел строку кода в этом файле : 

out_a, out_p, out_n = model(data_a), model(data_p), model(data_n)

Как вы можете видеть, функция forward() вызывается несколько раз до backward(). В моем тесте потребление графического процессора соответственно увеличится, поэтому происходит утечка памяти графического процессора. Вопросы здесь:

  1. Как решить эту проблему?
  2. Возможно ли использовать модель в качестве чистого экстрактора признаков для этих трех последовательных вызовов (с некоторой модификацией, конечно), а затем при четвертом вызове использовать ее как forward(), если "Да", как ее реализовать?

Ответы [ 2 ]

0 голосов
/ 31 октября 2018

В качестве альтернативы уже предложенному ответу вы также можете локально отключить вычисление градиента, используя torch.no_grad(): 

with torch.no_grad():
    out_a, out_p, out_n = model(data_a), model(data_p), model(data_n)

Это также предотвращает любые вычисления градиента в операторе with и, таким образом, значительно сокращает использование памяти:

Контекстный менеджер, который отключил расчет градиента.

Отключение вычисления градиента полезно для вывода, когда вы уверен, что ты не будешь звонить Tensor.backward(). Это уменьшит память потребление для вычислений, которые в противном случае имели бы requires_grad = True . В этом режиме результат каждого вычисления будет иметь require_grad = False , даже если входы имеют requires_grad = True .

Источник: https://pytorch.org/docs/stable/autograd.html#torch.autograd.no_grad

0 голосов
/ 31 октября 2018

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

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

model = FeatureExtractor().cuda()
mode.eval()

Далее, чтобы уменьшить потребление памяти графическим процессором, необходимо остановить накопление градиента и вычисление. 

for p in model.parameters():
    p.require_grads = False
out_a, out_p, out_n = model(data_a), model(data_p), model(data_n)

Тогда, я думаю, ваша проблема будет решена. Обратите внимание, что градиента «нет» можно пропустить через out_a, out_p и out_n.

...