Pytorch «CUDA вне памяти» только с одной тренировочной партией - PullRequest
Новая
       13

Pytorch «CUDA вне памяти» только с одной тренировочной партией

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

В настоящее время я пытаюсь создать сеть, аналогичную U-Net, и у меня возникают странные проблемы с загрузкой моих данных в графический процессор.В частности, когда я загружаю свою модель в GPU с 

model = CNN(NUM_LAYERS)
model.to(DEVICE)

, использование памяти на GPU увеличивается до 1063 MiB.Это говорит мне о том, что моя модель занимает примерно один ГБ памяти.Все идет нормально.Затем с помощью 

inst = (iter(dataloader)).__next__()
input_tensor = inst[0].to(DEVICE, dtype=torch.float)

я загружаю одну партию обучающих данных в графический процессор.Общее использование памяти GPU возрастает до 1093 MiB, что кажется правильным и говорит мне, что пакет из 8 обучающих образов имеет размер около 30 МБ.Теперь я выполняю 

output = model(input_tensor)

, и использование памяти увеличивается до 21123 MiB!И если я изменю размер пакета на что-то большее, скажем, 16, я получу CUDA out of memory (GPU 1; 31.72 GiB total capacity; 29.91 GiB already allocated; 359.75 MiB free; 402.84 MiB cached).

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

Кстати, загрузчик данных выглядит так: 

class AsyncCNNData(Dataset):
    def __init__(self, list_paths, resize_dims=None, transform=None, device='cpu'):
        'Initialization'
        super().__init__()

        self.list_paths = list_paths
        self.transform = transform
        self.to_tensor = transforms.ToTensor()
        self.device = device
        self.resize_dims=resize_dims

    def __len__(self):
        'Denotes the total number of samples'
        return len(self.list_paths)

    def __getitem__(self, index):
        'Generates one sample of data'
        # Select sample
        sample_path = self.list_paths[index]

        # Load data and get label
        fs = cv.FileStorage(sample_path, cv.FILE_STORAGE_READ)

        cti_list = []
        for i in range(0, NUM_LAYERS, 1):
            img_name = "layer_{0:03d}".format(i)
            img = (fs.getNode(img_name)).mat()
            diff = np.subtract(self.resize_dims, img.shape)/2
            img = cv.copyMakeBorder(img, int(diff[0]), int(diff[0]), int(diff[1]), int(diff[1]), cv.BORDER_CONSTANT, None, 1)
            cti_list.append(img)

        input_tensor = np.stack(cti_list)
        input_tensor = input_tensor[np.newaxis, :, :, :]

        img = (fs.getNode("frame")).mat()
        diff = np.subtract(self.resize_dims, img.shape)/2
        img = cv.copyMakeBorder(img, int(diff[0]), int(diff[0]), int(diff[1]), int(diff[1]), cv.BORDER_CONSTANT, None, 1)
        img = img[:, :, np.newaxis]
        output_tensor = self.to_tensor(img)

        if self.transform != None:
            input_tensor = self.transform(input_tensor)
        else:
            input_tensor = torch.from_numpy(input_tensor)

        return input_tensor, output_tensor, sample_path

и 

yaml_files = pd.read_csv(DATA_PATH).values.flatten().tolist()
data = AsyncCNNData(yaml_files, device=DEVICE, resize_dims=RESIZE)

train_size = int(len(data)*TRAIN_RATIO)
data_train, data_test = random_split(data, [train_size, len(data)-train_size])
dataloader = DataLoader(data_train, batch_size=BATCH_SIZE, shuffle=True)

Большое спасибо!

...