Как освободить всю память pytorch берется из памяти GPU - PullRequest
Новая
       36

Как освободить всю память pytorch берется из памяти GPU

0 голосов
/ 06 сентября 2018

У меня есть какой-то код высокого уровня, поэтому обучение модели и т. Д. Обернуто классом pipeline_network.Моя главная цель - тренировать новую модель каждый новый фолд. 

for train_idx, valid_idx in cv.split(meta_train[DEPTH_COLUMN].values.reshape(-1)):

        meta_train_split, meta_valid_split = meta_train.iloc[train_idx], meta_train.iloc[valid_idx]

        pipeline_network = unet(config=CONFIG, suffix = 'fold' + str(fold), train_mode=True)

Но затем я перехожу ко 2-му фолду, все выходит из памяти GPU: 

RuntimeError: cuda runtime error (2) : out of memory at /pytorch/torch/lib/THC/generic/THCStorage.cu:58

В конце эпохиЯ попытался вручную удалить этот конвейер безуспешно: 

 def clean_object_from_memory(obj): #definition
    del obj
    gc.collect()
    torch.cuda.empty_cache()

clean_object_from_memory( clean_object_from_memory) # calling

Вызов это тоже не помог: 

def dump_tensors(gpu_only=True):
        torch.cuda.empty_cache()
        total_size = 0
        for obj in gc.get_objects():
            try:
                if torch.is_tensor(obj):
                    if not gpu_only or obj.is_cuda:
                        del obj
                        gc.collect()
                elif hasattr(obj, "data") and torch.is_tensor(obj.data):
                    if not gpu_only or obj.is_cuda:
                        del obj
                        gc.collect()
            except Exception as e:
                pass

Как можно сбросить pytorch, после чего я перехожу к следующему сгибу

1 Ответ

0 голосов
/ 06 мая 2019

Попробуйте удалить объект с помощью del, а затем примените torch.cuda.empty_cache(). Многоразовая память будет освобождена после этой операции.

...