vits-fast-fine-tuning google colab

Sure! Here are the steps to fine-tune ViT-S on a custom dataset using Google Colab:

1. Open a new Google Colab notebook and select a GPU runtime environment.

2. Install the necessary libraries:

!pip install torch torchvision
!pip install timm

3. Download and prepare the custom dataset. You can use any dataset of your choice. Make sure to split it into training and validation sets.

4. Define the data loaders:

import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder

# Define the transformations
transform_train = transforms.Compose([
    transforms.RandomResizedCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

transform_val = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Define the data loaders
train_dataset = ImageFolder('path_to_train_data', transform=transform_train)
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)

val_dataset = ImageFolder('path_to_val_data', transform=transform_val)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

Replace 'path_to_train_data' and 'path_to_val_data' with the paths to your training and validation data folders, respectively.

5. Load the pre-trained ViT-S model:

import timm

model = timm.create_model('vit_small_patch16_224', pretrained=True)

6. Modify the last layer of the model to fit your custom dataset:

import torch.nn as nn

num_classes = len(train_dataset.classes)

model.head = nn.Sequential(
    nn.LayerNorm((768,)),
    nn.Linear(768, num_classes)
)

Replace '768' with the hidden size of the model you are using. For ViT-S, it is 768.

7. Define the optimizer and criterion:

import torch.optim as optim

optimizer = optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()

8. Fine-tune the model:

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

num_epochs = 10

for epoch in range(num_epochs):
    train_loss = 0.0
    val_loss = 0.0
    correct = 0
    total = 0

    # Train the model
    model.train()
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)

        optimizer.zero_grad()

        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        train_loss += loss.item() * inputs.size(0)

    # Evaluate the model on validation set
    model.eval()
    with torch.no_grad():
        for inputs, labels in val_loader:
            inputs, labels = inputs.to(device), labels.to(device)

            outputs = model(inputs)
            loss = criterion(outputs, labels)

            val_loss += loss.item() * inputs.size(0)

            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

    train_loss = train_loss / len(train_loader.dataset)
    val_loss = val_loss / len(val_loader.dataset)
    accuracy = 100 * correct / total

    print('Epoch: {} \tTraining Loss: {:.6f} \tValidation Loss: {:.6f} \tAccuracy: {:.2f}'.format(
        epoch+1, train_loss, val_loss, accuracy))

9. Save the model:

torch.save(model.state_dict(), 'path_to_save_model')

Replace 'path_to_save_model' with the path where you want to save the model.

That's it! You have fine-tuned ViT-S on your custom dataset using Google Colab.