本文的文字及图片来源于网络,仅供学习、交流使用,不具有任何商业用途,,版权归原作者所有，如有问题请及时联系我们以作处理

作者：weixin_46389668 来源：CSDN

本文链接：https://blog.csdn.net/weixin_46389668/article/details/111185617

pytorch数据集读取方法

**CIFAR10数据集的定义方法如下：**

`dataset_dir = '../../../dataset/'
torchvision.datasets.CIFAR10(dataset_dir, train=True, transform=None, target_transform=None, download=False) `

dataset_dir：存放数据集的路径。
train（bool，可选）–如果为True，则构建训练集，否则构建测试集。
transform：定义数据预处理，数据增强方案都是在这里指定。
target_transform：标注的预处理，分类任务不常用。
download：是否下载，若为True则从互联网下载，如果已经在dataset_dir下存在，就不会再次下载

读取示例1(从网上自动下载)

train_data = torchvision.datasets.CIFAR10('../../../dataset', 
                                                      train=True, 
                                                      transform=None,  
                                                      target_transform=None, 
                                                      download=True)          

读取示例2(示例1基础上附带数据增强)

# 读取训练集
custom_transform=transforms.transforms.Compose([
              transforms.Resize((64, 64)),    # 缩放到指定大小 64*64
              transforms.ColorJitter(0.2, 0.2, 0.2),    # 随机颜色变换
              transforms.RandomRotation(5),    # 随机旋转
              transforms.Normalize([0.485,0.456,0.406],    # 对图像像素进行归一化
                                   [0.229,0.224,0.225])])
train_data=torchvision.datasets.CIFAR10('../../../dataset', 
                                        train=True,                                       
                                        transform=custom_transforms,
                                        target_transform=None, 
                                        download=False) 

数据加载

# 读取数据集
train_data=torchvision.datasets.CIFAR10('../../../dataset', train=True, 
                                                      transform=None,  
                                                      target_transform=None, 
                                                      download=True)          
# 实现数据批量读取
train_loader = torch.utils.data.DataLoader(train_data,
                                           batch_size=2,
                                           shuffle=True,
                                           num_workers=4)        

自定义数据集及读取方法

简单的对pytorch读取数据一般化pipeline的描述，就是下面的这个流程：
图像数据 ? 图像索引文件 ? 使用Dataset构建数据集 ? 使用DataLoader读取数据
2.2.1 图像索引文件制作
下载MNIST的图像和标签数据到Dive-into-CV-PyTorch/dataset/MNIST/目录下，得到下面的压缩文件并解压暂存，以用来充当自己的图像数据集。

train-images-idx3-ubyte.gz: training set images (9912422 bytes) ? train-images-idx3-ubyte（解压后）
train-labels-idx1-ubyte.gz: training set labels (28881 bytes) ? train-labels-idx1-ubyte（解压后）
t10k-images-idx3-ubyte.gz:  test set images (1648877 bytes) ? t10k-images-idx3-ubyte（解压后）
t10k-labels-idx1-ubyte.gz:  test set labels (4542 bytes) ? t10k-labels-idx1-ubyte（解压后）

运行如下代码，实现图像数据的本地存储和索引文件的制作，我们将图像按照训练集和测试集分别存放，并且分别制作训练集和测试集的索引文件，在索引文件中将记录图像的文件名和标签信息。

import os
from skimage import io
import torchvision.datasets.mnist as mnist

# 数据文件读取
root = r'./MNIST/'  # MNIST解压文件根目录
train_set = (
    mnist.read_image_file(os.path.join(root, 'train-images-idx3-ubyte')),
    mnist.read_label_file(os.path.join(root, 'train-labels-idx1-ubyte'))
)
test_set = (
    mnist.read_image_file(os.path.join(root, 't10k-images-idx3-ubyte')),
    mnist.read_label_file(os.path.join(root, 't10k-labels-idx1-ubyte'))
)

# 数据量展示
print('train set:', train_set[0].size())
print('test set:', test_set[0].size())


def convert_to_img(save_path, train=True):
    '''
    将图片存储在本地，并制作索引文件
    @para: save_path  图像保存路径，将在路径下创建train、test文件夹分别存储训练集和测试集
    @para: train      默认True，本地存储训练集图像，否则本地存储测试集图像 
    '''
    if train:
        f = open(save_path + 'train.txt', 'w')
        data_path = save_path + '/train/'
        if (not os.path.exists(data_path)):
            os.makedirs(data_path)
        for i, (img, label) in enumerate(zip(train_set[0], train_set[1])):
            img_path = data_path + str(i) + '.jpg'
            io.imsave(img_path, img.numpy())
            int_label = str(label).replace('tensor(', '')
            int_label = int_label.replace(')', '')
            f.write(str(i)+'.jpg' + ',' + str(int_label) + '\n')
        f.close()
    else:
        f = open(save_path + 'test.txt', 'w')
        data_path = save_path + '/test/'
        if (not os.path.exists(data_path)):
            os.makedirs(data_path)
        for i, (img, label) in enumerate(zip(test_set[0], test_set[1])):
            img_path = data_path + str(i) + '.jpg'
            io.imsave(img_path, img.numpy())
            int_label = str(label).replace('tensor(', '')
            int_label = int_label.replace(')', '')
            f.write(str(i)+'.jpg' + ',' + str(int_label) + '\n')
        f.close()


# 根据需求本地存储训练集或测试集
save_path = r'./MNIST/mnist_data/'
convert_to_img(save_path, True)
convert_to_img(save_path, False)

2.2.2 构建自己的Dataset

from torch.utils.data.dataset import Dataset

class MyDataset(Dataset):  # 继承Dataset类
   def __init__(self):
       # 初始化图像文件路径或图像文件名列表等
       pass

   def __getitem__(self, index):
        # 1.根据索引index从文件中读取一个数据（例如，使用numpy.fromfile，PIL.Image.open，cv2.imread）
        # 2.预处理数据（例如torchvision.Transform）
        # 3.返回数据对（例如图像和标签）
       pass

   def __len__(self):
       return count  # 返回数据量

__init__() : 初始化模块，初始化该类的一些基本参数
__getitem__() : 接收一个index，这个index通常指的是一个list的index，这个list的每个元素就包含了图片数据的路径和标签信息,返回数据对（图像和标签）
__len__() : 返回所有数据的数量

import pandas as pd
import numpy as np
from PIL import Image
import torch
from torch.utils.data import Dataset
from torchvision import transforms

class MnistDataset(Dataset):

    def __init__(self, image_path, image_label, transform=None):
        super(MnistDataset, self).__init__()
        self.image_path = image_path  # 初始化图像路径列表
        self.image_label = image_label  # 初始化图像标签列表
        self.transform = transform  # 初始化数据增强方法

    def __getitem__(self, index):
        """
        获取对应index的图像，并视情况进行数据增强
        """
        image = Image.open(self.image_path[index])
        image = np.array(image)
        label = float(self.image_label[index])

        if self.transform is not None:
            image = self.transform(image)

        return image, torch.tensor(label)

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


def get_path_label(img_root, label_file_path):
    """
    获取数字图像的路径和标签并返回对应列表
    @para: img_root: 保存图像的根目录
    @para:label_file_path: 保存图像标签数据的文件路径 .csv 或 .txt 分隔符为','
    @return: 图像的路径列表和对应标签列表
    """
    data = pd.read_csv(label_file_path, names=['img', 'label'])
    data['img'] = data['img'].apply(lambda x: img_root + x)
    return data['img'].tolist(), data['label'].tolist()


# 获取训练集路径列表和标签列表
train_data_root = './dataset/MNIST/mnist_data/train/'
train_label = './dataset/MNIST/mnist_data/train.txt'
train_img_list, train_label_list = get_path_label(train_data_root, train_label)  
# 训练集dataset
train_dataset = MnistDataset(train_img_list,
                             train_label_list,
                             transform=transforms.Compose([transforms.ToTensor()]))

# 获取测试集路径列表和标签列表
test_data_root = './dataset/MNIST/mnist_data/test/'
test_label = './dataset/MNIST/mnist_data/test.txt'
test_img_list, test_label_list = get_path_label(test_data_root, test_label)
# 测试集sdataset
test_dataset = MnistDataset(test_img_list,
                            test_label_list,
                            transform=transforms.Compose([transforms.ToTensor()]))

2.2.3 使用DataLoader批量读取数据

使用 DataLoader 批量的读取数据，相当于帮我们完成一个batch的数据组装工作。

Dataloader 为一个迭代器，最基本的使用方法就是传入一个 Dataset 对象，在Dataloader中，会触发Dataset对象中的 gititem() 函数，逐次读取数据，并根据 batch_size 产生一个 batch 的数据，实现批量化的数据读取。

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, num_workers=0, collate_fn=default_collate, pin_memory=False, drop_last=False)

dataset：加载的数据集(Dataset对象)
batch_size：一个批量数目大小
shuffle:：是否打乱数据顺序
sampler： 样本抽样方式
num_workers：使用多进程加载的进程数，0代表不使用多进程
collate_fn： 将多个样本数据组成一个batch的方式，一般使用默认的拼接方式，可以通过自定义这个函数来完成一些特殊的读取逻辑。
pin_memory：是否将数据保存在pin memory区，pin memory中的数据转到GPU会快一些
drop_last：为True时，dataset中的数据个数不是batch_size整数倍时，将多出来不足一个batch的数据丢弃

Dataset是对本地数据读取逻辑的定义；而DataLoader是对Dataset对象的封装，执行调度，将一个batch size的图像数据组装在一起，实现批量读取数据。

对于图像分类问题，torchvision还提供了一种文件目录组织形式可供调用，即ImageFolder，因为利用了分类任务的特性，此时就不用再另行创建一份标签文件了。这种文件目录组织形式，要求数据集已经自觉按照待分配的类别分成了不同的文件夹，一种类别的文件夹下面只存放同一种类别的图片。