在 CIFAR10 數(shù)據(jù)集上微調(diào) Vision Transformer (ViT)
在這篇文章中，我們將對預(yù)訓(xùn)練的 Vision Transformer (ViT) 模型進行微調(diào)，以適應(yīng) CIFAR10 數(shù)據(jù)集。在之前的文章《在 CIFAR10 數(shù)據(jù)集上訓(xùn)練 Vision Transformer (ViT)》中，我們從頭開始創(chuàng)建了一個 ViT 模型，并在 CIFAR10 數(shù)據(jù)集上進行了訓(xùn)練。然而，模型的準確率僅達到了67%，沒有進行刻意的超參數(shù)微調(diào)。這是意料之中的，因為 ViT 模型的原始創(chuàng)建者指出，這些模型在小數(shù)據(jù)集上訓(xùn)練時，性能與卷積神經(jīng)網(wǎng)絡(luò)（CNNs）相比是中等的。
然而，當(dāng)在大型數(shù)據(jù)集上進行擴展時，它們開始與 CNNs 相當(dāng)，甚至更好。這就是為什么建議對已經(jīng)在大型數(shù)據(jù)集（如 ImageNet）上預(yù)訓(xùn)練的 ViT 模型進行微調(diào)。而這正是我們在這篇文章中將要做的事情。
訓(xùn)練循環(huán)
我們首先編寫訓(xùn)練和測試任何模型在 CIFAR10 數(shù)據(jù)集上的樣板代碼。您會注意到，我們在訓(xùn)練和測試圖像轉(zhuǎn)換中將圖像大小調(diào)整為224，注意 CIFAR10 的原始圖像大小是32。這是因為我們將要使用的模型需要輸入大小為224，因為它已經(jīng)在 ImageNet 上進行了訓(xùn)練。
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}")加載模型
現(xiàn)在我們必須從 torchvision.models 加載 ViT_b_16 模型。torchvision 中可用的所有 ViT 模型都列在以下鏈接中。如果您查看鏈接，您會發(fā)現(xiàn)有幾個帶有 b、l 和 h 標簽的模型。這些標簽對應(yīng)于我們擁有的基礎(chǔ)、大型和巨型模型大小。這些模型的架構(gòu)正是在第一篇 ViT 論文中發(fā)表的，標題為《 An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale》（原文鏈接：https://arxiv.org/abs/2010.11929）。與這些標簽相關(guān)的數(shù)字，如16、32和14，對應(yīng)于模型使用的補丁大小。所有這些模型都已在 ImageNet 上進行了訓(xùn)練。我們首先加載模型。默認提供的模型不是預(yù)訓(xùn)練的，為了確保我們加載一個預(yù)訓(xùn)練的模型，我們必須將權(quán)重參數(shù)傳遞為 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個類別的對數(shù)幾率，因為它已經(jīng)在 ImageNet 上進行了訓(xùn)練。然而，我們的數(shù)據(jù)集只包含10個類別。因此，我們需要將這個模型的頭部從1000個對數(shù)幾率更改為10個。加載模型的外層是“heads”層，這是一個序列層，其中只包含一個線性層。為了適應(yīng)模型，我們只需在保留層的輸入特征的同時，將一個新的線性層分配給“heads”層，并將外部特征替換為10。
model = vit_b_16(ViT_B_16_Weights.IMAGENET1K_V1)
model.heads = nn.Sequential(
nn.Linear(model.heads.head.in_features, 10)
)我們不是訓(xùn)練或加載模型中的變換器塊，我們可以凍結(jié)所有層，除了最后一個變換器層。通過這樣做，我們使微調(diào)過程的計算強度降低。我們最后將模型移動到 GPU 設(shè)備，并使用之前的訓(xùn)練循環(huán)進行訓(xù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 的測試數(shù)據(jù)集上測試我們的模型。您會發(fā)現(xiàn)，即使只訓(xùn)練了一個周期，模型也達到了非常高的準確率。這是因為在模型在 ImageNet 上訓(xùn)練時所打造的強大的特征。
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}%")
