PyTorch-DataLoader

news2024/11/17 8:49:52

DataLoader：从Dataset中取数据，怎么取，每次取多少可以由DataLoader中的参数进行设定，并将数据加载到神经网络中。

dataloader.py

import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

test_data = torchvision.datasets.CIFAR10('./dataset', train=False, transform=torchvision.transforms.ToTensor())
# 1.batch_size:每次取几个数据
# 2.shuffle:是否要打乱顺序
# 3.num_workers:进程数
# 4.drop_last:如9/4=2...1,若为True则最后一个不要了,若为False则最后一个数据也要。
test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=True, num_workers=0, drop_last=False)
# 测试集中的第一张图片以及target
img, target = test_data[0]
# print(img.shape, target)  # torch.Size([3, 32, 32]) 3
# print(test_data.classes[target])  # cat

writer = SummaryWriter('cifar10')
step = 0
for data in test_loader:
    imgs, targets = data
    writer.add_images('dataloader', imgs, step)
    step = step + 1
# for循环里的imgs是四张图片叠加在一起,所以用add_images
# torch.Size([4, 3, 32, 32]) tensor([0, 6, 0, 4])...
# torch.Size([4,3,32,32])-4张图,3通道,32×32大小的图片
writer.close()

result:

当drop_last为False时，最后一步的情况：

当drop_last为True时，最后一步的情况（只进行到155步）：

进行两轮，且若shuffle为False时：

dataloader.py

import torchvision
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

test_data = torchvision.datasets.CIFAR10('./dataset', train=False, transform=torchvision.transforms.ToTensor())
# 1.batch_size:每次取几个数据集
# 2.shuffle:是否要打乱顺序
# 3.num_workers:进程数
# 4.drop_last:如9/4=2...1,若为True则最后一个不要了,若为False则最后一个数据也要。
test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=False, num_workers=0, drop_last=True)
# 测试集中的第一张图片以及target
img, target = test_data[0]
# print(img.shape, target)  # torch.Size([3, 32, 32]) 3
# print(test_data.classes[target])  # cat

writer = SummaryWriter('cifar10')
for epoch in range(2):
    step = 0
    for data in test_loader:
        imgs, targets = data
        writer.add_images('Epoch:{}'.format(epoch), imgs, step)
        step = step + 1
# for循环里的imgs是四张图片叠加在一起,所以用add_images
# torch.Size([4, 3, 32, 32]) tensor([0, 6, 0, 4])...
# torch.Size([4,3,32,32])-4张图,3通道,32×32大小的图片
writer.close()