pytorch Dataloader 如何自定义collate_fn过滤脏数据-CFANZ编程社区

PyTorch Dataloader如何自定义collate_fn过滤脏数据

在使用PyTorch进行深度学习模型训练时，通常需要使用Dataloader来加载和处理数据。Dataloader是一个迭代器，用于将数据集分成一批一批的样本，以便于模型的训练。但是，在实际应用中，我们常常会遇到一些脏数据或异常数据，这些数据可能会影响模型的训练效果。因此，我们需要通过自定义collate_fn来过滤掉这些脏数据，以确保模型的训练数据的质量。

实际问题

假设我们有一个猫狗分类的数据集，其中包含了一些损坏的图像文件。这些损坏的图像文件无法正常加载，会导致模型训练过程中出现错误。为了解决这个问题，我们需要在加载数据时过滤掉这些损坏的图像文件。

示例

首先，我们需要准备一个包含损坏图像的数据集。在这个示例中，我们先手动创建一个包含3个正常图像和2个损坏图像的数据集。

import os
import shutil

# 创建数据集目录
dataset_dir = 'dataset'
os.makedirs(dataset_dir, exist_ok=True)

# 创建正常图像
normal_image_files = ['cat_1.jpg', 'dog_1.jpg', 'cat_2.jpg']
for image_file in normal_image_files:
    with open(os.path.join(dataset_dir, image_file), 'w') as f:
        f.write('Fake image data')

# 创建损坏图像
corrupted_image_files = ['cat_3.jpg', 'dog_2.jpg']
for image_file in corrupted_image_files:
    with open(os.path.join(dataset_dir, image_file), 'w') as f:
        f.write('Corrupted image data')

接下来，我们使用自定义的collate_fn来过滤掉损坏的图像文件。在collate_fn中，我们可以使用try-except语句来捕捉加载图像时的异常，然后将损坏的图像文件从数据集中剔除。

import torch
from PIL import Image

def collate_fn(batch):
    filtered_batch = []
    for image_file, label in batch:
        try:
            # 尝试加载图像文件
            image = Image.open(image_file)
            filtered_batch.append((image, label))
        except Exception as e:
            # 图像加载失败，跳过该图像文件
            print(f'Failed to load image file: {image_file}')
    return filtered_batch

# 创建数据集
dataset = [(os.path.join(dataset_dir, image_file), 0 if 'cat' in image_file else 1) for image_file in os.listdir(dataset_dir)]

# 创建数据加载器
dataloader = torch.utils.data.DataLoader(dataset, batch_size=2, collate_fn=collate_fn)

# 遍历数据集
for images, labels in dataloader:
    print(f'Batch size: {len(images)}')
    # 在这里进行模型的训练

在上面的示例中，我们首先定义了一个名为collate_fn的函数，它将数据集中的每个样本作为输入，并返回过滤后的样本。在collate_fn中，我们尝试加载图像文件，如果加载成功，则将图像文件和对应的标签添加到过滤后的样本中；如果加载失败，则跳过该图像文件。最后，我们将过滤后的样本作为数据加载器的输出。