Ai-Adam-Six-Sigma/melanoma-efficientnet-b4

Melanoma Classifier — EfficientNet-B4

Binary classifier for dermoscopic images: malignant (melanoma) vs benign (nevus).

Model

• Architecture: EfficientNet-B4 (timm library), 380x380 input
• Transfer learning: ImageNet pretrained weights, fine-tuned on dermoscopic data
• Output: 1 neuron (logit) → sigmoid → probability of the malignant class

Training Data

• Dataset: ISIC 2019 + 2020 (mel vs nevus), ~11,400 images after deduplication
• Class distribution: mel 44.6% / nevus 55.4%
• Preprocessing:
  • duplicate removal (MD5 hash)
  • class weighting (correcting the mild class imbalance)
  • augmentation correcting a resolution gap between classes (nevus images had on average ~3x higher resolution than mel in the source data — this augmentation reduces the risk of the model learning that difference instead of actual lesion features)
  • standard augmentations: rotation, flip, brightness/contrast jitter

Training

• Fine-tuning from ImageNet weights, AdamW, BCEWithLogitsLoss with pos_weight (class weighting)
• Early stopping (patience=5 epochs without recall improvement)
• Best checkpoint selected by recall on the malignant class (not by loss)

Metrics

Validation set (ISIC, 20% hold-out, stratified):

Metric	Value
Recall (malignant)	93.4%
Precision (malignant)	89.6%
AUC	0.980

Test on an independent sample of 505 real malignant images:

Metric	Value
Recall (malignant)	93.9% (474/505 detected)
Missed (false negatives)	31/505

Comparison with the previous version (EfficientNet-B0):

Model	Recall on 505 samples
EfficientNet-B0 (previous)	84.4%
EfficientNet-B4 (this model)	93.9%

Usage

import torch
import timm
from torchvision import transforms
from PIL import Image

model = timm.create_model('efficientnet_b4', pretrained=False, num_classes=1)
state_dict = torch.load('efficientnet_b4_melanoma_final.pt', map_location='cpu')
model.load_state_dict(state_dict)
model.eval()

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

img = Image.open('path/to/image.jpg').convert('RGB')
x = transform(img).unsqueeze(0)

with torch.no_grad():
    logit = model(x)
    prob_malignant = torch.sigmoid(logit).item()

print(f"Malignant probability: {prob_malignant:.2%}")

Limitations

• This model is a decision-support tool, not a substitute for medical diagnosis.
• A recall of 93.9% means roughly 6% of true malignant cases are still missed.
• Trained on dermoscopic images — performance on other image sources (e.g. regular phone photos) has not been evaluated.
• No formal testing has been done for demographic bias (skin type, age, lesion location).

Author

Adam Sobański — Ai-Adam-Six-Sigma