### PyTorch 数据集加载示例代码模板 #### 创建自定义数据集类 为了创建一个可以被 `DataLoader` 使用的数据集，通常需要继承 `torch.utils.data.Dataset` 类，并实现两个方法：`__len__()` 和 `__getitem__()`. 这些操作允许迭代器访问数据集中每一个样本。 ```python from torch.utils.data import Dataset, DataLoader import pandas as pd from torchvision import transforms as T class CustomDataset(Dataset): """Custom dataset class.""" def __init__(self, csv_file, transform=None): """ Args: csv_file (string): Path to the csv file with annotations. transform (callable, optional): Optional transform to be applied on a sample. """ self.annotations = pd.read_csv(csv_file) self.transform = transform def __len__(self): return len(self.annotations) def __getitem__(self, idx): if torch.is_tensor(idx): idx = idx.tolist() img_name = self.annotations.iloc[idx, 0] image = Image.open(img_name).convert('RGB') label = self.annotations.iloc[idx, 1] if self.transform: image = self.transform(image) return image, label ``` 此部分展示了如何通过 Pandas 来辅助读取 CSV 文件中的路径和标签信息[^1]. #### 定义转换链 在实际应用中，经常会对图像执行一系列预处理步骤。这些变换可以通过组合多个简单的变换来构建复杂的流水线： ```python transformations = T.Compose([ T.Resize((256, 256)), # 调整大小到固定尺寸 T.RandomHorizontalFlip(), # 随机水平翻转 T.ToTensor(), # 将PIL图片转化为张量 ]) ``` 这段代码说明了使用 `torchvision.transforms` 模块来进行常见的图像增强操作[^2]. #### 初始化数据集实例并设置 Data Loader 一旦有了上述组件，就可以轻松地初始化数据集对象以及相应的加载器: ```python dataset = CustomDataset('./path/to/annotations.csv', transformations) dataloader = DataLoader( dataset, batch_size=8, shuffle=True, num_workers=4 ) ``` 这里指出了 `DataLoader` 的主要参数及其作用，比如批量大小 (`batch_size`)、是否打乱顺序 (`shuffle`) 及工作进程数 (`num_workers`). 此外还强调了 `DataLoader` 实际上调用了 `__getitem__()` 方法以获取每一批次的数据. #### 训练循环内遍历 Dataloader 最后，在训练过程中可以直接利用 Python 的 for 循环语法来方便地枚举由 `Dataloader` 提供的一个个小批数据. ```python for epoch in range(num_epochs): running_loss = 0.0 for i, data in enumerate(dataloader, start=0): inputs, labels = data optimizer.zero_grad() outputs = net(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() ``` 这部分描述了一个典型的训练过程片段，其中包含了从 dataloader 获取输入和目标变量的过程[^3].