Torchvision数据集下载与使用

1.单张图片处理

import torchvision

train_set = torchvision.datasets.CIFAR10(root="./dataset", train=True, download=False)
test_set = torchvision.datasets.CIFAR10(root="./dataset", train=False, download=False)
img, label = test_set[4] # img是图片的信息, label是标签信息(用整数表示)
print(img)
print(label, test_set.classes[label]) # classes保存所有标签的信息,classes[label]返回真实标签
print(test_set.classes[label])
img.show()

2. 批量处理(对整个数据集)

import torchvision
from torch.utils.tensorboard import SummaryWriter

# 定义数据转换的操作流,应用在接下来的数据集中,实现对所有图片的转换,而不再需要每张图片逐一转换
dataset_transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor()
])

# 直接使用数据集的包torchvision.datasets.CIFAR10
# root      指定数据集的目录
# train     True时代表训练集,False时代表训练集
# transform 图片转换流
train_set = torchvision.datasets.CIFAR10(root="./dataset", train=True, transform=dataset_transform, download=False)
test_set = torchvision.datasets.CIFAR10(root="./dataset", train=False, transform=dataset_transform, download=False)

# Tensorboard显示
writer = SummaryWriter("logs")
for i in range(10):
    img, label = train_set[i]
    writer.add_image("train_set", img, i)

writer.close()

注意:cifar-10-batches-py一定要在代码的指定路径下，放入解压后的cifar-10-batches-py文件(上面的代码指定的路径是"./dataset")