### PyTorch Dataset 存储于 GPU 的最佳实践 在 PyTorch 中，默认情况下 `Dataset` 和 `DataLoader` 是设计用于加载数据到 CPU 内存中的。然而，在某些高性能计算场景下，可能希望直接将数据预加载至 GPU 以减少 I/O 开销并提高模型训练效率。以下是关于如何实现这一目标的最佳实践： #### 方法一：通过自定义 `Dataset` 类实现在初始化阶段加载数据到 GPU 可以创建一个继承自 `torch.utils.data.Dataset` 的子类，并在该类的 `__init__` 函数中预先将整个数据集加载到 GPU 上。 ```python import torch from torch.utils.data import Dataset, DataLoader class GPUDataset(Dataset): def __init__(self, data, labels, device='cuda'): self.data = torch.tensor(data).to(device) # 将数据移动到指定设备上 (e.g., 'cuda') self.labels = torch.tensor(labels).to(device) def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx], self.labels[idx] # 假设我们有一个简单的二维张量作为输入数据和标签 data = [[1, 2], [3, 4], [5, 6]] labels = [0, 1, 0] gpu_dataset = GPUDataset(data, labels) dataloader = DataLoader(gpu_dataset, batch_size=2, shuffle=True) ``` 这种方法适用于小型数据集能够完全驻留在显存的情况[^1]。对于较大的数据集，则不推荐此方法，因为可能会超出 GPU 显存容量限制。 #### 方法二：利用 pin_memory 提升性能配合异步传输机制 如果无法一次性把所有数据都放到 GPU 上，那么可以通过设置 `pin_memory=True` 参数来优化从主机内存向设备内存的数据拷贝过程。当启用 pinned memory 后，CUDA 能够更高效地执行同步或异步的数据传递操作。 ```python cpu_dataset = ... # 定义常规CPU上的dataset对象 pinned_dataloader = DataLoader(cpu_dataset, batch_size=32, num_workers=4, pin_memory=True) for inputs, targets in pinned_dataloader: inputs, targets = inputs.cuda(non_blocking=True), targets.cuda(non_blocking=True) ... ``` 这里需要注意的是，虽然每次迭代都会涉及额外的时间成本用来完成实际的数据迁移动作，但由于采用了 non-blocking 方式加上后台线程的支持，总体影响相对较小[^2]。 综上所述，针对不同规模大小的数据集可以选择不同的策略来进行处理；而对于那些特别庞大的数据源来说，通常建议采用第二种方案即借助 pinned memory 技术逐步按需上传所需部分而非试图全部塞入单一硬件资源之中。