PyTorch学习系列——加载数据2

Posted 2022-11-26 Vic时代

生成batch数据

现在有了由数据文件生成的结构数据MyDataset，那么怎么在训练时提供batch数据呢？PyTorch提供了生成batch数据的类。

PyTorch用类torch.utils.data.DataLoader加载数据，并对数据进行采样，生成batch迭代器。

class torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, sampler=None, num_workers=0, collate_fn=<function default_collate>, pin_memory=False, drop_last=False)

参数：
dataset：Dataset类型，从其中加载数据
batch_size：int，可选。每个batch加载多少样本
shuffle：bool，可选。为True时表示每个epoch都对数据进行洗牌
sampler：Sampler，可选。从数据集中采样样本的方法。
num_workers：int，可选。加载数据时使用多少子进程。默认值为0，表示在主进程中加载数据。
collate_fn：callable，可选。
pin_memory：bool，可选
drop_last：bool，可选。True表示如果最后剩下不完全的batch,丢弃。False表示不丢弃。

示例：

kwargs = 'num_workers': 1, 'pin_memory': True if args.cuda else 
train_loader = torch.utils.data.DataLoader(
    MyDataset(images, labels), batch_size=args.batch_size, shuffle=True, **kwargs)

其他用法：
len(train_loader) ：返回的是len(dataset)/batch_size