在 CIFAR10 數據集上微調 Vision Transformer (ViT)
在這篇文章中，我們將對預訓練的 Vision Transformer (ViT) 模型進行微調，以適應 CIFAR10 數據集。在之前的文章《在 CIFAR10 數據集上訓練 Vision Transformer (ViT)》中，我們從頭開始創(chuàng)建了一個 ViT 模型，并在 CIFAR10 數據集上進行了訓練。然而，模型的準確率僅達到了67%，沒有進行刻意的超參數微調。這是意料之中的，因為 ViT 模型的原始創(chuàng)建者指出，這些模型在小數據集上訓練時，性能與卷積神經網絡（CNNs）相比是中等的。
然而，當在大型數據集上進行擴展時，它們開始與 CNNs 相當，甚至更好。這就是為什么建議對已經在大型數據集（如 ImageNet）上預訓練的 ViT 模型進行微調。而這正是我們在這篇文章中將要做的事情。
訓練循環(huán)
我們首先編寫訓練和測試任何模型在 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 模型都列在以下鏈接中。如果您查看鏈接，您會發(fā)現有幾個帶有 b、l 和 h 標簽的模型。這些標簽對應于我們擁有的基礎、大型和巨型模型大小。這些模型的架構正是在第一篇 ViT 論文中發(fā)表的，標題為《 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 設備，并使用之前的訓練循環(huán)進行訓練。
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測試循環(huán)
我們最后在 CIFAR10 的測試數據集上測試我們的模型。您會發(fā)現，即使只訓練了一個周期，模型也達到了非常高的準確率。這是因為在模型在 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}%")
