### PyTorch 数据集加载与使用 #### 创建自定义数据集类 为了适应特定的文件夹结构，通常会创建一个继承 `torch.utils.data.Dataset` 的自定义数据集类。对于给定的文件夹结构： ```plaintext train ├── images └── labels └── 01.txt ``` 可以通过实现 `__init__`, `__len__` 和 `__getitem__` 方法来自定义数据集的行为。 ```python from torch.utils.data import Dataset import os from PIL import Image class CustomDataset(Dataset): def __init__(self, image_dir, label_dir, transform=None): self.image_dir = image_dir self.label_dir = label_dir self.transform = transform self.images = os.listdir(image_dir) def __len__(self): return len(self.images) def __getitem__(self, idx): img_path = os.path.join(self.image_dir, self.images[idx]) image = Image.open(img_path).convert('RGB') with open(os.path.join(self.label_dir, f"{os.path.splitext(self.images[idx])[0]}.txt"), 'r') as file: label = file.read() if self.transform: image = self.transform(image) return image, label ``` 此代码片段展示了如何构建一个简单的自定义数据集类[^4]。 #### 使用 DataLoader 加载数据 一旦有了数据集实例，就可以利用 `DataLoader` 来批量处理这些数据并提供多线程加速等功能。 ```python from torchvision.transforms import ToTensor from torch.utils.data import DataLoader dataset = CustomDataset( image_dir='path/to/train/images', label_dir='path/to/train/labels', transform=ToTensor() ) dataloader = DataLoader(dataset, batch_size=8, shuffle=True) ``` 这段代码说明了如何初始化 `CustomDataset` 并将其传递给 `DataLoader` 实现高效的数据读取[^1]。 #### 查看加载的数据样本 要查看由 `DataLoader` 提供的一个批次的数据，可以这样做: ```python import matplotlib.pyplot as plt images, labels = next(iter(dataloader)) print(f'Batch shape: {images.shape}') plt.figure(figsize=(12, 6)) for i in range(4): # 显示前四个图片作为例子 ax = plt.subplot(2, 2, i + 1) plt.imshow(images[i].permute(1, 2, 0)) plt.title(labels[i][:5]) # 只显示标签的一部分以便阅读 plt.axis('off') plt.show() ``` 上述脚本用于展示从 `val_loader` 中提取的一批图像及其对应的标签[^3]。