Управляйте памятью по-разному на поездах и тестируйте время

Question

В настоящее время я пишу модель сегментации, основанную на U-net с pytorch, и хочу использовать нечто похожее на инвертированный остаток, представленный в mobilenet v2, для повышения скорости модели на процессоре. код pytorch для мобильного netv2

Тогда я понимаю, что модель использует намного больше памяти на фазе поезда и на фазе тестирования. В то время как модель должна использовать больше памяти на фазе поезда, потому что все средние ступенчатые тензоры (карты характеристик) сохраняются, и с разделительной сверткой для каждой операции «свертки» создается больше тензоров. Но во время выполнения для сохранения пропуска соединения необходимо сохранить только несколько тензоров последнего шага, а все остальные тензоры можно удалить после создания следующего шага. Эффективность памяти должна быть одинаковой для u-сети с нормальной сверткой и u-сети с отделимой сверткой на этапе тестирования.

Я новичок в pytorch, поэтому я не знаю, как писать код, который предотвращает ненужные затраты памяти на время тестирования. Поскольку pytorch связан с python. Я думаю, я могу вручную удалить все ненужные тензоры в функции forward с помощью del. Но я предполагаю, что если я просто удалю переменные в форвардной функции, это повлияет на этап обучения. Есть ли здесь более продвинутые функциональные возможности Pytorch, которые могут оптимизировать использование памяти на этапе тестирования с помощью «сетевого графика»? Мне также любопытно, будет ли тензорный поток решать эти проблемы автоматически, поскольку он имеет более абстрактную и сложную логику построения графиков.

Answer 1

После прочтения официального кода pytorch для resnet я понимаю, что не должен давать имя всем переменным. Я не должен писать:

conv1 = self.conv1(x)
conv2 = self.conv2(conv1)

Я должен просто написать:

out = self.conv1(x)
out = self.conv2(out)

Таким образом, ничто не относится к obj, соответствующему conv1 после его использования, и python может его очистить.

, поскольку между блоками есть остаточные соединения, мне нужна еще одна переменная python для ссылки на переменную: aka

out = self.conv1(x)
residual_connect = out
out = self.conv2(out)
out = conv1 + out

Но на этапе декодирования необходим только выход. Поэтому я удалил residual_connect в начале этапа декодирования.

del residual_connect

Это похоже на взлом, и я удивлен, что это не вызвало проблем на тренировочном этапе. Использование оперативной памяти для моей модели теперь значительно уменьшено, но я считаю, что здесь должен быть более элегантный способ решения проблемы.

Answer 2

Поскольку вы использовали torch.no_grad() во время тестирования, вы автоматически запрашиваете у Context-менеджера отключить вычисление градиента, что приводит к меньшему использованию памяти, чем при обучении.Однако я обнаружил, что кэширующий распределитель занимает много памяти, которая может высвобождаться после каждого обновления модели во время обучения, что экономит много памяти.

Итак, вы можете использовать функцию torch.cuda.empty_cache () , и я нашел ее действительно полезной в моем случае.Кроме того, чтение управления памятью может помочь вам узнать другие важные вещи об управлении памятью GPU в PyTorch.