PyTorch④---DataLoader的使用

DataLoader的使用

送入神经网络时，一般用DataLoader进行打包。

import torchvision
#准备的测试数据集
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

test_data=torchvision.datasets.CIFAR10("dataset_CIFAR10",train=False,transform=torchvision.transforms.ToTensor())
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)
print(target)
''' dataset return img,target '''
''' batch_size=4: 64个64个进行打包 从dataset中取四个数据,img0,target0=dataset[0] img1,target1=dataset[1] img2,target2=dataset[2] img3,target3=dataset[3] 把img0,img1,img2,img3进行打包,返回imgs 把target0,1,2,3进行打包,返回targets 用for遍历获取 drop_last=False:不够64个会留下，True不够64个会舍去 '''
writer=SummaryWriter("dataloader")
step=0
for data in test_loader:
    imgs,targets=data
    #print(imgs.shape)
    #print(targets)
    writer.add_images("test_data",imgs,step)
    step=step+1

''' shuffle=True:循环两轮的图片不一样 shuffle=False:循环两轮的图片一样 '''
for epoch in range(2):
    step = 0
    for data in test_loader:
        imgs, targets = data
        # print(imgs.shape)
        # print(targets)
        writer.add_images("epoch:{}".format(epoch), imgs, step)
        step = step + 1
writer.close()