Блог

Я имею дело с pytorch в colab
Во время обучения, pytorch потребляет огромную память
после обучения, я сохранил модель и загрузил модель в другой ноутбук (примечание 2).
в примечании 2, после загрузки state_dict и всего остального, pytorch потребляет намного меньше памяти, чем в состоянии обучения.
Итак, мне интересно, что «бесполезные» данные хранятся в памяти графической карты во время обучения (в моем случае около 13 ГБ) …
Если да, то как мне удалить бесполезные данные после обучения?

plus. Я пытался удалить переменные, используемые во время обучения, но они были недостаточно большими (около 2 ГБ)

Этого следует ожидать во время обучения. В процессе обучения сами операции будут занимать память.

Например, рассмотрим следующую операцию —

 a = np.random.rand(100, 500, 300)
b = np.random.rand(200, 500, 300)
c = (a[:, None, :, :] * b[None, :, :, :]).sum(-1).sum(-1)
 

Объем памяти a, b и c по отдельности составляет около 400 МБ. Однако, если вы проверите

 %memit (a[:, None, :, :] * b[None, :, :, :]).sum(-1).sum(-1)
 

Это 23 ГБ! Сама строка занимает много памяти для фактического выполнения операции, потому что задействованы массивные промежуточные массивы. Эти массивы являются временными и автоматически удаляются после завершения операции. Таким образом, удаление некоторых переменных не сильно повлияет на уменьшение занимаемой площади.

Способ обойти это — использовать операции, оптимизированные для памяти.

Например, выполнение np.tensordot(a, b, ((1, 2), (1, 2))) вместо умножения путем трансляции оставляет гораздо больший объем памяти.

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