Diagnosing Diabetic Retinopathy from Fundus Images with PyTorch Lightning

Introduction

Diabetic Retinopathy (DR) is a leading cause of blindness globally. Early detection via retinal fundus imaging is critical for treatment. In this tutorial, we will train an automated DR classification model using PyTorch Lightning. PyTorch Lightning removes the boilerplate code of raw PyTorch, allowing us to focus purely on the deep learning architecture while effortlessly scaling to multiple GPUs or 16-bit precision.

Architecture Overview

graph TD
    Data[(APTOS Fundus Images)] --> DataLoader[PyTorch DataLoader]
    DataLoader --> Aug[Data Augmentation: Flip, Color Jitter, Resize]
    
    Aug --> Model[LightningModule: ResNet50 Backbone]
    Model --> Pred[Class Predictions 0-4]
    
    Pred --> Loss[CrossEntropy Loss]
    Loss --> Trainer[Lightning Trainer]
    
    Trainer --> Callbacks[Callbacks: Early Stopping, Model Checkpoint]
    Trainer -. Backprop .-> Model
        

Prerequisites

• Python 3.9+
• PyTorch, torchvision, pytorch-lightning
• The APTOS 2019 Blindness Detection dataset (or similar fundus imagery)

Step 1: Install Dependencies

pip install torch torchvision pytorch-lightning pandas Pillow scikit-learn

Step 2: Defining the PyTorch Dataset

Fundus images are large and varied in lighting. We need a custom Dataset to load the images, apply standard augmentations (like random flips and color jitter to simulate different camera sensors), and return the image tensor alongside its DR severity label (0-4).

import os
from PIL import Image
import pandas as pd
from torch.utils.data import Dataset
from torchvision import transforms

class FundusDataset(Dataset):
    def __init__(self, csv_file, img_dir, transform=None):
        self.labels_df = pd.read_csv(csv_file)
        self.img_dir = img_dir
        self.transform = transform

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

    def __getitem__(self, idx):
        img_name = os.path.join(self.img_dir, self.labels_df.iloc[idx, 0] + '.png')
        image = Image.open(img_name).convert('RGB')
        label = self.labels_df.iloc[idx, 1]

        if self.transform:
            image = self.transform(image)

        return image, label

# Standard ImageNet augmentations + resizing
train_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomVerticalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

Step 3: Creating the LightningModule

A LightningModule organizes our PyTorch code into 5 clear sections: initialization, forward pass, training step, validation step, and optimizer configuration. We will use a pre-trained ResNet50 as our backbone.

import pytorch_lightning as pl
import torch.nn as nn
import torch.optim as optim
import torchvision.models as models
import torch
import torchmetrics

class DRClassifier(pl.LightningModule):
    def __init__(self, num_classes=5, learning_rate=1e-4, class_weights=None):
        super().__init__()
        self.save_hyperparameters()
        
        # Load pre-trained ResNet50
        self.model = models.resnet50(weights=models.ResNet50_Weights.IMAGENET1K_V1)
        
        # Replace the final classification head
        num_ftrs = self.model.fc.in_features
        self.model.fc = nn.Linear(num_ftrs, num_classes)
        
        # Handle class imbalance if weights are provided
        weight_tensor = torch.tensor(class_weights).float() if class_weights else None
        self.criterion = nn.CrossEntropyLoss(weight=weight_tensor)
        
        # Use TorchMetrics for robust metric calculation
        self.train_acc = torchmetrics.Accuracy(task="multiclass", num_classes=num_classes)
        self.val_acc = torchmetrics.Accuracy(task="multiclass", num_classes=num_classes)
        self.val_kappa = torchmetrics.CohenKappa(task="multiclass", num_classes=num_classes, weights="quadratic")

    def forward(self, x):
        return self.model(x)

    def training_step(self, batch, batch_idx):
        images, labels = batch
        outputs = self(images)
        loss = self.criterion(outputs, labels)
        
        self.train_acc(outputs, labels)
        self.log('train_loss', loss, prog_bar=True)
        self.log('train_acc', self.train_acc, prog_bar=True, on_step=False, on_epoch=True)
        return loss

    def validation_step(self, batch, batch_idx):
        images, labels = batch
        outputs = self(images)
        loss = self.criterion(outputs, labels)
        
        self.val_acc(outputs, labels)
        self.val_kappa(outputs, labels)
        
        self.log('val_loss', loss, prog_bar=True)
        self.log('val_acc', self.val_acc, prog_bar=True)
        self.log('val_kappa', self.val_kappa, prog_bar=True)

    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=self.hparams.learning_rate)

Step 4: Training with the Lightning Trainer

The Trainer abstracts away the training loops, `.to(device)` calls, and precision settings.

from torch.utils.data import DataLoader
from pytorch_lightning import Trainer
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping

# Setup Dataloaders (assuming train.csv and val.csv exist)
train_dataset = FundusDataset(csv_file='train.csv', img_dir='train_images', transform=train_transform)
val_dataset = FundusDataset(csv_file='val.csv', img_dir='val_images', transform=train_transform)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)
val_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)

# Initialize Model
model = DRClassifier(num_classes=5)

# Setup Callbacks
# We monitor Quadratic Cohen's Kappa, the standard metric for DR challenges
checkpoint_callback = ModelCheckpoint(monitor='val_kappa', mode='max', dirpath='checkpoints', filename='dr-model-{epoch:02d}-{val_kappa:.2f}')
early_stop_callback = EarlyStopping(monitor='val_kappa', mode='max', patience=5)

# Initialize Trainer (Automatically uses GPU if available, and uses 16-bit precision to save memory)
trainer = Trainer(
    max_epochs=20,
    accelerator='auto',
    devices='auto',
    precision=16, 
    callbacks=[checkpoint_callback, early_stop_callback]
)

# Start Training
trainer.fit(model, train_loader, val_loader)

Troubleshooting: Class Imbalance

Medical datasets are notoriously imbalanced. In the APTOS dataset, the vast majority of images are "No DR" (Class 0), while "Severe DR" (Class 4) is very rare. If your model simply predicts Class 0 for every image, it will still achieve high accuracy, but it is useless clinically.

• Use Cohen's Kappa: Do not rely on plain Accuracy. Use Quadratic Weighted Cohen's Kappa, which heavily penalizes the model for misclassifying a Severe case as a Healthy case.
• Class Weights: Calculate the inverse frequency of each class in your training set and pass that array to the class_weights argument in our DRClassifier. This forces the CrossEntropyLoss to penalize errors on the rare classes more severely.
• Focal Loss: For extreme imbalance, replace CrossEntropyLoss with Focal Loss, which dynamically scales the loss based on prediction confidence.

Conclusion

By using PyTorch Lightning, we reduced hundreds of lines of complex GPU-management and training-loop code into a clean, reproducible object-oriented structure. By integrating TorchMetrics to track Quadratic Cohen's Kappa, we ensure our model is actually learning to detect Diabetic Retinopathy rather than just guessing the majority class. This ResNet50 model can now be easily evaluated, or converted via ONNX to run on edge devices like smartphone fundoscopy attachments to provide point-of-care DR screening.