【Pytorch学习】Dataset & Dataloader

1. Dataset

提供一种方式去获取数据及其label
功能有：获取每一个数据及其label；告诉我们总共有多少的数据

from torch.utils.data import Dataset
from PIL import Image
import os

class MyData(Dataset):

    def __init__(self, root_dir, label_dir):
        # 初始化
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.path = os.path.join(self.root_dir, self.label_dir)
        self.img_path = os.listdir(self.path)

    def __getitem__(self, idx):
        img_name = self.img_path[idx]
        img_item_path = os.path.join(self.root_dir, self.label_dir, img_name)
        img = Image.open(img_item_path)
        label = self.label_dir
        return img, label

    def __len__(self):
        return len(self.img_path)

root_dir = 'dataset/train'
ants_label_dir = 'ants'
bees_label_dir = 'bees'

# 实例化
ants_dataset = MyData(root_dir, ants_label_dir)
bees_dataset = MyData(root_dir, bees_label_dir)

train_dataset = ants_dataset + bees_dataset

2. Dataloader

作用：为后面的网络提供不同的获取图片的地址。
具体介绍可以参考以下代码：

import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

# 测试集
test_data = torchvision.datasets.CIFAR10("./datasets", train=False, transform=torchvision.transforms.ToTensor())

test_loader = DataLoader(dataset=test_data, batch_size=4, shuffle=True, num_workers=0, drop_last=False) # 随机取4张。drop_last为False代表若最后剩余的图片不足以一个batch，也不会舍去；为True表示会舍去最后不足的，保证所有的batch都是设定的大小4。shuffle设置为True会随机选取图片的。

# 测试数据集中第一张图片及label
img, target = test_data[0]
print(img.shape) # torch.Size([3, 32, 32])
print(target) # 3

# 展示
writer = SummaryWriter("CIFAR10") # save_dir
step = 0
for data in test_loader:
    imgs, targets = data
    print(imgs.shape) # torch.Size([4, 3, 32, 32]): 4张图片，3通道，32*32
    print(targets) # tensor([8, 5, 8, 0])：4张图片的labels综合在了一起
    writer.add_images("test_data", imgs, step)
    step = step + 1

writer.close()

在Terminal中输入代码tensorboard --logdir=CIFAR10，效果如下：