【Pytorch】torchvision中的数据集使用 - 学习笔记

视频网址

1 - 加载数据集：以CIFAR为例

网址

这里可以看看数据集的内容

1）如果没有下载数据集（不推荐）

import torchvision

train_set = torchvision.datasets.CIFAR10(root="./dataset", train=True, download=True)
test_set = torchvision.datasets.CIFAR10(root="./dataset", train=False, download=True)

print(test_set[0])

下载非常龟速…所以建议用迅雷下，在运行上面代码后可以发现输出一个网址https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz，在迅雷中新建下载任务，把这个链接输入进去，几秒钟搞定！

2）已经下好了数据集
把数据集解压到dataset文件夹中，然后运行

import torchvision

train_set = torchvision.datasets.CIFAR10(root="./dataset", train=True, download=False)
test_set = torchvision.datasets.CIFAR10(root="./dataset", train=False, download=False)

print(test_set[0])

输出结果为

(<PIL.Image.Image image mode=RGB size=32x32 at 0x29625500BB0>, 3)

为了更好看到结果，我们设置一下断点（鼠标左键）

我们来看看test_set[0]输出的两个参数：img和target分别是什么，接着上面的代码，输入

img, target = test_set[0]
print(img)
print(target) # label
print(test_set.class[target]) # 看看输出是什么类型
img.show() # 看看图片长啥样

输出

D:\Anaconda3\envs\pytorch\python.exe D:/研究生/代码尝试/P10_dataset_transform.py
(<PIL.Image.Image image mode=RGB size=32x32 at 0x10CDF470BB0>, 3)
['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
<PIL.Image.Image image mode=RGB size=32x32 at 0x10CDF470BB0>
3
cat

进程已结束，退出代码为 0

隐隐约约可以看出来是猫

2 - 使用Transform

import torchvision
from torch.utils.tensorboard import SummaryWriter

# 设置transform, PIL -> Tensor
dataset_transform = torchvision.transforms.Compose([
    torchvision.transforms.ToTensor()
])

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)

# print(test_set[0])
# print(test_set.classes)
#
# img, target = test_set[0]
# print(img)
# print(target)
# print(test_set.classes[target])
# img.show()

# 使用Tensorboard显示
writer = SummaryWriter("p10")
for i in range(10):
    img, target = test_set[i]
    writer.add_image("test_set", img, i)

writer.close()

打开tensorboard，显示

通过左右移动，可以看到不同的step的图片不一样~

Tips：如果数据集没有给下载地址，可以在dataset点击ctrl+鼠标左键，找到源码，查看下载地址~