Skin Lesion Classification with ViT

Train ConvNeXt and ViT models to 92%+ accuracy on ISIC 2019 skin lesions.

Introduction

30 Second Summary

Doctors examining skin lesions need to distinguish between 8 different types, some of which are incredibly rare. A model that gets 86% accuracy sounds good on paper, but in practice it means the rarest conditions get misclassified almost every time.

In this project, you will train two deep learning models to individually surpass 92% accuracy on the ISIC 2019 skin lesion dataset. You will fix the training mistakes that cap your current model at 86% and build a reusable pipeline that works within a 4GB GPU.

What You'll Build

You will run both a ConvNeXt and a Vision Transformer (ViT) against all 8 skin lesion types and watch them correctly identify even the rarest conditions that your baseline model struggles with.

By the end of this project, you'll have:

  • A balanced data pipeline where you can print a training batch and see every class represented fairly, despite the original dataset having a 51:1 imbalance between the most and least common classes.
  • A ConvNeXt Tiny model trained at 384x384 resolution with layer freezing and mixed precision that you can evaluate and see break past 92% accuracy on the test set.
  • A ViT-Small model trained with the same strategy, proving to yourself that the approach generalises across completely different architectures.
  • A side-by-side comparison showing per-class accuracy improvements over your original 86% baseline, revealing which classes each model handles best.
  • Secret Mission: Add Test-Time Augmentation (TTA) to push accuracy even higher without any retraining.

Are there any prerequisites?

This project assumes you already have the ISIC 2019 dataset downloaded and split into train/val/test CSVs with class weights computed from a previous project. You also need a GPU that supports CUDA (this project uses an RTX 3050 with 4GB VRAM).

You should have Python, PyTorch, torchvision, and Jupyter Notebook already installed and working.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters.

Your current ConvNeXt model reaches 86% accuracy on ISIC 2019. In this project, you will train two models (ConvNeXt Tiny and ViT-Small) past 92% by fixing the three things holding that baseline back: training all layers at the same learning rate, using 224x224 resolution, and having no oversampling for rare classes.

Set Up Your Environment

Your current ConvNeXt model tops out at 86% accuracy because it trains all layers at the same rate, uses 224x224 images, and never properly oversamples minority classes.

To break past 92%, you need a better training pipeline and a second model architecture. This step gets your environment ready by installing the timm library (which gives you access to ViT models), importing everything you need, and loading your existing ISIC 2019 preprocessing artifacts.

In this step, get ready to:
  • Install the timm library and verify GPU mixed precision support.
  • Import all libraries needed for the training pipeline.
  • Load your dataset splits, class weights, and label mapping from previous work.
Install the timm library

The timm (PyTorch Image Models) library provides access to hundreds of pretrained vision models, including the ViT-Small architecture you will train in later steps.

  • In a new Jupyter cell, install timm by running:
!pip install timm

What does this command do?

The ! prefix runs a shell command from inside a Jupyter cell. This installs the timm package from PyPI into your current Python environment.

Before you run the next cell, do you think your GPU will support mixed precision training?

  • Verify the installation and check GPU support by running this in a new cell:
import timm
import torch
print(f"timm version: {timm.__version__}")
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"Mixed Precision supported: {torch.cuda.is_available()}")

You should see your timm version number, your GPU name (NVIDIA GeForce RTX 3050), and True for mixed precision support.

Not seeing your GPU?

If torch.cuda.is_available() returns False, your PyTorch installation may not include CUDA support. Reinstall PyTorch with the correct CUDA version from the official PyTorch website.

If !pip install timm fails, check your internet connection and try again.

help me fix my GPU or timm installation

Import all libraries for the project

This project uses PyTorch for model training, torchvision for transforms and the ConvNeXt model, and timm for the ViT model.

  • Import all required libraries by running this cell:
import os
import copy
import json
import random
import time
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
from torch.amp import autocast, GradScaler
from torchvision import transforms
from torchvision.models import convnext_tiny, ConvNeXt_Tiny_Weights
import timm
from PIL import Image
from sklearn.metrics import classification_report
import matplotlib.pyplot as plt

What do these imports cover?

  • torch.amp provides mixed precision training (FP16) which halves GPU memory usage.
  • WeightedRandomSampler handles class imbalance by oversampling rare lesion types.
  • GradScaler prevents underflow when gradients are computed in half precision.
  • convnext_tiny and timm give you access to both model architectures you will train.
  • Run the cell. If no errors appear, all libraries are installed correctly.

Seeing an ImportError?

If from torch.amp import autocast, GradScaler fails, you may be on an older PyTorch version. This import path requires PyTorch 2.0 or newer.

If from sklearn.metrics import classification_report fails, install scikit-learn with !pip install scikit-learn in a new cell.

help me fix my import errors in Jupyter

Load your preprocessing artifacts

You already have the ISIC 2019 dataset split into train, validation, and test sets with precomputed class weights from your previous project. This cell loads those artifacts and sets the configuration constants for training.

  • Set up your configuration and load the artifacts by running this cell:
# ============================================================
# Configuration
# ============================================================
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)

IMG_SIZE = 384
BATCH_SIZE = 4
ACCUMULATION_STEPS = 4
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ============================================================
# Load Artifacts
# ============================================================
train_df = pd.read_csv("artifacts/train.csv")
val_df = pd.read_csv("artifacts/val.csv")
test_df = pd.read_csv("artifacts/test.csv")
class_weights_tensor = torch.load("artifacts/class_weights.pt")
with open("artifacts/label_mapping.json", "r") as f:
    label_mapping = json.load(f)

What does this code do?

  • SEED = 42 locks randomness so your results are reproducible across runs.
  • IMG_SIZE = 384 sets the image resolution. Larger than your original 224x224, this captures finer lesion details.
  • BATCH_SIZE = 4 with ACCUMULATION_STEPS = 4 simulates a batch size of 16 while only holding 4 images in GPU memory at once.
  • class_weights_tensor contains per-class weights computed from the imbalanced ISIC 2019 distribution.
  • Verify your artifacts loaded correctly by running this in a new cell:
print(f"Training Images: {len(train_df)}")
print(f"Classes: {label_mapping}")
print(f"Device: {device}")

You should see approximately 20,000+ training images, 8 class names in the label mapping, and cuda as the device.

Seeing a FileNotFoundError?

If loading fails, your artifacts/ folder may not be in the same directory as your notebook. Confirm the folder exists and contains train.csv, val.csv, test.csv, class_weights.pt, and label_mapping.json.

help me find my preprocessing artifacts

✔️ Awesome, I've got everything!

Great. Double check you have saved your notebook with all three cells above.

ⓧ I'd like to double check the full code

Scroll up and confirm you have these three cells in order:

  • Cell 1: The imports block (18 lines starting with import os through import matplotlib.pyplot as plt).
  • Cell 2: The configuration and artifact loading block (starting with the # Configuration header through label_mapping = json.load(f)).
  • Cell 3: The verification prints (Training Images, Classes, Device).

If any cell is missing or shows an error, re-run it from the code blocks above.

Your environment is ready with both model libraries, GPU acceleration, and all your preprocessing artifacts loaded. Next up, you will build the improved data pipeline that fixes the 51:1 class imbalance killing your accuracy.

Build the Improved Data Pipeline

Your current pipeline feeds the model 224x224 images with basic augmentations and no oversampling. The ISIC 2019 dataset has a severe 51:1 class imbalance (NV has 12,875 images while DF has just 239).

Without oversampling, your model barely sees rare classes during training, which caps accuracy at ~86%. In this step, you will upgrade to 384x384 resolution to capture fine-grained lesion details and add WeightedRandomSampler so every class gets fair exposure during training.

In this step, get ready to:
  • Create stronger image transforms at 384x384 resolution with multiple augmentation strategies.
  • Build a balanced data pipeline using WeightedRandomSampler to oversample minority classes.
  • Define reusable training and validation functions with mixed precision and gradient accumulation.
Create stronger transforms at 384x384

Larger input images preserve fine-grained texture details that distinguish similar lesion types. Combined with aggressive augmentation, the model learns to generalize across rotations, color shifts, and scale variations.

  • Define the training and validation transforms in a new Jupyter cell by running this code:
# Training transforms: aggressive augmentation at 384x384
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(IMG_SIZE, scale=(0.7, 1.0)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomVerticalFlip(),
    transforms.RandomRotation(30),
    transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.1),
    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Validation transforms: just resize and normalize (no augmentation)
val_transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

What does this code do?

RandomResizedCrop randomly crops AND resizes the image in one step. This teaches the model to handle lesions at different scales, from close-up to far away.

The validation transform only resizes without randomness. You want consistent, repeatable results when evaluating. The ImageNet normalization values (mean=[0.485, 0.456, 0.406]) match the pretrained weights you will load later.

  • Run the cell. It should complete with no output or errors.

Seeing a NameError for IMG_SIZE?

Make sure you ran the configuration cell from Step 1 that defines IMG_SIZE = 384. In Jupyter, each variable must be defined in a cell you have already executed.

help me fix a NameError for IMG_SIZE in my transforms cell

Create the dataset class and WeightedRandomSampler

The ISIC 2019 dataset has a severe imbalance. NV makes up 51% of all images. If you sample uniformly, a batch of 16 images might contain zero examples of DF or VASC. WeightedRandomSampler assigns higher sampling probability to minority classes so every batch has balanced representation.

  • Define the Dataset class and create the three dataset objects by running this code in a new cell:
class SkinCancerDataset(Dataset):
    def __init__(self, dataframe, transform=None):
        self.dataframe = dataframe
        self.transform = transform

    def __len__(self):
        return len(self.dataframe)

    def __getitem__(self, idx):
        # Load image path and label from the dataframe
        image_path = self.dataframe.iloc[idx]["image_path"]
        label = self.dataframe.iloc[idx]["label_encoded"]
        image = Image.open(image_path).convert("RGB")
        if self.transform:
            image = self.transform(image)
        return image, label

train_dataset = SkinCancerDataset(train_df, transform=train_transform)
val_dataset = SkinCancerDataset(val_df, transform=val_transform)
test_dataset = SkinCancerDataset(test_df, transform=val_transform)

What does this code do?

SkinCancerDataset is a custom PyTorch Dataset that reads image_path and label_encoded columns from a DataFrame. Each call to __getitem__ loads one image, applies the transform, and returns the tensor plus its label.

The test dataset uses val_transform (no augmentation) because you want deterministic evaluation on test images.

  • Run the cell. It should complete without errors.
  • Now create the WeightedRandomSampler and DataLoaders by running this code in a new cell:
# Compute per-sample weights for WeightedRandomSampler
labels = train_df["label_encoded"].values
class_counts = np.bincount(labels)
class_weights_sampler = 1.0 / class_counts
sample_weights = class_weights_sampler[labels]
sample_weights = torch.DoubleTensor(sample_weights)

sampler = WeightedRandomSampler(
    weights=sample_weights,
    num_samples=len(sample_weights),
    replacement=True
)

train_loader = DataLoader(
    train_dataset, batch_size=BATCH_SIZE,
    sampler=sampler, num_workers=0, pin_memory=True
)
val_loader = DataLoader(
    val_dataset, batch_size=BATCH_SIZE,
    shuffle=False, num_workers=0, pin_memory=True
)
test_loader = DataLoader(
    test_dataset, batch_size=BATCH_SIZE,
    shuffle=False, num_workers=0, pin_memory=True
)

What does this code do?

Each sample gets a weight equal to 1 / (count of its class). A DF image (239 total) gets a weight ~54x higher than an NV image (12,875 total). The sampler uses these weights to pick which images appear in each batch.

replacement=True means the same minority-class image can appear multiple times in one epoch. This is necessary because minority classes have so few images that without replacement, they would be exhausted early.

The batch size is 4 (set by BATCH_SIZE) because 384x384 images at FP32 would fill the 4GB GPU. Gradient accumulation (next substep) compensates by simulating a batch of 16.

  • Run the cell. It should complete without errors.

Before you move on, do you think the sampler produces a balanced distribution across all 8 classes, or will some classes still dominate?

  • Verify balanced sampling by running this code in a new cell:
# Verify balanced sampling (first 1000 samples)
print("Verifying balanced sampling (first 1000 samples):")
sample_labels = []
for i, (_, lbl) in enumerate(train_loader):
    sample_labels.extend(lbl.numpy())
    if i >= 250:
        break
unique, counts = np.unique(sample_labels, return_counts=True)
for cls, cnt in zip(unique, counts):
    print(f"  Class {cls}: {cnt} samples")

You should see roughly similar sample counts across all 8 classes. Instead of NV dominating with 51% of samples, each class now appears with approximately equal frequency.

Still seeing one class dominate?

Check that sample_weights has the same length as your training DataFrame. If np.bincount(labels) returns fewer entries than expected, your label_encoded column may have gaps in its integer values.

help me debug why my WeightedRandomSampler is not balancing classes properly

Define the training functions with mixed precision and gradient accumulation

Mixed precision (FP16) halves memory usage so 384x384 images fit on your 4GB GPU. Gradient accumulation simulates a larger effective batch size by accumulating gradients over multiple small batches before updating weights.

  • Define the training function by running this code in a new cell:
def train_one_epoch(model, dataloader, criterion, optimizer, scaler, device, accumulation_steps):
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    optimizer.zero_grad()

    for i, (images, labels) in enumerate(dataloader):
        images = images.to(device)
        labels = labels.to(device)

        # Forward pass in FP16 to save memory
        with autocast(device_type='cuda', dtype=torch.float16):
            outputs = model(images)
            loss = criterion(outputs, labels)
            loss = loss / accumulation_steps

        # Backward pass with scaled gradients
        scaler.scale(loss).backward()

        # Only update weights every accumulation_steps batches
        if (i + 1) % accumulation_steps == 0:
            scaler.step(optimizer)
            scaler.update()
            optimizer.zero_grad()

        running_loss += loss.item() * accumulation_steps
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    epoch_loss = running_loss / len(dataloader)
    epoch_accuracy = 100.0 * correct / total
    return epoch_loss, epoch_accuracy

What does this code do?

autocast(device_type='cuda', dtype=torch.float16) runs the forward pass in half-precision. This uses roughly half the GPU memory compared to FP32, which is how 384x384 images fit on a 4GB card.

The loss is divided by accumulation_steps before backprop. Gradients accumulate for 4 mini-batches, then one optimizer step fires. The effective batch size becomes 4 x 4 = 16 without ever holding 16 images in memory at once.

GradScaler prevents FP16 underflow by scaling gradients up before backprop and scaling them back down before the optimizer step.

  • Run the cell. It defines the function without producing output.
  • Define the validation function and prediction helper by running this code in a new cell:
def validate_one_epoch(model, dataloader, criterion, device):
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0

    with torch.no_grad():
        for images, labels in dataloader:
            images = images.to(device)
            labels = labels.to(device)

            with autocast(device_type='cuda', dtype=torch.float16):
                outputs = model(images)
                loss = criterion(outputs, labels)

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

    epoch_loss = running_loss / len(dataloader)
    epoch_accuracy = 100.0 * correct / total
    return epoch_loss, epoch_accuracy

How is validation different from training?

model.eval() disables dropout and batch normalization updates. torch.no_grad() stops PyTorch from tracking gradients, which saves memory and speeds up inference.

No gradient accumulation or scaler is needed here because you are only measuring performance, not updating weights.

  • Run the cell. It defines the function without producing output.
  • Define the prediction helper function by running this code in a new cell:
def get_predictions(model, dataloader, device):
    model.eval()
    all_preds = []
    all_labels = []
    with torch.no_grad():
        for images, labels in dataloader:
            images = images.to(device)
            with autocast(device_type='cuda', dtype=torch.float16):
                outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            all_preds.extend(predicted.cpu().numpy())
            all_labels.extend(labels.numpy())
    return np.array(all_preds), np.array(all_labels)

What does this code do?

get_predictions runs the model on an entire DataLoader and collects all predictions and true labels into NumPy arrays. You will use this in Step 5 to generate classification reports for both models.

  • Run the cell. Confirm all three functions are defined by running this quick check in a new cell:
print(f"train_one_epoch: {type(train_one_epoch)}")
print(f"validate_one_epoch: {type(validate_one_epoch)}")
print(f"get_predictions: {type(get_predictions)}")

You should see three lines, each showing <class 'function'>. Your training pipeline is ready.

Seeing a NameError for autocast or GradScaler?

Make sure you ran the imports cell from Step 1. Both autocast and GradScaler are imported from torch.amp. If you restarted your kernel, re-run all cells from the top.

help me fix a NameError in my training function definitions

✔️ Awesome, I've got everything!

Great. Confirm you have saved all cells in your notebook.

ⓧ I'd like to double check the full code

Scroll through your notebook and confirm these sections appear in order, each in their own cell:

  • Transforms cell: train_transform and val_transform defined with transforms.Compose.
  • Dataset cell: SkinCancerDataset class plus train_dataset, val_dataset, test_dataset instantiation.
  • Sampler cell: sample_weights computed, WeightedRandomSampler created, three DataLoaders defined.
  • Training function cell: train_one_epoch with autocast and gradient accumulation.
  • Validation cell: validate_one_epoch with torch.no_grad() and FP16 inference.
  • Predictions cell: get_predictions returning NumPy arrays of predictions and labels.

Your data pipeline now handles the 51:1 class imbalance and fits 384x384 images on your 4GB GPU. Next up, you will train ConvNeXt Tiny with a two-phase strategy that preserves pretrained features while adapting to skin lesions.

Train ConvNeXt Tiny with Layer Freezing

Your data pipeline is ready to feed balanced batches into a model. The next challenge is training ConvNeXt Tiny so it actually learns skin lesion features without destroying the useful patterns it already knows from ImageNet.

Your original training trains ALL layers with the same learning rate from epoch 1. This destroys the good pretrained features in early layers. The fix is a two-phase approach: first freeze the backbone and train only the new classifier head (5 epochs), then unfreeze everything with a much lower learning rate for the backbone (15 epochs).

In this step, get ready to:
  • Run Phase 1 training with the backbone frozen, letting only the classifier head learn skin lesion classes.
  • Run Phase 2 training with the backbone unfrozen and differential learning rates, pushing past 92% accuracy.
Phase 1: freeze backbone, train classifier head only

Transfer learning works best when you preserve the pretrained backbone's knowledge of edges, textures, and shapes. By freezing those layers first, only the brand-new classifier head updates its weights to map ConvNeXt's features to your 8 skin lesion classes.

  • Create a new Jupyter cell and paste the following code to load ConvNeXt Tiny, freeze its backbone, and train for 5 epochs:
# Define loss function with class weights
criterion = nn.CrossEntropyLoss(weight=class_weights_tensor.to(device))

# Load pretrained ConvNeXt Tiny
model_convnext = convnext_tiny(weights=ConvNeXt_Tiny_Weights.DEFAULT)
num_features = model_convnext.classifier[2].in_features
model_convnext.classifier[2] = nn.Linear(num_features, 8)

# Phase 1: Freeze backbone
for param in model_convnext.features.parameters():
    param.requires_grad = False

model_convnext = model_convnext.to(device)
scaler = GradScaler()
optimizer_phase1 = optim.AdamW(model_convnext.classifier.parameters(), lr=1e-3, weight_decay=1e-4)

print("ConvNeXt Phase 1: Training classifier head")
for epoch in range(5):
    train_loss, train_acc = train_one_epoch(
        model_convnext, train_loader, criterion, optimizer_phase1, scaler, device, ACCUMULATION_STEPS
    )
    val_loss, val_acc = validate_one_epoch(model_convnext, val_loader, criterion, device)
    print(f"  Epoch [{epoch+1}/5] | Train Acc: {train_acc:.2f}% | Val Acc: {val_acc:.2f}%")

What does this code do?

This cell sets up and runs the first training phase:

  • criterion uses your precomputed class weights so the loss function penalizes mistakes on rare classes (like Dermatofibroma) more heavily.
  • model_convnext.classifier[2] replaces ConvNeXt's original 1000-class ImageNet head with a new 8-class linear layer for your skin lesion types.
  • param.requires_grad = False on all backbone parameters means only the classifier head's weights update during Phase 1.
  • GradScaler() handles mixed precision scaling so FP16 gradients stay numerically stable.
  • Run the cell. This takes a few minutes per epoch on your GPU.

You should see 5 epochs printed, with validation accuracy climbing to around 80-85% by the final epoch.

Why is 80-85% already good for Phase 1?

Only the classifier head is learning here. The frozen backbone extracts features it learned from ImageNet, and the head learns to map those features to skin lesion classes. Reaching 80-85% with a frozen backbone means the pretrained features are highly relevant to this task.

Seeing CUDA out of memory?

Make sure no other notebooks or processes are using your GPU. Restart your Jupyter kernel to free any leftover memory from previous experiments. Your BATCH_SIZE should be 4 (not higher) for a 4GB GPU.

Help me debug a CUDA out of memory error during ConvNeXt Phase 1 training with batch size 4 on a 4GB GPU.

Phase 2: unfreeze backbone with differential learning rates

Phase 1 proved the backbone features are useful. Now you unfreeze the backbone so it can adapt those features specifically to skin lesions. The key is using a 10x lower learning rate for the backbone than for the classifier. This prevents the large updates that would destroy the pretrained knowledge.

Before you run Phase 2, think about this: do you expect accuracy to jump immediately in epoch 1, or climb gradually over 15 epochs? Why?

  • Create a new cell and paste the following code to unfreeze the backbone and fine-tune with differential learning rates for 15 epochs:
# Phase 2: Unfreeze backbone
for param in model_convnext.features.parameters():
    param.requires_grad = True

# Differential learning rates: backbone gets 10x lower LR
optimizer_phase2 = optim.AdamW([
    {"params": model_convnext.features.parameters(), "lr": 1e-5},
    {"params": model_convnext.classifier.parameters(), "lr": 1e-4},
], weight_decay=1e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer_phase2, T_max=15, eta_min=1e-7)

print("\nConvNeXt Phase 2: Fine-tuning full model")
best_convnext_acc = 0.0
best_convnext_weights = None

for epoch in range(15):
    train_loss, train_acc = train_one_epoch(
        model_convnext, train_loader, criterion, optimizer_phase2, scaler, device, ACCUMULATION_STEPS
    )
    val_loss, val_acc = validate_one_epoch(model_convnext, val_loader, criterion, device)
    scheduler.step()
    print(f"  Epoch [{epoch+1}/15] | Train Acc: {train_acc:.2f}% | Val Acc: {val_acc:.2f}%")
    if val_acc > best_convnext_acc:
        best_convnext_acc = val_acc
        best_convnext_weights = copy.deepcopy(model_convnext.state_dict())

print(f"\nConvNeXt Best Accuracy: {best_convnext_acc:.2f}%")
torch.save(best_convnext_weights, "best_convnext_384.pth")

What does this code do?

This cell runs the full fine-tuning phase:

  • Setting requires_grad = True unfreezes the backbone so ALL layers can now update.
  • The backbone gets lr=1e-5 while the classifier gets lr=1e-4. This 10x difference lets early layers make small adjustments without overwriting ImageNet features.
  • CosineAnnealingLR gradually reduces the learning rate following a cosine curve, preventing overshooting as the model converges.
  • copy.deepcopy(model_convnext.state_dict()) saves the weights from whichever epoch had the best validation accuracy, not just the final epoch.
  • Run the cell. Phase 2 takes longer per epoch because the entire model is now updating.

You should see accuracy climb steadily over 15 epochs. The final line prints the best validation accuracy, which should exceed 92%. The best model weights are saved to best_convnext_384.pth.

Why did this break past 86%?

Three changes compound here. The 384x384 resolution captures fine lesion details that 224x224 misses. The WeightedRandomSampler ensures the model actually trains on rare classes like DF and VASC. The two-phase freezing strategy preserves useful pretrained features instead of destroying them on epoch 1.

Accuracy stuck below 90% after Phase 2?

Check that Phase 1 reached 80-85% before you started Phase 2. If Phase 1 accuracy was low, restart the kernel and re-run from the beginning of this step. Also verify that class_weights_tensor loaded correctly from your artifacts.

Help me debug low ConvNeXt accuracy after Phase 2 fine-tuning.

✔️ Awesome, I've got everything!

Great. Confirm you see the best_convnext_384.pth file saved and that your printed best accuracy exceeds 92%.

ⓧ I'd like to double check the full code

Scroll up and confirm you have both cells in order:

  • Cell 1 (Phase 1): Starts with criterion = nn.CrossEntropyLoss(...), loads ConvNeXt, freezes backbone, trains 5 epochs.
  • Cell 2 (Phase 2): Starts with for param in model_convnext.features.parameters(), unfreezes backbone, trains 15 epochs with differential LR, saves best_convnext_384.pth.

Your ConvNeXt Tiny is now trained and exceeding 92%. Next up, you will train a ViT-Small model with the same strategy to prove this approach works across different architectures.

Train ViT-Small with the Same Strategy

Your ConvNeXt model just proved that the two-phase training strategy works. It broke through 92% by learning local features like textures, edges, and color transitions within skin lesions.

Now you need a second model that captures something fundamentally different: global features. A Vision Transformer (ViT) processes the entire image as a sequence of patches, letting it learn long-range spatial relationships across the full lesion. Training ViT-Small with the same pipeline proves your approach generalizes across architectures and gives you the second model needed for the fusion project.

In this step, get ready to:
  • Load ViT-Small from timm with a frozen backbone and train the classification head.
  • Unfreeze the full model and fine-tune with differential learning rates.
  • Verify ViT-Small also exceeds 92% validation accuracy.
Phase 1: load ViT-Small from timm and freeze backbone

The timm library provides the vit_small_patch16_384.augreg_in21k_ft_in1k model. This ViT-Small was pretrained on ImageNet-21k (14 million images) and fine-tuned on ImageNet-1k, with 22.2M parameters designed for 384x384 input.

  • In a new Jupyter cell, load the pretrained ViT-Small, freeze all layers except the head, and run Phase 1 training by pasting this code:
# Load ViT-Small pretrained on ImageNet-21k, fine-tuned on ImageNet-1k
# 22.2M params, designed for 384x384 input
model_vit = timm.create_model(
    'vit_small_patch16_384.augreg_in21k_ft_in1k',
    pretrained=True,
    num_classes=8
)

# Freeze all layers except the classification head
for name, param in model_vit.named_parameters():
    if "head" not in name:
        param.requires_grad = False

model_vit = model_vit.to(device)
scaler_vit = GradScaler()

optimizer_vit_phase1 = optim.AdamW(
    model_vit.head.parameters(), lr=1e-3, weight_decay=1e-4
)

print("\nViT-Small Phase 1: Training head")
for epoch in range(5):
    train_loss, train_acc = train_one_epoch(
        model_vit, train_loader, criterion, optimizer_vit_phase1,
        scaler_vit, device, ACCUMULATION_STEPS
    )
    val_loss, val_acc = validate_one_epoch(
        model_vit, val_loader, criterion, device
    )
    print(f"  Epoch [{epoch+1}/5] | Train Acc: {train_acc:.2f}% | Val Acc: {val_acc:.2f}%")

What does this code do?

timm.create_model() downloads the pretrained ViT-Small and replaces the final classification layer with a new 8-class head (matching your skin lesion classes).

The freezing loop checks each parameter's name. Anything without "head" in the name gets frozen. This is different from ConvNeXt where you froze model_convnext.features directly. Each library structures its model differently, so the freeze logic adapts.

A separate scaler_vit is created because each GradScaler tracks its own scaling state for a specific model's training run.

This cell will take several minutes to run. You should see epoch-by-epoch output with validation accuracy climbing toward 80-85%.

Seeing an out-of-memory error?

ViT-Small at 384x384 is tight on a 4GB GPU. Make sure no other Python processes are using your GPU memory. Restart your Jupyter kernel, re-run all previous cells, then try this cell again.

help me debug a CUDA out-of-memory error when loading ViT-Small on a 4GB GPU

Phase 2: unfreeze and fine-tune the full ViT

Phase 1 taught the head to map ViT's frozen features to your 8 classes. Now you unfreeze the entire transformer and fine-tune it with the same differential learning rate strategy you used for ConvNeXt: the backbone learns at 10x lower rate than the head.

Before you run Phase 2, think about this: do you expect ViT to reach the same accuracy as ConvNeXt, higher, or lower? They see the same data but process it differently.

  • In a new cell, unfreeze the full model and run Phase 2 training by pasting this code:
# UNFREEZE all layers for full fine-tuning
for param in model_vit.parameters():
    param.requires_grad = True

# Differential learning rates: backbone 10x lower than head
optimizer_vit_phase2 = optim.AdamW([
    {"params": [p for n, p in model_vit.named_parameters() if "head" not in n], "lr": 1e-5},
    {"params": model_vit.head.parameters(), "lr": 1e-4},
], weight_decay=1e-4)

scheduler_vit = optim.lr_scheduler.CosineAnnealingLR(
    optimizer_vit_phase2, T_max=15, eta_min=1e-7
)

print("\nViT-Small Phase 2: Fine-tuning full model")
best_vit_acc = 0.0
best_vit_weights = None

for epoch in range(15):
    train_loss, train_acc = train_one_epoch(
        model_vit, train_loader, criterion, optimizer_vit_phase2,
        scaler_vit, device, ACCUMULATION_STEPS
    )
    val_loss, val_acc = validate_one_epoch(
        model_vit, val_loader, criterion, device
    )
    scheduler_vit.step()

    print(f"  Epoch [{epoch+1}/15] | Train Acc: {train_acc:.2f}% | Val Acc: {val_acc:.2f}%")

    if val_acc > best_vit_acc:
        best_vit_acc = val_acc
        best_vit_weights = copy.deepcopy(model_vit.state_dict())

print(f"\nViT-Small Best Accuracy: {best_vit_acc:.2f}%")
torch.save(best_vit_weights, "best_vit_384.pth")

What does this code do?

  • The list comprehension [p for n, p in model_vit.named_parameters() if "head" not in n] collects all backbone parameters (attention layers, patch embedding, norm layers) into one group with a learning rate of 1e-5.
  • CosineAnnealingLR gradually decays the learning rate from its initial value down to 1e-7 over 15 epochs, following a cosine curve. This prevents overfitting in later epochs.
  • copy.deepcopy(model_vit.state_dict()) saves the best weights in memory whenever validation accuracy improves. The final torch.save writes the best checkpoint to disk as best_vit_384.pth.

Phase 2 runs 15 epochs and will take significantly longer than Phase 1 since all 22.2M parameters are being updated. Watch the epoch output as it prints.

Once training finishes, you should see the final line reporting ViT-Small's best validation accuracy above 92%.

Why does ViT need the same strategy as ConvNeXt?

Both models share the same core problem: they were pretrained on millions of natural images (dogs, cars, landscapes) and must adapt to dermatoscopic skin lesions. The two-phase approach (freeze then unfreeze with lower backbone LR) prevents catastrophic forgetting regardless of architecture.

The key difference is in the freezing mechanics. ConvNeXt organizes layers under model.features and model.classifier[2]. ViT from timm uses model.head as the classification layer, with everything else being the backbone.

Training seems stuck at low accuracy?

If Phase 2 accuracy plateaus below 88%, confirm that Phase 1 finished successfully with accuracy around 80-85%. A weak Phase 1 means the head never learned good initial mappings, so Phase 2 starts from a poor foundation.

Also verify that your criterion still uses the class weights tensor. If you accidentally re-defined it without weights, the model will ignore rare classes.

help me debug why my ViT-Small is not reaching 92% accuracy during Phase 2 training

✔️ Awesome, I've got everything!

Great. Double check you see best_vit_384.pth saved in your working directory and that the printed best accuracy exceeds 92%.

ⓧ I'd like to double check the full code

Scroll up and confirm you have two cells for ViT training, in this order:

  • Cell 1 (Phase 1): loads vit_small_patch16_384.augreg_in21k_ft_in1k with num_classes=8, freezes all non-head parameters, creates scaler_vit and optimizer_vit_phase1, trains for 5 epochs.
  • Cell 2 (Phase 2): unfreezes all parameters, creates optimizer_vit_phase2 with differential LRs (1e-5 backbone, 1e-4 head), creates scheduler_vit with T_max=15, trains for 15 epochs, saves best weights to best_vit_384.pth.

You now have two independently trained models, each exceeding 92% accuracy. Next up, you will evaluate both side by side and compare their per-class strengths to understand what each model handles best.

Evaluate and Compare Both Models

Both ConvNeXt and ViT have completed training, and their best weights are saved to disk. The training logs showed promising validation accuracy, but you need to confirm performance on the held-out test set that neither model has ever seen.

Running evaluation on the test set reveals per-class accuracy for all 8 skin lesion types. This tells you exactly which classes each architecture handles better. That information is critical for the feature fusion step in your next project.

In this step, get ready to:
  • Generate test predictions for both trained models.
  • Print per-class classification reports and a summary comparison.
  • Confirm both models are saved and ready for the fusion project.
Generate test predictions for both models

You already have the get_predictions function defined from earlier. Now you need to load each model's best weights and run inference on the entire test set.

  • Load the best saved weights into each model and generate predictions by running this cell:
# Load best weights from training
model_convnext.load_state_dict(best_convnext_weights)
model_vit.load_state_dict(best_vit_weights)

# Generate predictions on the held-out test set
preds_convnext, labels_test = get_predictions(model_convnext, test_loader, device)
preds_vit, _ = get_predictions(model_vit, test_loader, device)

print(f"ConvNeXt predictions: {len(preds_convnext)} samples")
print(f"ViT predictions: {len(preds_vit)} samples")

What does this code do?

load_state_dict(best_convnext_weights) replaces the model's current parameters with the best checkpoint saved during training. This ensures you evaluate the epoch that performed best on validation, not necessarily the final epoch.

get_predictions runs inference in eval mode with no gradient tracking and mixed precision. It returns numpy arrays of predicted classes and true labels for every test image.

The cell will take a minute or two to run inference on the full test set. You should see output confirming the number of predictions matches your test set size.

Seeing a RuntimeError?

  • If you see a shape mismatch error, make sure you did not restart the kernel between training and evaluation. The model architectures must match the saved weights exactly.
  • If you see a CUDA out of memory error, restart the kernel and re-run just the evaluation cells. Training left optimizer states in GPU memory that evaluation does not need.

Help me debug my model evaluation error

Print classification reports and compare accuracy

With predictions in hand, you can now see exactly how each model performs on every class. The classification_report from scikit-learn breaks down precision, recall, and F1-score per class.

Before you run the next cell, think about this: do you expect both models to exceed 92%, or will one architecture outperform the other? Which minority classes (DF, VASC, AK) do you think improved the most compared to the 86% baseline?

  • Print the full classification reports and summary comparison by running this cell:
# Map class indices back to readable names
class_names = list(label_mapping.keys())

# Per-class report for ConvNeXt
print("\n" + "=" * 60)
print("ConvNeXt Tiny Results")
print("=" * 60)
print(classification_report(labels_test, preds_convnext, target_names=class_names))

# Per-class report for ViT
print("=" * 60)
print("ViT-Small Results")
print("=" * 60)
print(classification_report(labels_test, preds_vit, target_names=class_names))

# Overall accuracy comparison
acc_convnext = 100.0 * np.mean(preds_convnext == labels_test)
acc_vit = 100.0 * np.mean(preds_vit == labels_test)
print("=" * 60)
print("SUMMARY")
print("=" * 60)
print(f"ConvNeXt Tiny Accuracy: {acc_convnext:.2f}%")
print(f"ViT-Small Accuracy:     {acc_vit:.2f}%")
print(f"Previous baseline:      86.00%")
print(f"Target:                 >92.00%")

What does this output tell you?

The classification_report shows precision, recall, and F1-score for each of the 8 skin lesion classes. Look at the minority classes (DF, VASC, AK) specifically. These are the classes your 86% baseline struggled with because of the 51:1 imbalance.

The SUMMARY section compares both models against the 86% baseline and the 92% target. Both architectures should exceed 92%, proving that the combination of WeightedRandomSampler, 384x384 resolution, and two-phase layer freezing works across different architectures.

You should see both models reporting accuracy above 92% in the SUMMARY block. The per-class reports should show significant improvements on DF, VASC, and AK compared to the 86% baseline where those classes had low recall.

Why do ConvNeXt and ViT have different per-class strengths?

ConvNeXt excels at local features like textures, color gradients, and edge patterns. ViT excels at global features like long-range spatial relationships across the entire lesion. Some lesion types are best identified by local texture (where ConvNeXt wins), while others need global context (where ViT wins).

This difference is exactly why the next project fuses both models. Their complementary strengths combined can push accuracy even higher.

Accuracy below 92%?

  • If accuracy is in the 88-91% range, the model may need a few more Phase 2 epochs. Try increasing the Phase 2 epoch count from 15 to 20 and re-running.
  • If accuracy is significantly lower (below 85%), check that the WeightedRandomSampler is working correctly. Re-run the balanced sampling verification from Step 2 and confirm roughly equal class counts.
  • Make sure you loaded best_convnext_weights (not the final epoch weights). The best checkpoint is the one that achieved peak validation accuracy.

Help me understand why my model accuracy is below 92%

Confirm models are saved for the fusion project

Both models were saved during training with torch.save. You now have two independently trained classifiers ready to serve as feature extractors in your next project.

  • Confirm both model files exist by running this cell:
# Confirm saved model files
import os
print("Models saved for fusion project:")
print(f"  - best_convnext_384.pth: {os.path.getsize('best_convnext_384.pth') / 1e6:.1f} MB")
print(f"  - best_vit_384.pth: {os.path.getsize('best_vit_384.pth') / 1e6:.1f} MB")

What are these models used for next?

  • best_convnext_384.pth is your local feature extractor. It captures textures, edges, and color transitions in skin lesions.
  • best_vit_384.pth is your global feature extractor. It captures long-range spatial relationships across the entire lesion image.

In the next project, you will extract features from both models, fuse them with an attention mechanism, and add a quantum classification layer on top.

You should see both files listed with their sizes (ConvNeXt around 110 MB, ViT around 85 MB). These confirm your trained weights are persisted to disk.

✔️ Awesome, I've got everything!

Great. Both models exceeded 92% and are saved. You have completed the training pipeline.

ⓧ I'd like to double check the full code

Scroll through your notebook and confirm you have the following sections in order:

  • Configuration cell (SEED, IMG_SIZE, BATCH_SIZE, ACCUMULATION_STEPS, device).
  • Load Artifacts cell (CSVs, class_weights, label_mapping).
  • Transforms cell (train_transform, val_transform at 384x384).
  • Dataset and WeightedRandomSampler cell (SkinCancerDataset, sampler, data loaders).
  • Training Functions cell (train_one_epoch, validate_one_epoch, get_predictions).
  • ConvNeXt Phase 1 and Phase 2 training cells.
  • ViT-Small Phase 1 and Phase 2 training cells.
  • Evaluation cell (load_state_dict, get_predictions, classification_report, SUMMARY).

Secret mission

Add Test-Time Augmentation (TTA)

Your models already break 92%, but there is still accuracy left on the table. Test-Time Augmentation runs multiple augmented views of each test image through the model, averages their predictions, and squeezes out an extra 1-2% without any retraining. Can you implement a TTA pipeline that boosts both ConvNeXt and ViT even further?

Clean Up Your Resources

Clean Up Your Resources

Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project runs entirely locally, so there are no ongoing costs.

Resources you used:

  • Trained model weights: best_convnext_384.pth and best_vit_384.pth.
  • Training script: skin_lesion_classification.py.
  • TTA inference script: tta_inference.py.
  • Python package: timm (PyTorch Image Models library).

Keep everything running

No action needed. Choose this if you plan to use these models in your next project (feature fusion + quantum classification).

  • Your model weights (best_convnext_384.pth and best_vit_384.pth) are saved locally and ready to load in a future project.
  • The timm library remains installed in your Python environment for future ViT work.
  • Your Jupyter Notebook kernel stays active with all variables in memory.

Pause - I'll come back to this later

Shut down running processes to free up GPU memory while keeping all files intact.

  • In Jupyter Notebook, click Kernel in the menu bar, then select Shut Down All Kernels. This releases GPU memory.
  • Your saved model files (.pth) and scripts remain on disk, ready to reload when you return.

Delete - I don't want to use this again

Remove all project files and the installed library to start fresh.

  • Shut down your Jupyter Notebook kernel by clicking Kernel, then Shut Down All Kernels.
  • Delete the trained model weights by running this in your terminal:
rm best_convnext_384.pth best_vit_384.pth
  • Delete the project scripts by running:
rm skin_lesion_classification.py tta_inference.py
  • Uninstall the timm library by running:
pip uninstall timm -y
  • Confirm the files are gone by running:
ls *.pth skin_lesion_classification.py tta_inference.py 2>/dev/null || echo "All project files removed."

Nice Work!

Nice Work!

You just took your skin lesion classifier from 86% to over 92% accuracy by applying professional-grade training techniques to both ConvNeXt and ViT architectures.

You've learned how to:

  • Build a balanced data pipeline using WeightedRandomSampler that overcomes severe class imbalance (51:1 ratio) so rare lesion types get fair training representation.
  • Train models with a two-phase fine-tuning strategy (freeze backbone, then unfreeze with differential learning rates) that preserves pretrained features while adapting to skin lesion classification.
  • Fit 384x384 resolution models on a 4GB GPU using mixed precision (FP16) and gradient accumulation, producing two high-accuracy models ready for feature fusion.
  • Secret Mission: Implemented Test-Time Augmentation (TTA) to squeeze additional accuracy from both models without any retraining.

Ready to quiz yourself?