【Pytorch学习】数据集使用

1. 下载数据集

pytorch的官网：https://pytorch.org/，点击菜单栏的Docs中的torchvision，将会看到有许多Package Reference。其中的datasets中有很多数据集，接下来的操作以CIFAR10数据集为例。

具体操作代码如下。download设置为True将会开始在root路径下下载数据集。如果下载过慢，可以将Run端的链接复制一下，自己去迅雷下载。（建议不管之前有没有下载数据集，download都设置为True，因为如果已经存在，不会再去下载一遍了，会提示Files already downloaded and verified）

import torchvision

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

2. 了解数据集

import torchvision

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

print(test_set[0]) 
# (<PIL.Image.Image image mode=RGB size=32x32 at 0x7F6CA0F8C2E8>, 3)
print(test_set.classes) 
# ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']

img, target = test_set[0]
print(img) 
# <PIL.Image.Image image mode=RGB size=32x32 at 0x7F715EDB70F0>
print(target) 
# 3， 对应的是classes中的cat
print(test_set.classes[target]) 
# cat
img.show()

3. 显示数据集

import torchvision
from torch.utils.tensorboard import SummaryWriter

dataset_transform = torchvision.transforms.Compose([torchvision.transforms.ToTensor()])

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

print(test_set[0]) # tensor([...])

# 使用TensorBoard显示
writer = SummaryWriter("CIFAR10")
for i in range(10):
    img, target = test_set[i]
    writer.add_image("test_test", img, i) # i为global step

writer.close()