Erfan-Nourbakhsh/KD-OCT-EfficientNet

KD-OCT: Knowledge Distillation for OCT Image Classification

This model is part of the KD-OCT project, introduced in the paper KD-OCT: Efficient Knowledge Distillation for Clinical-Grade Retinal OCT Classification.

Model Description

KD-OCT is a knowledge distillation framework designed to compress a high-performance ConvNeXtV2-Large teacher model into a lightweight EfficientNet-B2 student model. This approach is optimized for the clinical-grade classification of retinal Optical Coherence Tomography (OCT) images, specifically targeting conditions like age-related macular degeneration (AMD) and choroidal neovascularization (CNV).

• Task: Multi-class classification (Normal, Drusen, CNV, and DME)
• Student Architecture: EfficientNet-B2
• Teacher Architecture: ConvNeXtV2-Large
• Goal: Enabling real-time deployment on edge devices for clinical screening.

Training Details

• Framework: PyTorch
• Method: Real-time knowledge distillation using a combined loss (soft teacher knowledge transfer + hard ground-truth supervision).
• Optimization: Focal loss for class imbalance and Stochastic Weight Averaging (SWA).
• Datasets: Evaluated on Noor Eye Hospital (NEH) and UCSD datasets.

Usage

The following code snippet demonstrates how to load the model and perform inference using PyTorch:

import torch
from torchvision import transforms

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

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

# Inference (example for a single image tensor 'input_tensor')
with torch.no_grad():
    # input_tensor = transform(image).unsqueeze(0)
    output = model(input_tensor)
    prediction = torch.argmax(output, dim=1)

Citation

If you use this model in your research, 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}
}

Links

• Paper: Hugging Face Paper Page
• Repository: Official GitHub Repository

License

This project is licensed under the MIT License.