Как загрузить набор данных CelebA в Google Colab с помощью Torch Vision, не исчерпывая памяти?

Я следую руководству по DCGAN. Каждый раз, когда я пытаюсь загрузить набор данных CelebA, torchvision использует всю мою оперативную память (12 ГБ), и среда выполнения дает сбой. Я ищу способы, как я могу загружать и применять преобразования к набору данных, не перегружая ресурсы времени выполнения.

Воспроизводить

Вот часть кода, вызывающая проблемы.

# Root directory for the dataset
data_root = 'data/celeba'
# Spatial size of training images, images are resized to this size.
image_size = 64

celeba_data = datasets.CelebA(data_root,
                              download=True,
                              transform=transforms.Compose([
                                  transforms.Resize(image_size),
                                  transforms.CenterCrop(image_size),
                                  transforms.ToTensor(),
                                  transforms.Normalize(mean=[0.5, 0.5, 0.5],
                                                       std=[0.5, 0.5, 0.5])
                              ]))

Полную записную книжку можно найти здесь

Среда

  • Версия PyTorch: 1.7.1 + cu101

  • Отладочная сборка: False

  • CUDA используется для сборки PyTorch: 10.1

  • ROCM, использованный для сборки PyTorch: N / A

  • ОС: Ubuntu 18.04.5 LTS (x86_64)

  • Версия GCC: (Ubuntu 7.5.0-3ubuntu1 ~ 18.04) 7.5.0

  • Версия Clang: 6.0.0-1ubuntu2 (теги / RELEASE_600 / final)

  • Версия CMake: версия 3.12.0

  • Версия Python: 3.6 (64-разрядная среда выполнения)

  • Доступен ли CUDA: True

  • Версия среды выполнения CUDA: 10.1.243

  • Модели и конфигурация графического процессора: GPU 0: Tesla T4

  • Версия драйвера Nvidia: 418.67

  • Версия cuDNN: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5

  • Версия среды выполнения HIP: N / A

  • Версия среды выполнения MIOpen: N / A

Версии соответствующих библиотек:

  • [pip3] numpy == 1.19.4
  • [pip3] torch == 1.7.1 + cu101
  • [pip3] torchaudio == 0.7.2
  • pip3] torchsummary == 1.5.1
  • [pip3] torchtext == 0.3.1
  • [pip3] torchvision == 0.8.2 + cu101
  • [conda] Не удалось собрать

Дополнительный контекст

Вот некоторые из вещей, которые я пробовал:

  • Скачивание и загрузка набора данных по отдельным строкам. например:
# Download the dataset only
datasets.CelebA(data_root, download=True)
# Load the dataset here
celeba_data = datasets.CelebA(data_root, download=False, transforms=...)
  • Использование класса набора данных ImageFolder вместо класса CelebA. например:
# Download the dataset only
datasets.CelebA(data_root, download=True)
# Load the dataset using the ImageFolder class
celeba_data = datasets.ImageFolder(data_root, transforms=...)

Проблема с памятью сохраняется в любом из случаев.

См. также:  Как вручную экранировать литералы SQL в Python

Вам необходимо использовать dataloader   —  person Kinyugo    schedule 01.01.2021

Похоже, что руководство было обновлено, чтобы продемонстрировать использование класса torchvision.datasets.ImageFolder. Используя рецепт @kinyugo download + unzip (но не их собственный набор данных), остальная часть учебника проходит нормально на colab без ошибок памяти.   —  person Kinyugo    schedule 25.06.2021

Понравилась статья? Поделиться с друзьями:
IT Шеф
Комментарии: 2
  1. Kinyugo

    Попробуйте следующее:

    from torchvision.datasets import ImageFolder
    from torch.utils.data import DataLoader
    from torchvision import transforms
    
    # Root directory for the dataset
    data_root = 'data/celeba'
    # Spatial size of training images, images are resized to this size.
    image_size = 64
    # batch size
    batch_size = 10
    
    transform=transforms.Compose([
                                  transforms.Resize(image_size),
                                  transforms.CenterCrop(image_size),
                                  transforms.ToTensor(),
                                  transforms.Normalize(mean=[0.5, 0.5, 0.5],
                                                       std=[0.5, 0.5, 0.5])
    
    dataset = ImageFolder(data_root, transform)
    
    data_loader = DataLoader(dataset=dataset, batch_size=batch_size, shuffle=True, num_workers=8, drop_last=True)
    

    Более подробную информацию о классе Dataloader можно найти здесь. Приведенный выше ответ любезно предоставлен этим блокнотом kaggle .

    Чтобы перебрать этот data_loader, используйте следующий цикл for: for i, (images, labels) in enumerate(data_loader):; где изображения — это набор изображений. person Kinyugo; 01.01.2021

    Проблема с памятью сохраняется даже после этого. person Kinyugo; 01.01.2021

  2. Kinyugo

    Решить проблему с памятью не удалось. Однако я придумал обходной путь — настраиваемый набор данных. Вот моя реализация:

    import os
    import zipfile 
    import gdown
    import torch
    from natsort import natsorted
    from PIL import Image
    from torch.utils.data import Dataset
    from torchvision import transforms
    
    ## Setup
    # Number of gpus available
    ngpu = 1
    device = torch.device('cuda:0' if (
        torch.cuda.is_available() and ngpu > 0) else 'cpu')
    
    ## Fetch data from Google Drive 
    # Root directory for the dataset
    data_root = 'data/celeba'
    # Path to folder with the dataset
    dataset_folder = f'{data_root}/img_align_celeba'
    # URL for the CelebA dataset
    url = 'https://drive.google.com/uc?id=1cNIac61PSA_LqDFYFUeyaQYekYPc75NH'
    # Path to download the dataset to
    download_path = f'{data_root}/img_align_celeba.zip'
    
    # Create required directories 
    if not os.path.exists(data_root):
      os.makedirs(data_root)
      os.makedirs(dataset_folder)
    
    # Download the dataset from google drive
    gdown.download(url, download_path, quiet=False)
    
    # Unzip the downloaded file 
    with zipfile.ZipFile(download_path, 'r') as ziphandler:
      ziphandler.extractall(dataset_folder)
    
    ## Create a custom Dataset class
    class CelebADataset(Dataset):
      def __init__(self, root_dir, transform=None):
        """
        Args:
          root_dir (string): Directory with all the images
          transform (callable, optional): transform to be applied to each image sample
        """
        # Read names of images in the root directory
        image_names = os.listdir(root_dir)
    
        self.root_dir = root_dir
        self.transform = transform 
        self.image_names = natsorted(image_names)
    
      def __len__(self): 
        return len(self.image_names)
    
      def __getitem__(self, idx):
        # Get the path to the image 
        img_path = os.path.join(self.root_dir, self.image_names[idx])
        # Load image and convert it to RGB
        img = Image.open(img_path).convert('RGB')
        # Apply transformations to the image
        if self.transform:
          img = self.transform(img)
    
        return img
    
    ## Load the dataset 
    # Path to directory with all the images
    img_folder = f'{dataset_folder}/img_align_celeba'
    # Spatial size of training images, images are resized to this size.
    image_size = 64
    # Transformations to be applied to each individual image sample
    transform=transforms.Compose([
        transforms.Resize(image_size),
        transforms.CenterCrop(image_size),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5],
                              std=[0.5, 0.5, 0.5])
    ])
    # Load the dataset from file and apply transformations
    celeba_dataset = CelebADataset(img_folder, transform)
    
    ## Create a dataloader 
    # Batch size during training
    batch_size = 128
    # Number of workers for the dataloader
    num_workers = 0 if device.type == 'cuda' else 2
    # Whether to put fetched data tensors to pinned memory
    pin_memory = True if device.type == 'cuda' else False
    
    celeba_dataloader = torch.utils.data.DataLoader(celeba_dataset,
                                                    batch_size=batch_size,
                                                    num_workers=num_workers,
                                                    pin_memory=pin_memory,
                                                    shuffle=True)
    

    Эта реализация эффективна с точки зрения памяти и работает для моего случая использования, даже во время обучения использованная память в среднем составляет около (4 ГБ). Однако я был бы признателен за дополнительную интуицию относительно того, что могло вызвать проблемы с памятью.

Добавить комментарий

;-) :| :x :twisted: :smile: :shock: :sad: :roll: :razz: :oops: :o :mrgreen: :lol: :idea: :grin: :evil: :cry: :cool: :arrow: :???: :?: :!: