在PyTorch中，加载数据并运行批处理操作的核心组件是`DataLoader`和`Dataset`类。下面详细介绍如何使用它们来进行数据加载及批量处理。 ### 使用 PyTorch 加载数据并运行批处理操作 #### 1. `Dataset` 类 这是用户自定义的数据集表示形式的基础类。你需要创建一个继承自 `torch.utils.data.Dataset` 的新类，并实现两个基本方法： - **`__len__(self)`**: 返回数据集中样本的数量。 - **`__getitem__(self, idx)`**: 根据给定索引获取单个样本及其标签（如果有的话）。 例如，假设我们有一个简单的图像分类任务，则可以构建如下的 `CustomImageDataset`: ```python from torch.utils.data import Dataset, DataLoader import os from PIL import Image import torchvision.transforms as transforms class CustomImageDataset(Dataset): def __init__(self, annotations_file, img_dir, transform=None, target_transform=None): self.img_labels = pd.read_csv(annotations_file) self.img_dir = img_dir self.transform = transform self.target_transform = target_transform def __len__(self): return len(self.img_labels) def __getitem__(self, idx): img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0]) image = Image.open(img_path).convert('RGB') label = self.img_labels.iloc[idx, 1] if self.transform: image = self.transform(image) if self.target_transform: label = self.target_transform(label) return image, label ``` #### 2. `DataLoader` 一旦有了定制化的 `Dataset` 实现之后，就可以利用 `DataLoader` 来简化对这个数据源的操作了。它提供了一种高效的方式从 `Dataset` 中抽取小批次(batch size)的数据供训练过程使用，并支持并发读取、随机打乱等功能。 以下是创建 `DataLoader` 并开始迭代访问数据的示例代码片段： ```python # 定义转换规则 transform = transforms.Compose([ transforms.Resize((64, 64)), # 将图片大小统一到指定尺寸 transforms.ToTensor() # 转换为张量类型 ]) # 创建dataset实例 train_dataset = CustomImageDataset( annotations_file='path/to/csv', img_dir='path/to/images', transform=transform, ) # 设置batch_size和其他参数 train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) for images, labels in train_loader: print(images.shape) # 输出形状 [Batch Size x Channels x Height x Width] break # 打印一次就停止循环测试效果 ``` 通过这种方式，您可以轻松地将大量数据分为较小的部分进行有效管理，非常适合深度学习模型的大规模训练需求。