Erfan-Nourbakhsh/KD-OCT-ConvNeXtV2

KD-OCT: Knowledge Distillation for OCT Image Classification

This model is part of the KD-OCT project, which provides a framework for efficient, clinical-grade retinal Optical Coherence Tomography (OCT) image classification. It was introduced in the paper KD-OCT: Efficient Knowledge Distillation for Clinical-Grade Retinal OCT Classification.

The official implementation is available at the GitHub Repository.

Model Description

KD-OCT is a novel knowledge distillation framework designed to compress a high-performance ConvNeXtV2-Large teacher model into a lightweight EfficientNet-B2 student model. The framework is optimized to classify OCT B-scans into clinical categories:

• Normal
• Drusen
• Choroidal Neovascularization (CNV)
• Diabetic Macular Edema (DME) (in the 4-class UCSD variant)

By employing real-time distillation with a combined loss (soft teacher knowledge transfer and hard ground-truth supervision), the student model achieves near-teacher accuracy while significantly reducing model size and inference time, facilitating deployment on edge devices for AMD screening.

Training Details

• Framework: PyTorch
• Architecture: EfficientNet-B2 (Student) distilled from ConvNeXtV2-Large (Teacher)
• Datasets: Noor Eye Hospital (NEH) and UCSD OCT datasets
• Training Method: Knowledge Distillation using combined KL Divergence and Cross-Entropy/Focal Loss
• Optimizer: AdamW with Cosine Annealing

Usage

Below is an example of how to perform inference with the model. Note that according to the paper, the images should be resized to 384x384 for inference.

import torch
from torchvision import transforms
from PIL import Image

# Load model
model = torch.load("model.pth")
model.eval()

# Prepare image
transform = transforms.Compose([
    transforms.Resize((384, 384)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                       std=[0.229, 0.224, 0.225])
])

# Assume 'image' is a PIL image
input_tensor = transform(image).unsqueeze(0)

# Inference
with torch.no_grad():
    output = model(input_tensor)
    prediction = torch.argmax(output, dim=1)

Citation

If you use this model or the KD-OCT framework, please cite:

@article{nourbakhsh2025kd,
  title={KD-OCT: Efficient Knowledge Distillation for Clinical-Grade Retinal OCT Classification},
  author={Nourbakhsh, Erfan and Sanjari, Nasrin and Nourbakhsh, Ali},
  journal={arXiv preprint arXiv:2512.09069},
  year={2025}
}

License

This project is licensed under the MIT License.