在PyTorch中，自定义数据集通常涉及到两个主要的类：`torch.utils.data.Dataset` 和 `torch.utils.data.DataLoader`。`Dataset` 类用于表示数据集，你需要重写其中的 `__len__` 和 `__getitem__` 方法。而 `DataLoader` 类用于加载数据，它可以自动地将数据集分批、打乱、多线程处理等。 以下是一个简单的例子，展示了如何自定义数据集并使用 `DataLoader` 加载它。 ### 步骤 1: 自定义 Dataset 首先，你需要从 `torch.utils.data.Dataset` 继承并定义你自己的数据集类。你需要实现 `__len__` 和 `__getitem__` 方法。 ```python from torch.utils.data import Dataset import torch class CustomDataset(Dataset): def __init__(self, data, labels): """ Args: data: 列表或数组，包含你的特征数据 labels: 列表或数组，包含与数据对应的标签 """ self.data = data self.labels = labels def __len__(self): # 返回数据集中的样本数量 return len(self.data) def __getitem__(self, idx): # 根据索引idx获取样本 sample = self.data[idx] label = self.labels[idx] # 根据需要处理数据，例如转换为Tensor sample = torch.tensor(sample, dtype=torch.float32) label = torch.tensor(label, dtype=torch.long) return sample, label ``` ### 步骤 2: 使用 DataLoader 加载数据 一旦你定义了数据集类，就可以使用 `DataLoader` 来加载数据了。你可以设置批量大小、是否打乱数据、是否多线程加载等参数。 ```python from torch.utils.data import DataLoader # 假设你已经有了一些数据和标签 data = [[1, 2], [3, 4], [5, 6], [7, 8]] # 示例数据 labels = [0, 1, 0, 1] # 示例标签 # 实例化你的数据集 dataset = CustomDataset(data, labels) # 创建 DataLoader # batch_size 表示每个批次加载的样本数 # shuffle=True 表示在每个epoch开始时打乱数据 # num_workers=0 表示不使用额外的进程来加载数据 dataloader = DataLoader(dataset, batch_size=2, shuffle=True, num_workers=0) # 使用 DataLoader 遍历数据 for batch_idx, (data, target) in enumerate(dataloader): print(f"Batch {batch_idx}, Data: {data}, Target: {target}") ``` 在上面的代码中，`DataLoader` 将自动从 `CustomDataset` 中获取数据，并根据 `batch_size` 分批处理。`shuffle=True` 参数使得每个epoch开始时数据会被打乱，这有助于模型训练时的泛化能力。 这样，你就成功地在PyTorch中自定义了数据集，并使用 `DataLoader` 高效地加载了数据。