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