[PyTorch小试牛刀]实战六·准备自己的数据集用于训练（基于猫狗大战数据集）

[PyTorch小试牛刀]实战六·准备自己的数据集用于训练（基于猫狗大战数据集）

在上面几个实战中，我们使用的是Pytorch官方准备好的FashionMNIST数据集进行的训练与测试。本篇博文介绍我们如何自己去准备数据集，以应对更多的场景。

我们此次使用的是猫狗大战数据集，开始之前我们要先把数据处理一下，形式如下

datas
│
└───train
│ │
│ └───cats
│ │ │ cat1000.jpg
│ │ │ cat1001.jpg
│ │ │ …
│ └───dogs
│ │ │ dog1000.jpg
│ │ │ dog1001.jpg
│ │ │ …
└───valid
│ │
│ └───cats
│ │ │ cat0.jpg
│ │ │ cat1.jpg
│ │ │ …
│ └───dogs
│ │ │ dog0.jpg
│ │ │ dog1.jpg
│ │ │ …

train数据集中有23000张数据，valid数据集中有2000数据用于验证网络性能

代码部分
1.采用隐形字典形式，代码简练，不易理解

import torch as t
import torchvision as tv
import os

data_dir = "./datas"

BATCH_SIZE = 100

EPOCH = 10

transform = {
    x:tv.transforms.Compose(
        [tv.transforms.Resize([64,64]),tv.transforms.ToTensor()]#tv.transforms.Resize 用于重设图片大小
    ) 
    for x in ["train","valid"]
}

datasets = {
    x:tv.datasets.ImageFolder(root = os.path.join(data_dir,x),transform=transform[x])
    for x in ["train","valid"]
}

dataloader = {
    x:t.utils.data.DataLoader(dataset= datasets[x],
        batch_size=BATCH_SIZE,
        shuffle=True
    ) 
    for x in ["train","valid"]
}

b_x,b_y = next(iter(dataloader["train"]))

print(b_x.shape,b_y.shape)

index_classes = datasets["train"].class_to_idx

print(index_classes)

2.采用显性字典形式，代码稍多，易于理解

import torch as t
import torchvision as tv


data_dir = "./datas"

BATCH_SIZE = 100

EPOCH = 10

transform = {
    "train":tv.transforms.Compose(
        [tv.transforms.Resize([64,64]),tv.transforms.ToTensor()]
    ),
    "valid":tv.transforms.Compose(
        [tv.transforms.Resize([64,64]),tv.transforms.ToTensor()]
    ),
}

datasets = {
    "train":tv.datasets.ImageFolder(root = os.path.join(data_dir,"train"),transform=transform["train"]),
    "vaild":tv.datasets.ImageFolder(root = os.path.join(data_dir,"vaild"),transform=transform["vaild"]),
}
dataloader = {
    "train":t.utils.data.DataLoader(dataset= datasets["train"],
        batch_size=BATCH_SIZE,
        shuffle=True
    ),
    "valid":t.utils.data.DataLoader(dataset= datasets["valid"],
        batch_size=100,
        shuffle=True
    )
}

b_x,b_y = next(iter(dataloader["train"]))

print(b_x.shape,b_y.shape)

index_classes = datasets["train"].class_to_idx

print(index_classes)

输出结果

torch.Size([100, 3, 64, 64]) torch.Size([100])
{'cats': 0, 'dogs': 1}