Как очистить переменную, чтобы освободить оперативную память во время вывода?

Question

Я пытаюсь обучить unet на небольших изображениях и выполнить выводы для больших. моя идея состояла в том, чтобы тренироваться на моем GPU и сделать этап вывода на CPU. проблема в том, что мне не хватает оперативной памяти. во время прямого прохода я попытался удалить переменную, которая мне больше не нужна, но память не выделяется снова. Что я делаю неправильно? Вот мой прямой проход:

def forward(self, x):
    x1 = self.in_conv(x)
    del x
    x2 = self.down1(x1)
    x3 = self.down2(x2)
    x4 = self.down3(x3)
    x5 = self.down4(x4)
    x_n = self.up1(x5, x4)
    x_n = self.up2(x_n, x3)
    x_n = self.up3(x_n, x2)
    x_n = self.up4(x_n, x1)
    return self.out_conv(x_n)

первая строка

x1 = self.in_conv(x)

- та, которая занимает большую часть оперативной памяти. однако после этого мне не нужен x, а «del x» не помогает.

Answer 1

Вам нужно использовать torch.no_grad () при выводе.

Удаление x подобно тому, как вы это сделали, не освободит память, так как вызывающая сторона предположительно имеет ссылку на входной тензор.

Answer 2

Попробуйте использовать модуль g c ( сборщик мусора )

import gc
gc.collect()