PyTorch学习系列(一)——加载数据2

生成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

猜你喜欢

转载自blog.csdn.net/VictoriaW/article/details/80314497