如何使用 Torch Vision 在 Google Colab 上加载 CelebA 数据集，而不会耗尽内存？

如何解决如何使用 Torch Vision 在 Google Colab 上加载 CelebA 数据集，而不会耗尽内存？

我正在学习关于 DCGAN 的教程。每当我尝试加载 CelebA 数据集时，torchvision 会耗尽我所有运行时的内存（12GB）并且运行时崩溃。我正在寻找如何在不占用运行时资源的情况下加载和应用数据集转换的方法。

复制

这是导致问题的代码部分。

# 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])
                              ]))

可以找到完整的笔记本here

环境

PyTorch 版本：1.7.1+cu101
是否为调试版本：False
用于构建 PyTorch 的 CUDA：10.1
用于构建 PyTorch 的 ROCM：不适用
操作系统：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 是否可用：正确
CUDA 运行时版本：10.1.243
GPU 型号和配置：GPU 0：Tesla T4
Nvidia 驱动程序版本：418.67
cuDNN 版本：/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
HIP 运行时版本：不适用
MIOpen 运行时版本：不适用

附加上下文

我尝试过的一些事情是：

在单独的行上下载和加载数据集。例如：

# 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=...)

在任何一种情况下，内存问题仍然存在。

解决方法

我没有设法找到内存问题的解决方案。但是，我想出了一个解决方法，自定义数据集。这是我的实现：

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])
])
# 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)

此实现具有内存效率，适用于我的用例，即使在训练期间使用的内存平均约为（4GB）。但是，对于可能导致内存问题的原因，我希望能有进一步的直觉。

尝试以下操作：

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),0.5])

dataset = ImageFolder(data_root,transform)

data_loader = DataLoader(dataset=dataset,shuffle=True,num_workers=8,drop_last=True)

Dataloader 类的更多详细信息可以查询 here。以上回答，礼貌this kaggle notebook。

如何使用 Torch Vision 在 Google Colab 上加载 CelebA 数据集，而不会耗尽内存？

如何解决如何使用 Torch Vision 在 Google Colab 上加载 CelebA 数据集，而不会耗尽内存？

复制

环境

附加上下文

解决方法

相关推荐