How do I improve the Accuracy of the Vision Transformer Model? (Code Review)

[Kart1290]

Hi everyone, I followed the general steps in Section 11 of the course for deployment but instead of efficient net I used a vision transformer. However, I cannot seem to push the model accuracy beyond 75%
The things I have tried include:

1. Amount of data - 20%, 30%, 50%
2. Learning Rate - 1e-3, 5e-4, 1e-5
3. Epochs - 5, 10, 20, 40
4. For regularization I used: Label Smoothing, Weight Decay and TrivialAugmentWide (for the training data)
   I have shared the relevant parts of the code below, was hoping someone could point out the error or any misunderstandings I may have.

Thanks :)

def create_vit_model(num_classes: int) -> Tuple[torch.nn.Module, torchvision.transforms.Compose]:
  vit_weights = ViT_B_16_Weights.DEFAULT
  vit_model = vit_b_16(weights=vit_weights)
  vit_transforms = vit_weights.transforms()

  for p in vit_model.parameters():
    p.requires_grad = False

  vit_model.heads = torch.nn.Sequential(
      torch.nn.Linear(in_features=768,
                      out_features=num_classes)
  )
  return (vit_model, vit_transforms)

vit_model, vit_transforms = create_vit_model(101)
vit_model = vit_model.to(device)

train_transform = Compose(transforms=[
    torchvision.transforms.TrivialAugmentWide(),
    vit_transforms])

train_dataset = Food101(root='data',
                        split='train',
                        transform=train_transform,
                        download=True)

test_dataset = Food101(root='data',
                        split='test',
                        transform=vit_transforms,
                        download=True)


class_names = train_dataset.classes

train_labels = [train_dataset[i][1] for i in range(len(train_dataset))] # IS THERE A WAY TO IMPROVE THIS???
test_labels = [test_dataset[i][1] for i in range(len(test_dataset))]

train_indices, _, _, _ = train_test_split(range(len(train_dataset)),
                                    train_labels,
                                    test_size=0.5, # taking 50% of the data
                                    shuffle=True,
                                    stratify=train_labels)

test_indices, _, _, _, = train_test_split(range(len(test_dataset)),
                                   test_labels,
                                   test_size=0.5, # taking 50% of the data
                                   shuffle=True,
                                   stratify=test_labels)

train_dataset_subset = Subset(train_dataset, train_indices)
test_dataset_subset = Subset(test_dataset, test_indices)

BATCH_SIZE = 32
train_dataloader = DataLoader(train_dataset_subset, BATCH_SIZE, shuffle=True, num_workers=os.cpu_count(), pin_memory=True) # pin memory will optimize the transfer from cpu to gpu
test_dataloader = DataLoader(test_dataset_subset, BATCH_SIZE, shuffle=False, num_workers=os.cpu_count(), pin_memory=True)

def train_step(model: torch.nn.Module,
               dataloader: torch.utils.data.DataLoader,
               optimizer: torch.optim.Optimizer,
               loss_fn: torch.nn.Module,
               device: str)-> Tuple[float, float]:
  # Put the model in train
  model.train()
  # track loss and accuracy
  train_loss, correct_preds, total_preds = 0.0, 0, 0
  # Iterate through the data
  for X, y in dataloader:
    # Move data to device
    X, y = X.to(device), y.to(device)
    # Forward pass
    y_logits = model(X)
    # Calculate the loss and accuracy
    loss = loss_fn(y_logits, y)
    train_loss += loss.item()
    correct_preds += torch.sum(torch.softmax(y_logits, dim=-1).argmax(dim=-1) == y).item()
    total_preds += len(y)
    # Zero grad
    optimizer.zero_grad()
    # Backpropagation
    loss.backward()
    # Step
    optimizer.step()
  # Epoch loss and accuracy
  train_loss /= len(dataloader)
  acc = (correct_preds / total_preds) * 100

  return (train_loss, acc)

def test_step(model: torch.nn.Module,
              dataloader: torch.utils.data.DataLoader,
              loss_fn: torch.nn.Module,
              device: str) -> Tuple[float, float]:
  # track loss and acc per epoch
  test_loss, correct_preds, total_preds = 0.0, 0, 0
  # evaluation
  model.eval() # disables drop out and uses trained norm
  with torch.inference_mode(): # mem optim, grad tracking off and read only
    for X, y in dataloader:
      # Move data to device
      X, y = X.to(device), y.to(device)
      # forward pass
      y_logits = model(X)
      # calculate the loss and accuracy
      loss = loss_fn(y_logits, y)
      test_loss += loss.item()
      correct_preds += torch.sum(torch.softmax(y_logits, dim=-1).argmax(dim=-1) == y).item()
      total_preds += len(y)
  test_loss /= len(dataloader)
  acc = (correct_preds / total_preds) * 100

  return (test_loss, acc)

from tqdm.auto import tqdm

def train(model: torch.nn.Module,
          train_dataloader: torch.utils.data.DataLoader,
          test_dataloader: torch.utils.data.DataLoader,
          optimizer: torch.optim.Optimizer,
          loss_fn: torch.nn.Module,
          epochs: int,
          device: str):
  # results
  results = []
  for e in tqdm(range(epochs)):
    epoch_result = {}
    train_loss_per_epoch, train_acc_per_epoch = train_step(model, train_dataloader, optimizer, loss_fn, device)
    test_loss_per_epoch, test_acc_per_epoch = test_step(model, test_dataloader, loss_fn, device)
    epoch_result['epoch'] = e+1
    epoch_result['train_loss'] = train_loss_per_epoch
    epoch_result['train_acc'] = train_acc_per_epoch
    epoch_result['test_loss'] = test_loss_per_epoch
    epoch_result['test_acc'] = test_acc_per_epoch
    print(f'[INFO] Epoch: {e+1} | train_loss: {train_loss_per_epoch: .4f} | train_acc: {train_acc_per_epoch: .4f} | test_loss: {test_loss_per_epoch: .4f} | test_acc: {test_acc_per_epoch: .4f}')
    results.append(epoch_result)
  return results