### PyTorch 自定义数据集代码示例 以下是基于 `torch.utils.data.Dataset` 类构建自定义数据集的一个完整代码示例。此代码展示了如何实现 `__init__`, `__len__`, 和 `__getitem__` 方法，以便加载和预处理数据。 ```python import torch from torch.utils.data import Dataset, DataLoader class CustomDataset(Dataset): def __init__(self, data, labels, transform=None): """ 初始化数据集。 参数: data (list or numpy array): 数据列表或数组。 labels (list or numpy array): 标签列表或数组。 transform (callable, optional): 可选的数据转换操作。 """ self.data = data self.labels = labels self.transform = transform def __len__(self): """ 返回数据集中样本的数量。 """ return len(self.data) def __getitem__(self, index): """ 获取指定索引处的单个样本及其标签。 参数: index (int): 样本索引。 返回: tuple: 单个样本及其对应的标签。 """ sample = self.data[index] label = self.labels[index] if self.transform is not None: sample = self.transform(sample) return sample, label # 示例：创建并使用自定义数据集 if __name__ == "__main__": # 假设我们有一些简单的二维数据和对应标签 X = [[1, 2], [3, 4], [5, 6], [7, 8]] # 数据 y = [0, 1, 0, 1] # 标签 dataset = CustomDataset(data=X, labels=y) # 创建数据集实例 dataloader = DataLoader(dataset, batch_size=2, shuffle=True) # 使用DataLoader批量加载数据 for batch_data, batch_labels in dataloader: print("Batch Data:", batch_data) print("Batch Labels:", batch_labels) ``` #### 关键点说明 - **初始化 (`__init__`) 方法**：用于存储数据和标签，并设置可选的变换操作[^3]。 - **长度 (`__len__`) 方法**：返回数据集中样本的总数[^3]。 - **获取项 (`__getitem__`) 方法**：根据给定索引返回相应的样本和标签。 - **DataLoader 的作用**：将数据集划分为批次，并支持随机打乱等功能[^4]。 --- ###