在CIFAR10 数据集上微调 Vision Transformer (ViT)

点击下方卡片，关注“小白玩转Python”公众号

在这篇文章中，我们将对预训练的 Vision Transformer (ViT) 模型进行微调，以适应 CIFAR10 数据集。在昨天的文章《在 CIFAR10 数据集上训练 Vision Transformer (ViT)》中，我们从头开始创建了一个 ViT 模型，并在 CIFAR10 数据集上进行了训练。然而，模型的准确率仅达到了67%，没有进行刻意的超参数微调。这是意料之中的，因为 ViT 模型的原始创建者指出，这些模型在小数据集上训练时，性能与卷积神经网络（CNNs）相比是中等的。然而，当在大型数据集上进行扩展时，它们开始与 CNNs 相当，甚至更好。这就是为什么建议对已经在大型数据集（如 ImageNet）上预训练的 ViT 模型进行微调。而这正是我们在这篇文章中将要做的事情。

训练循环

我们首先编写训练和测试任何模型在 CIFAR10 数据集上的样板代码。您会注意到，我们在训练和测试图像转换中将图像大小调整为224，注意 CIFAR10 的原始图像大小是32。这是因为我们将要使用的模型需要输入大小为224，因为它已经在 ImageNet 上进行了训练。

transform_train = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])


transform_test = transforms.Compose([
    transforms.Resize(224),
    transforms.ToTensor(),
    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
])


train_set = CIFAR10(root='./datasets', train=True, download=True, transform=transform_train)
test_set = CIFAR10(root='./datasets', train=False, download=True, transform=transform_test)


train_loader = DataLoader(train_set, shuffle=True, batch_size=64)
test_loader = DataLoader(test_set, shuffle=False, batch_size=64)

n_epochs = 10
lr = 0.0001


optimizer = Adam(model.parameters(), lr=lr)
criterion = CrossEntropyLoss()


for epoch in range(n_epochs):
    train_loss = 0.0
    for i,batch in enumerate(train_loader):
        x, y = batch
        x, y = x.to(device), y.to(device)
        y_hat = model(x)
        loss = criterion(y_hat, y)


        batch_loss = loss.detach().cpu().item()
        train_loss += batch_loss / len(train_loader)


        optimizer.zero_grad()
        loss.backward()
        optimizer.step()


        if i%100==0:
          print(f"Batch {i}/{len(train_loader)} loss: {batch_loss:.03f}")


    print(f"Epoch {epoch + 1}/{n_epochs} loss: {train_loss:.03f}")

加载模型

现在我们必须从 torchvision.models 加载 ViT_b_16 模型。torchvision 中可用的所有 ViT 模型都列在以下链接中。如果您查看链接，您会发现有几个带有 b、l 和 h 标签的模型。这些标签对应于我们拥有的基础、大型和巨型模型大小。这些模型的架构正是在第一篇 ViT 论文中发表的，标题为《 An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale》（原文链接：https://arxiv.org/abs/2010.11929）。与这些标签相关的数字，如16、32和14，对应于模型使用的补丁大小。所有这些模型都已在 ImageNet 上进行了训练。我们首先加载模型。默认提供的模型不是预训练的，为了确保我们加载一个预训练的模型，我们必须将权重参数传递为 ViT_B_16_Weights.IMAGENET1K_V1。

from torchvision.models import ViT_B_16_Weights, vit_b_16


model = vit_b_16(ViT_B_16_Weights.IMAGENET1K_V1)

默认情况下，这个模型输出来自1000个类别的对数几率，因为它已经在 ImageNet 上进行了训练。然而，我们的数据集只包含10个类别。因此，我们需要将这个模型的头部从1000个对数几率更改为10个。加载模型的外层是“heads”层，这是一个序列层，其中只包含一个线性层。为了适应模型，我们只需在保留层的输入特征的同时，将一个新的线性层分配给“heads”层，并将外部特征替换为10。

model = vit_b_16(ViT_B_16_Weights.IMAGENET1K_V1)


model.heads = nn.Sequential(
    nn.Linear(model.heads.head.in_features, 10)
)

我们不是训练或加载模型中的变换器块，我们可以冻结所有层，除了最后一个变换器层。通过这样做，我们使微调过程的计算强度降低。我们最后将模型移动到 GPU 设备，并使用之前的训练循环进行训练。

model = vit_b_16(ViT_B_16_Weights.IMAGENET1K_V1)


model.heads = nn.Sequential(
    nn.Linear(model.heads.head.in_features, 10)
)


# Freeze all layers
for param in model.parameters():
    param.requires_grad = False


# Unfreeze the last encoder layer and the head
for param in model.encoder.layers[-1].parameters():
    param.requires_grad = True
for param in model.heads.parameters():
    param.requires_grad = True

测试循环

我们最后在 CIFAR10 的测试数据集上测试我们的模型。您会发现，即使只训练了一个周期，模型也达到了非常高的准确率。这是因为在模型在 ImageNet 上训练时所打造的强大的特征。

with torch.no_grad():
    correct, total = 0, 0
    test_loss = 0.0
    for batch in tqdm(test_loader, desc="Testing"):
        x, y = batch
        x, y = x.to(device), y.to(device)
        y_hat = model(x)
        loss = criterion(y_hat, y)
        test_loss += loss.detach().cpu().item() / len(test_loader)


        correct += torch.sum(torch.argmax(y_hat, dim=1) == y).detach().cpu().item()
        total += len(x)
    print(f"Test loss: {test_loss:.2f}")
    print(f"Test accuracy: {correct / total * 100:.2f}%")

·  END  ·

🌟 想要变身计算机视觉小能手？快来「小白玩转Python」公众号！

回复“Python视觉实战项目”，解锁31个超有趣的视觉项目大礼包！🎁

本文仅供学习交流使用，如有侵权请联系作者删除