【新手入门】PyTorch中的torchvision数据集：一站式数据获取方案

人工智能

2023-09-29 22:20:44

探索 Torchvision 数据集：你的深度学习数据宝藏

踏入深度学习的奇妙世界，数据是至关重要的第一步。在 PyTorch 中，Torchvision 库为你提供了丰富的内置数据集，涵盖广泛的视觉任务，从图像分类到目标检测再到语义分割。本文将带你深入了解这个宝库，帮助你释放深度学习的无限潜力。

Torchvision 数据集一览

Torchvision 提供了 10 个精心挑选的数据集，为你提供高品质且多样化的数据，助力你的深度学习项目：

COCO (Common Objects in Context)： 一个包含超过 20 万张图像的大型数据集，涵盖 80 个目标类别，适用于图像分类和目标检测。
Caltech 101： 包含 101 个物体类别的图像分类数据集，每个类别约有 50 张图像。
CIFAR10： 一个经典的图像分类数据集，包含 10 个物体类别，每个类别有 6000 张图像。
FER2013： 一个表情识别数据集，包含 35887 张面部表情图像，涵盖 7 种表情类别。
iNaturalist： 一个庞大的图像分类数据集，包含超过 40 万张图像，涵盖超过 10000 个物种。
MNIST： 一个手写数字识别数据集，包含 7 万张手写数字图像，分为训练集和测试集。
FashionMNIST： 一个服饰识别数据集，包含 7 万张服饰图像，分为训练集和测试集。
CelebA： 一个包含超过 20 万张名人面部图像的数据集，涵盖 40 个属性。
LSUN： 一个大型场景理解数据集，包含超过 100 万张图像，涵盖多种场景类别。
AVA： 一个动作识别数据集，包含超过 10 万张视频片段，涵盖 80 个动作类别。

如何使用 Torchvision 数据集

使用 Torchvision 数据集很简单，只需几行代码即可：

import torchvision.datasets as datasets

# 加载 CIFAR10 数据集
train_dataset = datasets.CIFAR10(
    root='./data',
    train=True,
    download=True
)

# 加载训练集和测试集
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_dataset = datasets.CIFAR10(
    root='./data',
    train=False,
    download=True
)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=64, shuffle=False)