English

SkinDistilViT: Lightweight Vision Transformer for Skin Lesion Classification

Computer Vision and Pattern Recognition 2023-08-21 v1 Machine Learning

Abstract

Skin cancer is a treatable disease if discovered early. We provide a production-specific solution to the skin cancer classification problem that matches human performance in melanoma identification by training a vision transformer on melanoma medical images annotated by experts. Since inference cost, both time and memory wise is important in practice, we employ knowledge distillation to obtain a model that retains 98.33% of the teacher's balanced multi-class accuracy, at a fraction of the cost. Memory-wise, our model is 49.60% smaller than the teacher. Time-wise, our solution is 69.25% faster on GPU and 97.96% faster on CPU. By adding classification heads at each level of the transformer and employing a cascading distillation process, we improve the balanced multi-class accuracy of the base model by 2.1%, while creating a range of models of various sizes but comparable performance. We provide the code at https://github.com/Longman-Stan/SkinDistilVit.

Keywords

Cite

@article{arxiv.2308.08669,
  title  = {SkinDistilViT: Lightweight Vision Transformer for Skin Lesion Classification},
  author = {Vlad-Constantin Lungu-Stan and Dumitru-Clementin Cercel and Florin Pop},
  journal= {arXiv preprint arXiv:2308.08669},
  year   = {2023}
}

Comments

Accepted at ICANN 2023