温馨提示×

在PyTorch中如何加载和处理数据集

小樊
94
2024-03-05 18:13:55
栏目: 编程语言

在PyTorch中,通常通过使用torch.utils.data.Datasettorch.utils.data.DataLoader来加载和处理数据集。

首先,创建一个自定义的数据集类,继承自torch.utils.data.Dataset,并实现__len____getitem__方法。在__getitem__方法中,可以根据索引加载和预处理数据。

import torch
from torch.utils.data import Dataset

class CustomDataset(Dataset):
    def __init__(self, data):
        self.data = data
        
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        sample = self.data[idx]
        # 进行数据预处理
        return sample

然后,实例化自定义数据集类并使用torch.utils.data.DataLoader创建一个数据加载器,指定批量大小和是否打乱数据。

data = [...]  # 数据集

dataset = CustomDataset(data)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=64, shuffle=True)

最后,可以通过迭代数据加载器来访问数据集中的数据。

for batch in dataloader:
    # 处理批量数据
    pass

0