2.分类训练之数据准备-自定义数据集(CIFAR10数据加载)

前言

          pytorch中有许多公开的数据集可以加载，但是换成自己数据集，想训练自己定义的类别怎样做？这里以CIFAR10解析出来的图片数据为例。

代码分析

# 数据增强库
from torchvision import transforms
# 数据加载和读取的库
from torch.utils.data import DataLoader, Dataset
import os
from PIL import Image
import numpy as np
import glob

# 类别名
label_name = ["airplane", "automobile", "bird","cat", "deer", "dog","frog", "horse", "ship", "truck"]

# 将类别保存到字典中
label_dict = {}

# 将字符串类别转化成数字
for idx, name in enumerate(label_name):
    label_dict[name] = idx


# 返回通过PIL读取的图片数据并转换成RGB格式
def default_loader(path):
    return Image.open(path).convert("RGB")


# MyDataset需要完成三个函数的定义
class MyDataset(Dataset):
    # 初始化函数
    def __init__(self, im_list,
                       # 数据增强
                       transform=None,
                       #使用PIL进行图片数据的读取
                       loader = default_loader):
        # 初始化类
        super(MyDataset, self).__init__()

        # 定义数据列表
        imgs = []
        for im_item in im_list:
            # 获取每一张图片的路径，格式如下
            #"/home/user/pytorch_code_classes/06/cifar-10-python/cifar-10-batches-py/train/airplane/aeroplane_s_000021.png"
            # 获取类别名
            im_label_name = im_item.split("/")[-2]
            # 保存图片路径和类别名对应的id
            imgs.append([im_item, label_dict[im_label_name]])

        self.imgs = imgs
        self.transform = transform
        self.loader = loader

    # 定义数据的读取和增强，返回数据和类别
    def __getitem__(self, index):
        # index：训练时传入的索引值
        im_path, im_label = self.imgs[index]
        # 读取图片
        im_data = self.loader(im_path)
        # 数据增强
        if self.transform is not None:
            im_data = self.transform(im_data)

        return im_data, im_label

    # 返回样本的数量
    def __len__(self):
        return len(self.imgs)


# 训练集图片路径
im_train_list = glob.glob("/home/jf/pytorch_code_classes/06/cifar-10-python/cifar-10-batches-py/train/*/*.png")
# 测试集图片路径
im_test_list = glob.glob("/home/jf/pytorch_code_classes/06/cifar-10-python/cifar-10-batches-py/test/*/*.png")

# 训练集数据增强
train_transform = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.RandomVerticalFlip(),
        # transforms.RandomRotation(90),
        # transforms.ColorJitter(brightness=0.2, contrast=0.2, hue=0.2),
        # transforms.RandomGrayscale(0.2),
        # transforms.RandomCrop(28),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465),
                             (0.2023, 0.1994, 0.2010)),
    ])
# 测试集数据增强
test_transform = transforms.Compose([
        transforms.CenterCrop((28, 28)),
        transforms.ToTensor(),
        transforms.Normalize((0.4914, 0.4822, 0.4465),
                             (0.2023, 0.1994, 0.2010)),
    ])

# 定义数据
train_dataset = MyDataset(im_train_list,transform=train_transform)
test_dataset = MyDataset(im_test_list,transform =transforms.ToTensor())
# 加载数据
train_loader = DataLoader(dataset=train_dataset,
                               batch_size=128,
                               shuffle=True,
                               num_workers=4)

test_loader = DataLoader(dataset=test_dataset,
                               batch_size=128,
                               shuffle=False,
                               num_workers=4)

print("num_of_train", len(train_dataset))
print("num_of_test", len(test_dataset))

1）自定义类别名

        # 类别名
        label_name = ["airplane", "automobile", "bird","cat", "deer", "dog","frog", "horse", "ship", "truck"]

       # 将类别保存到字典中
       label_dict = {}

      # 将字符串类别转化成id
      for idx, name in enumerate(label_name):
          label_dict[name] = idx

2）定义MyDataset类别，需要完成三个函数的定义

    1.  def __init__：处理数据 添加图片的路径和id

    2.  def __getitem__: 定义数据的读取和增强，返回数据和类别
    3. def __len__:返回样本的数量

3)创建对象，加载到DataLoader中

具体代码解析见注释，对你有用的话，请点一下关注，谢谢！！！