时间戳

首页 > 解决方案 > 创建火车和测试数据加载器

python - 创建火车和测试数据加载器

问题描述

我实际上有一个RealPhotos包含 17000jpg张照片的目录。我有兴趣创建一个火车数据加载器和一个测试数据加载器

ls RealPhotos/
2007_000027.jpg  2008_007119.jpg  2010_001501.jpg  2011_002987.jpg
2007_000032.jpg  2008_007120.jpg  2010_001502.jpg  2011_002988.jpg
2007_000033.jpg  2008_007123.jpg  2010_001503.jpg  2011_002992.jpg
2007_000039.jpg  2008_007124.jpg  2010_001505.jpg  2011_002993.jpg
2007_000042.jpg  2008_007129.jpg  2010_001511.jpg  2011_002994.jpg
2007_000061.jpg  2008_007130.jpg  2010_001514.jpg  2011_002996.jpg
2007_000063.jpg  2008_007131.jpg  2010_001515.jpg  2011_002997.jpg
2007_000068.jpg  2008_007133.jpg  2010_001516.jpg  2011_002999.jpg
2007_000121.jpg  2008_007134.jpg  2010_001518.jpg  2011_003002.jpg
2007_000123.jpg  2008_007138.jpg  2010_001520.jpg  2011_003003.jpg
...

我知道我可以继承 TensorDataset 以使其与未标记的数据兼容

class UnlabeledTensorDataset(TensorDataset):
    """Dataset wrapping unlabeled data tensors.
    
    Each sample will be retrieved by indexing tensors along the first
    dimension.
    
    Arguments:
        data_tensor (Tensor): contains sample data.
    """
    def __init__(self, data_tensor):
        self.data_tensor = data_tensor
    
    def __getitem__(self, index):
        return self.data_tensor[index]

以及用于训练自动编码器的一些东西

X_train     = rnd.random((300,100))
train       = UnlabeledTensorDataset(torch.from_numpy(X_train).float())
train_loader= data_utils.DataLoader(train, batch_size=1)

for epoch in range(50):
    for batch in train_loader:
        data = Variable(batch)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, data)

标签： pythonpytorchdatasetdataloader

解决方案

您首先需要定义一个 Dataset ( torch.utils.data.Dataset)，然后您可以在其上使用 DataLoader。您的训练数据集和测试数据集之间没有区别，您可以定义一个通用数据集，该数据集将查看特定目录并将每个索引映射到一个唯一文件。

class MyDataset(Dataset):
    def __init__(self, directory):
        self.files = os.listdir(directory)

    def __getitem__(self, index):
        img = Image.open(self.files[index]).convert('RGB')
        return T.ToTensor()(img)

哪里T是指torchvision.transform和图像是从导入的PIL。

然后，您可以使用

data_set = MyDataset('./RealPhotos')

从那里您可以使用torch.utils.data.random_split执行拆分：

train_len = int(len(data_set)*0.7)
train_set, test_set = random_split(data_set, [train_len, data_set - train_len])

然后torch.utils.data.DataLoader像你一样使用：

train_loader = DataLoader(train_set, batch_size=1, shuffle=True)
test_loader = DataLoader(test_set, batch_size=16, shuffle=False)

推荐阅读