PyTorch快速入门教程【小土堆】之DataLoader的使用

EnochChen_2024-12-30 13:33

dataset数据集，相当于一副扑克，dataloader数据加载器相当于我们的手，选择摸几张牌，怎么摸牌

python 复制代码

import torchvision

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

test_data = torchvision.datasets.CIFAR10(root="./dataset", train=False, transform=torchvision.transforms.ToTensor()
                                         )  # root为数据集放置位置，train为true说明训练集，transform为ToTensor方法，将图片转化为tensor格式

test_loader = DataLoader(dataset=test_data, batch_size=4, shuffle=True, num_workers=0,
                         drop_last=False)  # 每批次取4个数据，shuffle为true说明每次选取的数据要相同，num_workers为0说明只选择主进程加载，多个进程会加载更快，但在window下>0有可能会报错，取数据时有可能最后余下几个drop_last为true，则舍去

# 测试数据集中第一张图片及target
img, target = test_data[0]
print(img.shape)  # torch.Size([3, 32, 32])
print(target)  # 3

writer = SummaryWriter("dataloader")
step = 0
for data in test_loader:
    imgs, targets = data
    # print (imgs.shape) # torch.Size([4 3, 32, 32]),4为batch_size
    # print(targets) # tensor([2, 3, 6, 8]),4张图片的target进行了打包
    writer.add_images("test_data", imgs, step)
    step += 1

writer.close()