torchvision中的数据使用

兮℡檬，2025-07-04 21:10

1、下载数据集

在pytorch官网中找到docs选择Domains，在该页面中有各种数据类型的数据集

在左边菜单栏中选择datasets

python 复制代码

import torchvision
train_set=torchvision.datasets.CIFAR10(root='/data',train=True,download=True)
test_set=torchvision.datasets.CIFAR10(root='./data',train=False,download=True)

2、Dataloader的使用

Dataloader参数介绍

dataset ：加载的数据集，必须是 torch.utils.data.Dataset 的子类实例。
batch_size：每个批次的数据样本数，默认值为1。
shuffle：是否在每个周期开始时打乱数据，默认为 False。
sampler：定义从数据集中抽取样本的策略，如果指定，则忽略 shuffle 参数。
num_workers：用于数据加载的子进程数量，默认为0，表示数据将在主进程中加载。
collate_fn：如何将多个数据样本整合成一个批次，通常不需要指定。
pin_memory：如果为 True，会将数据放置到 GPU 上去，默认为 False。
drop_last：如果数据集大小不能被批次大小整除，是否丢弃最后一个不完整的批次，默认为 False。

python 复制代码

test_loader=DataLoader(dataset=test_set,batch_size=4,shuffle=True,num_workers=0,drop_last=False)
#获取一张图片的信息
img,target=test_set[0]
print(img.shape)
print(target)

writer=SummaryWriter("dataloader")
#taet_loader是一个迭代对象，用for循环进行迭代
step=0
for data in test_loader:
    imgs,targets=data
    # print(imgs.shape)
    # print(targets)
    writer.add_image("test_data",imgs,step,dataformats='NCHW')
    step+=1

writer.close()

添加轮次

python 复制代码

for epoch in range(2):
    step=0
    for data in test_loader:
        imgs,targets=data
        # print(imgs.shape)
        # print(targets)
        writer.add_image("Epoch:{}".format(epoch),imgs,step,dataformats='NCHW')
        step+=1