[英]How do I load the CelebA dataset on Google Colab, using torch vision, without running out of memory?
I am following a tutorial on DCGAN .我正在关注关于DCGAN的教程。 Whenever I try to load the CelebA dataset, torchvision uses up all my run-time's memory(12GB) and the runtime crashes.每当我尝试加载 CelebA 数据集时,torchvision 都会耗尽我所有的运行时内存(12GB)并且运行时崩溃。 Am looking for ways on how I can load and apply transformations to the dataset without hogging my run-time's resources.我正在寻找如何在不占用运行时资源的情况下加载和应用转换到数据集的方法。
Here is the part of the code that is causing issues.这是导致问题的代码部分。
# 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])
]))
The full notebook can be found here完整的笔记本可以在这里找到
PyTorch version: 1.7.1+cu101 PyTorch版本:1.7.1+cu101
Is debug build: False是否调试构建:False
CUDA used to build PyTorch: 10.1 CUDA 用于构建 PyTorch:10.1
ROCM used to build PyTorch: N/A ROCM 用于构建 PyTorch:不适用
OS: Ubuntu 18.04.5 LTS (x86_64)操作系统:Ubuntu 18.04.5 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0 GCC 版本:(Ubuntu 7.5.0-3ubuntu1~18.04)7.5.0
Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final) Clang 版本:6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.12.0 CMake版本:版本3.12.0
Python version: 3.6 (64-bit runtime) Python 版本:3.6(64 位运行时)
Is CUDA available: True CUDA 是否可用:真
CUDA runtime version: 10.1.243 CUDA 运行时版本:10.1.243
GPU models and configuration: GPU 0: Tesla T4 GPU型号及配置:GPU 0:特斯拉T4
Nvidia driver version: 418.67英伟达驱动版本:418.67
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5 cuDNN 版本:/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5
HIP runtime version: N/A HIP 运行时版本:不适用
MIOpen runtime version: N/A MIOpen 运行时版本:不适用
Versions of relevant libraries:相关库的版本:
Some of the things I have tried are:我尝试过的一些事情是:
# Download the dataset only
datasets.CelebA(data_root, download=True)
# Load the dataset here
celeba_data = datasets.CelebA(data_root, download=False, transforms=...)
ImageFolder
dataset class instead of the CelebA
class.使用 ImageFolder 数据集ImageFolder
而不是CelebA
class。 eg:例如:# Download the dataset only
datasets.CelebA(data_root, download=True)
# Load the dataset using the ImageFolder class
celeba_data = datasets.ImageFolder(data_root, transforms=...)
The memory problem is still persistent in either of the cases.在这两种情况下,memory 问题仍然存在。
I did not manage to find a solution to the memory problem.我没有找到解决 memory 问题的方法。 However, I came up with a workaround, custom dataset.但是,我想出了一个解决方法,自定义数据集。 Here is my implementation:这是我的实现:
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)
This implementation is memory efficient and works for my use case, even during training the memory used averages around(4GB).这个实现是 memory 高效的,并且适用于我的用例,即使在训练 memory 使用的平均值约为(4GB)时也是如此。 I would however, appreciate further intuition as to what might be causing the memory problems.但是,我希望进一步了解可能导致 memory 问题的原因。
Try the following:尝试以下操作:
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)
More details for Dataloader
class may be looked up here . Dataloader
class 的更多详细信息可在此处查找。 The above answer, courtesy this kaggle notebook .以上答案,由kaggle notebook提供。
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.