您的位置:

TinyImageNet——一个用于视觉分类的挑战性数据集

一、数据集介绍

TinyImageNet是一个用于视觉分类的挑战性数据集,由Stanford CS231n课程所提供。数据集包含200个类别,每个类别有500张训练图像、50张验证图像和50张测试图像,大小为64×64像素。图像来自ImageNet数据集,但是通常使用更小的版本,通常可以通过随机裁剪、水平翻转等数据增强技术进行处理。

数据集类别非常多样化,包括常见的动物类别如“狗(Dog)”、“猫(Cat)”,以及不太常见的类别如“蘑菇(Mushroom)”、“校车(School Bus)”等。此外,数据集还包含人工标签类别,如“裙子(Dress)”、“钢笔(Pen)”等。

import torch
import torchvision

data_transforms = {
    'train': torchvision.transforms.Compose([
        torchvision.transforms.RandomResizedCrop(64),
        torchvision.transforms.RandomHorizontalFlip(),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': torchvision.transforms.Compose([
        torchvision.transforms.Resize(64),
        torchvision.transforms.CenterCrop(64),
        torchvision.transforms.ToTensor(),
        torchvision.transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

train_dataset = torchvision.datasets.ImageFolder(root='tiny-imagenet-200/train', transform=data_transforms['train'])
val_dataset = torchvision.datasets.ImageFolder(root='tiny-imagenet-200/val', transform=data_transforms['val'])
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=2)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=128, shuffle=False, num_workers=2)

上面的代码展示了如何使用PyTorch中的torchvision库加载数据集和进行数据增强。在训练集上,我们使用了随机裁剪和水平翻转,而在验证集上则使用了中心裁剪。同时使用了标准化技术来使每个像素的值归一化到0-1之间。使用DataLoader来快速加载数据并自动处理数据的拆分和批处理。

二、经典模型在TinyImageNet上的应用

许多经典的深度学习模型都已经在不同的数据集上进行了预训练和微调,并在TinyImageNet上取得了较好的表现。下面是几个在TinyImageNet上表现较好的模型:

VGG

在TinyImageNet上,VGG表现出色。最好的模型使用VGG16,在测试集上获得了50.88%的准确率。关于VGG的更多信息可以参考K. Simonyan, A. Zisserman. Very Deep Convolutional Networks for Large-Scale Image Recognition, 2014。

import torch.nn as nn
import torchvision.models as models

model = models.vgg16_bn(pretrained=True)
for param in model.parameters():
    param.requires_grad = False

num_features = model.classifier[-1].in_features
features = list(model.classifier.children())[:-1]
features.extend([nn.Linear(num_features, 200)])
model.classifier = nn.Sequential(*features)

上面的代码展示了如何在PyTorch中使用VGG进行微调,这里选择了VGG16模型,并使用预训练的权重进行初始化。通过修改VGG最后一层的输出数量,使其与TinyImageNet的类别数量相匹配。接下来,我们将冻结前面的层,并只对新添加的层进行训练。

ResNet

ResNet也是一种流行的深度学习模型,它在TinyImageNet上的表现也很不错。最好的模型使用ResNet-152,在测试集上获得了51.05%的准确率。关于ResNet的更多信息可以参考K. He, X. Zhang, S. Ren, J. Sun. Deep Residual Learning for Image Recognition, 2016。

import torch.nn as nn
import torchvision.models as models

model = models.resnet152(pretrained=True)
for param in model.parameters():
    param.requires_grad = False

num_features = model.fc.in_features
model.fc = nn.Linear(num_features, 200)

上面的代码展示了如何在PyTorch中使用ResNet进行微调,这里选择了ResNet-152模型,并使用预训练的权重进行初始化。通过修改最后一层的输出数量,使其与TinyImageNet的类别数量相匹配。接下来,我们将冻结前面的层,并只对新添加的层进行训练。

DenseNet

相对于VGG和ResNet,DenseNet是一种比较新的模型结构,但是在TinyImageNet上也表现很好。最好的模型使用DenseNet-201,在测试集上获得了51.72%的准确率。关于DenseNet的更多信息可以参考G. Huang, Z. Liu, K. Q. Weinberger, L. van der Maaten. Densely Connected Convolutional Networks, 2016。

import torch.nn as nn
import torchvision.models as models

model = models.densenet201(pretrained=True)
for param in model.parameters():
    param.requires_grad = False

num_features = model.classifier.in_features
model.classifier = nn.Linear(num_features, 200)

上面的代码展示了如何在PyTorch中使用DenseNet进行微调,这里选择了DenseNet-201模型,并使用预训练的权重进行初始化。通过修改最后一层的输出数量,使其与TinyImageNet的类别数量相匹配。接下来,我们将冻结前面的层,并只对新添加的层进行训练。

三、小结

TinyImageNet是一个用于视觉分类的挑战性数据集,数据集类别多样化,具有一定的难度。本文介绍了如何在PyTorch中使用常用的经典模型在TinyImageNet上进行微调。希望此文能对使用PyTorch进行深度学习开发的读者有所帮助。