使用PyTorch和TorchVision对自定义数据集进行训练有效测试拆分

时间:2020-05-15 04:47:05

标签: python pytorch torchvision

我有一些用于二进制分类任务的图像数据,并且图像被分为2个文件夹,分别为data / model_data / class-A和data / model_data / class-B。

总共有N张图像。我想将70/20/10分配给火车/ val /测试。 我正在使用PyTorch和Torchvision完成任务。这是我到目前为止的代码。

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils, datasets, models

data_transform = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

model_dataset = datasets.ImageFolder(root, transform=data_transform) 
train_count = int(0.7 * total_count) 
valid_count = int(0.2 * total_count)
test_count = total_count - train_count - valid_count
train_dataset, valid_dataset, test_dataset = torch.utils.data.random_split(model_dataset, (train_count, valid_count, test_count))
train_dataset_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKER)  
valid_dataset_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKER) 
test_dataset_loader  = torch.utils.data.DataLoader(test_dataset , batch_size=BATCH_SIZE, shuffle=False,num_workers=NUM_WORKER)
dataloaders = {'train': train_dataset_loader, 'val': valid_dataset_loader, 'test': test_dataset_loader}

由于两种原因,我认为这不是正确的方法。

  • 我正在对所有拆分应用相同的变换。 (显然,这不是我想要做的!解决方案很可能是答案here。)
  • 通常人们首先将原始数据分为测试/训练,然后他们 将火车分成火车/火车,而我直接将 将原始数据输入火车/验证/测试。 (这正确吗?)

所以,我的问题是,我在做什么正确? (可能不是)
而且如果不正确,我该如何编写数据加载程序以实现所需的拆分,以便可以对train / test / val中的每一个应用单独的转换?

1 个答案:

答案 0 :(得分:2)

通常人们首先将原始数据分为测试/训练和 然后他们将火车分成火车/火车,而我直接 将原始数据分为训练/验证/测试。 (对吗?)

是的,它是完全正确的,可读的并且总的来说还不错

我将相同的转换应用于所有拆分。 (这不是什么 我想做,很明显!解决方案很可能是 在这里回答。)

是的,这个答案是可能的,但是它太冗长了。您可以使用第三方工具torchdata,只需安装即可:

pip install torchdata

可以找到文档here(也免责声明:我是作者)。

它使您可以轻松地将转换映射到任何torch.utils.data.Dataset(在本例中为train)。您的代码看起来像这样(只需更改两行,检查注释,还可以对代码进行格式化以使其更容易遵循):

import torch
import torchvision

import torchdata as td

data_transform = torchvision.transforms.Compose(
    [
        torchvision.transforms.RandomResizedCrop(224),
        torchvision.transforms.RandomHorizontalFlip(),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize(
            mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
        ),
    ]
)

# Single change, makes an instance of torchdata.Dataset
# Works just like PyTorch's torch.utils.data.Dataset, but has
# additional capabilities like .map, cache etc., see project's description
model_dataset = td.datasets.WrapDataset(torchvision.datasets.ImageFolder(root))
# Also you shouldn't use transforms here but below
train_count = int(0.7 * total_count)
valid_count = int(0.2 * total_count)
test_count = total_count - train_count - valid_count
train_dataset, valid_dataset, test_dataset = torch.utils.data.random_split(
    model_dataset, (train_count, valid_count, test_count)
)

# Apply transformations here only for train dataset

train_dataset = train_dataset.map(data_transform)

# Rest of the code goes the same

train_dataset_loader = torch.utils.data.DataLoader(
    train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKER
)
valid_dataset_loader = torch.utils.data.DataLoader(
    valid_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKER
)
test_dataset_loader = torch.utils.data.DataLoader(
    test_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKER
)
dataloaders = {
    "train": train_dataset_loader,
    "val": valid_dataset_loader,
    "test": test_dataset_loader,
}

是的,我同意在拆分之前指定transform不太清楚,而IMO则更易读。