pytorch——torchvision中的数据集的使用

首先我们打开pytorch的官方网页

比如我们打开其中的torchvision，我们可以看到其中有很多的数据集：

 其中有很多是我们常用的、用来学习的数据集。接着的torchvision.models是一些常见的神经网络：

---

我们以CIFAR为例：

root是存在或下载保存的路径；train如果是True，则作为训练集，否则是测试集；transform可以对数据集进行一个变化； download会设置下载数据集。

import torchvision

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

查看test_set第一个元素：

>>>test_set[0]
Out[3]: (<PIL.Image.Image image mode=RGB size=32x32 at 0x29103869F48>, 3)  

知道，元素是元组，并且含有的图片是PIL类型。那么我们可以利用一个compose函数：

import torchvision

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

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

print(test_set[0])

 输出结果为一个tensor类型，如此，我们可以使用Tensorboard进行一个可视化：

import torchvision
from torch.utils.tensorboard import SummaryWriter

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

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

writer = SummaryWriter("logs")
for i in range(10):
    img, target = test_set[i]
    writer.add_image("pic", img, i)

writer.close()

打开TensorBoard可以看到：