DataLoader源代码剖析

前言

• dataloader  本质是一个可迭代对象，使用  iter()  访问，不能使用  next()  访问；
• 使用   iter(dataloader)  返回的是一个迭代器，然后可以使用   next  访问；
• 也可以使用  `for inputs, labels in dataloaders`  进行可迭代对象的访问；
• 一般我们实现一个 datasets 对象，传入到  dataloader  中；然后内部使用  yeild  返回每一次  batch  的数据；

 

pytorch 的数据加载到模型的操作顺序是这样的：

　　① 创建一个 Dataset 对象
　　② 创建一个 DataLoader 对象
　　③ 循环这个 DataLoader 对象，将img, label加载到模型中进行训练

dataset = MyDataset()
dataloader = DataLoader(dataset)
num_epoches = 100
for epoch in range(num_epoches):
    for img, label in dataloader:
        ....

　　所以，作为直接对数据进入模型中的关键一步， DataLoader非常重要。

 

DataLoader

　　先介绍一下DataLoader(object)的参数：

• dataset(Dataset): 传入的数据集
• batch_size(int, optional): 每个batch有多少个样本
• shuffle(bool, optional): 在每个epoch开始的时候，对数据进行重新排序
• sampler(Sampler, optional): 自定义从数据集中取样本的策略，如果指定这个参数，那么shuffle必须为False
• batch_sampler(Sampler, optional): 与sampler类似，但是一次只返回一个batch的indices（索引），需要注意的是，一旦指定了这个参数，那么batch_size,shuffle,sampler,drop_last就不能再制定了（互斥——Mutually exclusive）
• num_workers (int, optional): 这个参数决定了有几个进程来处理data loading。0意味着所有的数据都会被load进主进程。（默认为0）
• collate_fn (callable, optional): 将一个list的sample组成一个mini-batch的函数
• pin_memory (bool, optional)： 如果设置为True，那么data loader将会在返回它们之前，将tensors拷贝到CUDA中的固定内存（CUDA pinned memory）中.
• drop_last (bool, optional): 如果设置为True：这个是对最后的未完成的batch来说的，比如你的batch_size设置为64，而一个epoch只有100个样本，那么训练的时候后面的36个就被扔掉了…
• 如果为False（默认），那么会继续正常执行，只是最后的batch_size会小一点。
• timeout(numeric, optional): 如果是正数，表明等待从worker进程中收集一个batch等待的时间，若超出设定的时间还没有收集到，那就不收集这个内容了。这个numeric应总是大于等于0。默认为0
• worker_init_fn (callable, optional): 每个worker初始化函数 If not None, this will be called on each
• worker subprocess with the worker id (an int in [0, num_workers - 1]) as
• input, after seeding and before data loading. (default: None)

 

例子：

from torch.utils.data import DataLoader
from torchvision import transforms,datasets
import matplotlib.pyplot  as plt

def get_datas(batch_size):
    data_tf = transforms.Compose(
        [
            transforms.ToTensor(),
            transforms.Normalize(0,1)
        ]
    )
    train_dataset = datasets.MNIST(root = "./data",train=True,transform=data_tf,download = True)
    train_loader = DataLoader(train_dataset,shuffle =True,batch_size = batch_size,drop_last = True )
    return train_loader

train_loader = get_datas(2)
print(len(train_loader))
for i, data in enumerate(train_loader):
    # f输出两张图片
    if i>0:
        break
    inputs, labels = data
    for j in range(len(inputs)):
        print(inputs[j].shape)
        decode_img = inputs[j].squeeze()  #去掉1
        print(decode_img.shape)
        decode_img = decode_img.data.cpu().numpy() * 255
        plt.imshow(decode_img.astype('uint8'), cmap='gray')
        plt.show()



train_loader = get_datas(2)
print(len(train_loader))
i=0
for imgs, labels in iter(train_loader):
    # f输出两张图片
    if i>0:
        break
    for j in range(len(imgs)):
        decode_img =imgs[j].squeeze()  #去掉1
        print(decode_img.shape)
        decode_img = decode_img.data.cpu().numpy() * 255
        plt.imshow(decode_img.astype('uint8'), cmap='gray')
        plt.show()
    i = i+1