【Pytorch】Load your own dataset

pytorch 在载入数据时用torchvision.datasets.ImageFolder 配合 torch.utils.data.DataLoader 很方便，但是只能遍历图片和图片的标签，无法灵活的获取图片的其他信息，比如图片的名字，本文介绍如何定义自己的 ImageFolder，在使用 Dataloader 时实现获取图片名字的功能！

---

---

1 ImageFolder and DataLoader

以分类为例，用 pytorch 的 torchvision.datasets.ImageFolder 配合 torch.utils.data.DataLoader 即可对数据按类别进行读取、预处理、分成 batch

import torchvision
import torch

train_dataset = torchvision.datasets.ImageFolder(
    train_data_pth,
    transforms.Compose([
        transforms.Resize(input_size,interpolation=2), # resize
        transforms.ToTensor(), # ToTensor
        normalize,])) # Normalization

train_loader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size, # set batchsize
    shuffle=False,
    num_workers=n_worker,
    pin_memory=True)

参考

• pytorch训练自己图像分类数据集
• 使用pytorch测试单张图片（test single image with pytorch）

---

ImageFolder 中 train_data_pth 是存放数据集的文件夹，文件结构应该如下

train_data_pth
	class1
		xxx.jpg
		...
	class2
		xxx.jpg
		...
	...
	classn
		xxx.jpg
		...

Dataloader 的参数介绍如下

• dataset：加载的数据集(Dataset对象)
• batch_size：batch size
• shuffle：是否将数据打乱
• sampler： 样本抽样，后续会详细介绍
• num_workers：使用多进程加载的进程数，0 代表不使用多进程
• collate_fn： 如何将多个样本数据拼接成一个 batch，一般使用默认的拼接方式即可’
• pin_memory：是否将数据保存在pin memory 区，pin memory 中的数据转到 GPU 会快一些
• drop_last：dataset中的数据个数可能不是 batch_size 的整数倍，drop_last 为 True 会将多出来不足一个batch的数据丢弃

参考 pytorch之DataLoader()函数

官网中 Dataloader 的介绍如下（https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader）

强调下，num_workers 表示一次可以装载 num_workers 个 batch，而不是一次装载一个 batch。

在训练和测试时，可用如下循环来对数据进行操作

for batch_images, batch_labels in train_loader:
	pass

把数据集的文件夹建立好，直接调用 ImageFolder 和 DataLoader 来进行数据的载入分批读取确实很方便，但是如果我们想知道哪些图片分类错误了， train_loader loader 中仅有 image（图片） 和 label 属性，没有 image name（图片名称） 属性，有些力不从心！

因此，我们可以自己写 ImageFolder 来实现读取 image、label、image name 的功能，当然熟悉这个流程后，以后可以进行更个性化的操作！

参考 从零开始深度学习Pytorch笔记（11）—— DataLoader类

---

补充

collate_fn参数用于是否需要以自定义的方式组织一个batch, 例子中将一个mini-batch的数据组织成numpy.ndarray的类型. 默认情况下collate_fn=None时,数据以元组的方式返回.

比如将多个numpy小数组组合成一个大的numpy数组:

def ssd_dataset_collate(batch):
    # print('ssd_dataset_collate函数被执行...')
    images = []
    bboxes = []
    for img, box in batch:
        images.append(img)
        bboxes.append(box)
    images = np.array(images)
    bboxes = np.array(bboxes)
    return images, bboxes


gen = DataLoader(train_dataset, \
                        batch_size=Batch_size, \
                        num_workers=8, \
                        pin_memory=True,\
                        drop_last=True, \
                        collate_fn=ssd_dataset_collate)
                        # collate_fn=None)

来自 torch.utils.data.DataLoader中的collate_fn的使用

2 OwnFolder and DataLoader

自己写数据读取和预处理，来替代 torchvision.datasets.ImageFolder 的功能，具体实现如下 class Own_Dataset 所示

class Own_Dataset(Dataset):
    def __init__(self, image_label_list, transform=None):
        super().__init__()
        self.samples_list = image_label_list  # xxx.jpg class1 
        self.transform = transform  # pre-processing of data

    def __getitem__(self, index):
        img_name = self.samples_list[index][0] # absolute path of image name
        with open(img_name,"rb") as f:
            img = Image.open(f).convert("RGB") # load image
        label = self.samples_list[index][1] # image label
        if img is None:
            print(img_name)
        if self.transform is not None:
            img = self.transform(img)
        return img, label, img_name

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

其中 image_label_list 为列表，存放着图片的绝对路径以及标签信息，格式如下

[(/train_data_pth/calss1/1.jpg，class1)，
(/train_data_pth/calss1/2.jpg，class1)，
...，
(/train_data_pth/calssn/m.jpg，classn)]

想实现更多功能，在 def __getitem__(self, index): 中定义即可，

__getitem__：实例[idx] 时触发

参考：【python】类（11）

配合 DataLoader 使用

train_loader = torch.utils.data.DataLoader(
    Own_Dataset(image_label_list=val_list,
               transform=transforms.Compose([
               	   transforms.Resize(input_size,interpolation=2), # resize
                   transforms.ToTensor(),
                   normalize,])),
    batch_size=test_batch_size,
    shuffle=False,
    num_workers=n_worker,
    pin_memory=True)

训练测试时，就可以访问图片，类别以及图片名信息了，如下所示

for batch_images, batch_labels,batch_names in train_loader:
	pass

3 transforms

下面介绍部分 torchvision.transforms 方法

更多的 torchvision.transforms 方法可以参考官网介绍

https://pytorch.org/docs/stable/torchvision/transforms.html

train_dataset = datasets.ImageFolder(
    train_data_pth,
    transforms.Compose([
        transforms.Resize(scale_size,interpolation=2),
        transforms.RandomRotation(5),
        transforms.ColorJitter(brightness=0.1,contrast=0.1,
                               saturation=0.1,hue=0.1),
        transforms.FiveCrop(input_size),
        transforms.Lambda(lambda crops: torch.stack([transforms.ToTensor()(crop) for crop in crops])),
        transforms.Lambda(lambda crops: torch.stack([transforms.Normalize(
            mean = [0.5,0.5,0.5],
            std = [0.5,0.5,0.5])(crop) for crop in crops]))
    ]))

input_size 和 scale_size 写成元组的形式，eg，（224，224） 和 （256，256）

Normalize 时注意 mean 和 std 一定要除以 255，值介于 0~1 之间

FiveCrop 或者 TenCrop 时，测试代码也需要进行相应的调整，如下

原来

out = net(batch_images)

现在

bs, ncrops, c, h, w, = batch_images.size()
result = net(batch_images.view(-1,c,h,w))
out = result.view(bs,ncrops,-1).mean(1)

计算数据集的均值和标准差

import os
import cv2
import numpy as np
from torch.utils.data import Dataset
from PIL import Image


def compute_mean_and_std(dataset):
    # 输入PyTorch的dataset，输出均值和标准差
    mean_r = 0
    mean_g = 0
    mean_b = 0

    for img, _ in dataset:
        img = np.asarray(img) # change PIL Image to numpy array
        mean_r += np.mean(img[:, :, 0])
        mean_g += np.mean(img[:, :, 1])
        mean_b += np.mean(img[:, :, 2])

    mean_r /= len(dataset)
    mean_g /= len(dataset)
    mean_b /= len(dataset)

    diff_r = 0
    diff_g = 0
    diff_b = 0

    N = 0

    for img, _ in dataset:
        img = np.asarray(img)

        diff_r += np.sum(np.power(img[:, :, 0] - mean_r, 2))
        diff_g += np.sum(np.power(img[:, :, 1] - mean_g, 2))
        diff_b += np.sum(np.power(img[:, :, 2] - mean_b, 2))

        N += np.prod(img[:, :, 0].shape)

    std_r = np.sqrt(diff_r / N)
    std_g = np.sqrt(diff_g / N)
    std_b = np.sqrt(diff_b / N)

    mean = (mean_r.item() / 255.0, mean_g.item() / 255.0, mean_b.item() / 255.0)
    std = (std_r.item() / 255.0, std_g.item() / 255.0, std_b.item() / 255.0)
    return mean, std