Pytorch自定义Dataset和DataLoader去除不存在和空数据的操作

PyTorch是一个流行的深度学习框架，可实现自定义数据集的灵活性和效率。在本攻略中，我们将学习如何自定义PyTorch的数据集和数据加载器，并使用它们来去除存在或空数据的条目。

自定义数据集

自定义数据集需要继承PyTorch的Dataset类，并重写其中的__len__和__getitem__方法。其中，__len__方法用于返回数据集的长度，而__getitem__方法提供了索引访问数据样本的功能。下面是一个自定义数据集的示例，该数据集从给定目录中读取所有图像文件，并返回图像的Tensor表示和其标签。

import os
from PIL import Image
from torch.utils.data import Dataset

class ImageDataset(Dataset):
    def __init__(self, root_dir):
        self.images = []
        self.labels = []
        for dir_name in os.listdir(root_dir):
            label = int(dir_name)
            for img_file in os.listdir(os.path.join(root_dir, dir_name)):
                img_path = os.path.join(root_dir, dir_name, img_file)
                self.images.append(img_path)
                self.labels.append(label)

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        img_path = self.images[idx]
        label = self.labels[idx]
        with open(img_path, 'rb') as f:
            img = Image.open(f)
            img = img.convert('RGB')
            return img, label

在这个示例中，我们首先在__init__方法中读取所有图像文件和它们的标签。然后，在__getitem__方法中使用PIL库读取图像，并将其转换为RGB格式的Tensor。最后，返回图像Tensor和标签。

自定义数据加载器

数据加载器可对自定义数据集进行批量加载和并行化处理。在PyTorch中，可以使用DataLoader类来创建数据加载器。下面是一个自定义数据加载器的示例，该数据加载器从给定的自定义数据集读取数据，同时实现了去除任何空数据的操作。

from torch.utils.data import DataLoader

class ImageDataLoader(DataLoader):
    def __init__(self, dataset, batch_size, shuffle=True, **kwargs):
        super().__init__(dataset, batch_size, shuffle, **kwargs)

    def __iter__(self):
        batch = []
        for item in super().__iter__():
            if item is None:
                continue
            batch.append(item)
        yield from batch

在这个示例中，我们首先创建一个继承自DataLoader的子类ImageDataLoader。然后在__iter__方法中，我们首先调用基类的__iter__方法，以获取每个批次的数据条目。但是，如果有任何条目为空，我们将跳过它们并继续处理下一个条目。最后，我们返回一个列表，其中包含所有非空条目的Tensor。

示例

下面是两个示例，演示如何使用上述自定义数据集和数据加载器去除存在或空数据的操作。

示例1：去除不存在的数据

假设我们的自定义数据集中包含多个图像，但是其中一个图像被删除或移动，因此不再存在。为了去除这样的无效数据项，我们可以在自定义数据集的__getitem__方法中添加异常处理。如果无法读取图像，则返回空值。然后，使用自定义数据加载器去除空值。

class ImageDataset(Dataset):
    def __init__(self, root_dir):
        self.images = []
        self.labels = []
        for dir_name in os.listdir(root_dir):
            label = int(dir_name)
            for img_file in os.listdir(os.path.join(root_dir, dir_name)):
                img_path = os.path.join(root_dir, dir_name, img_file)
                if os.path.exists(img_path):
                    self.images.append(img_path)
                    self.labels.append(label)

    def __getitem__(self, idx):
        img_path = self.images[idx]
        label = self.labels[idx]
        try:
            with open(img_path, 'rb') as f:
                img = Image.open(f)
                img = img.convert('RGB')
                return img, label
        except:
            return None, None

dataset = ImageDataset('data')
data_loader = ImageDataLoader(dataset, batch_size=8)
for images, labels in data_loader:
    print('Batch size:', len(images))

在这个示例中，我们首先在自定义数据集的__init__方法中检查每个图像是否存在。然后，在__getitem__方法中，我们使用异常处理来捕获无法读取图像的情况，并返回空值。最后，我们使用ImageDataLoader实例来加载数据，并使用if item is None语句在__iter__方法中去除空值。

示例2：去除空数据

假设我们的自定义数据集中包含多个图像文件夹，但其中一个图像文件夹为空。为了去除这样的空数据项，我们可以在自定义数据集的__init__方法中检查每个图像文件夹是否为空。如果为空，则跳过该文件夹，并以此不将其包含在数据集中。然后，我们可以使用自定义数据加载器去除空值。

class ImageDataset(Dataset):
    def __init__(self, root_dir):
        self.images = []
        self.labels = []
        for dir_name in os.listdir(root_dir):
            if len(os.listdir(os.path.join(root_dir, dir_name))) == 0:
                continue
            label = int(dir_name)
            for img_file in os.listdir(os.path.join(root_dir, dir_name)):
                img_path = os.path.join(root_dir, dir_name, img_file)
                self.images.append(img_path)
                self.labels.append(label)

    def __getitem__(self, idx):
        img_path = self.images[idx]
        label = self.labels[idx]
        with open(img_path, 'rb') as f:
            img = Image.open(f)
            img = img.convert('RGB')
            return img, label

dataset = ImageDataset('data')
data_loader = ImageDataLoader(dataset, batch_size=8)
for images, labels in data_loader:
    print('Batch size:', len(images))

在这个示例中，我们首先在自定义数据集的__init__方法中检查每个图像文件夹是否为空。如果是，则跳过该文件夹，并以此不将其图像包含在数据集中。然后，我们可以使用ImageDataLoader实例来加载数据，并使用if item is None语句在__iter__方法中去除空值。

总结

在本攻略中，我们学习了如何使用PyTorch自定义数据集和数据加载器，并使用这些工具实现了去除存在或空数据的操作。自定义数据集需要继承PyTorch的Dataset类，并重写其中的__len__和__getitem__方法。自定义数据加载器需要继承PyTorch的DataLoader类，并重写其中的__iter__方法。最后，我们实现了两个示例，演示了如何使用自定义数据集和数据加载器去除无效或空数据项。