【pytorch】学习笔记 | DataSet和DateLoader

持续coding，保持happy

前言

我们在深度学习的时候，最开始总有一个问题。
如何把自己本地的数据，包装成对应的DataSet和DataLoader。
一般不会遇到这种问题，主要是我任务的开源数据量很小，或者有一些有代码，却没有数据。

数据集

首先看一下我的数据集，它是图片分类的二分类的任务。train是训练数据，val是验证数据。ants里面全部都是蚂蚁的图片，bees全部都是蜜蜂的图片。
现在还是比较好理解的，我需要把这些数据读进去，放入dataset中。

官网代码详解

1.标记csv文件

这是官网link提供的参考数据和代码，它的数据表如下所示，左边是图片名，右边是标签名。

2.官方代码

import os
import pandas as pd
from torch.utils.data import Dataset
from torchvision.io import read_image


class CustomImageDataset(Dataset):
    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform

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

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = read_image(img_path)
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label

在写自己的Dataset的时候，需要重写3个方法，len，init，getitem。

我们可能需要想象一下，读取左边的图片地址，然后再磁盘上面读取对应的文件（图片）和标签。

def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
    self.img_labels = pd.read_csv(annotations_file)
    self.img_dir = img_dir
    self.transform = transform
    self.target_transform = target_transform

上面代码是初始化一些参数，比如：

annotations_file 文件名
img_dir 放图片的地址
transform 图片可能需要转换成对应的格式
target_transform label可能需要转换成对应的格式

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

上面代码就是返回一个dataset的长度。

def __getitem__(self, idx):
    img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
    image = read_image(img_path)
    label = self.img_labels.iloc[idx, 1]
	if self.transform:
		image = self.transform(image)
	if self.target_transform:
		label = self.target_transform(label)
    return image, label

上面代码是再磁盘中读取文件，并进行一些相应的转变：

img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0]) # 拼接完成图片的地址
image = read_image(img_path) # 读取图片
label = self.img_labels.iloc[idx, 1] # 读取label
if self.transform: # 必要的转变
image = self.transform(image)
if self.target_transform:
label = self.target_transform(label)

2.自己的数据集和代码

import os
from torch.utils.data import Dataset
from PIL import Image
from torchvision.io import read_image


class MyData(Dataset):
    def __init__(self, root_dir, label_dir, transform=None, target_transform=None):
        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 __len__(self):
        return len(self.img_path)

    def __getitem__(self, idx):
        img_name = self.img_path[idx]
        img_item_path = os.path.join(self.root_dir, self.label_dir, self.img_path[idx])
        image = Image.open(img_item_path)
        # 另一种方式读取3维向量
        # image = read_image(img_item_path)
        label = self.label_dir
        return image, label


if __name__ == "__main__":
    train_root_dir = ".data/hymenoptera_data/train"
    test_root_dir = ".data/hymenoptera_data/val"
    ants_label = "ants"
    bees_label = "bees"
    train_ants_dataset = MyData(train_root_dir, ants_label)
    train_bees_dataset = MyData(train_root_dir, bees_label)
    test_ants_dataset = MyData(test_root_dir, ants_label)
    test_bees_dataset = MyData(test_root_dir, bees_label)
    print("done")