Back to Search
Start Over
SkinDistilViT: Lightweight Vision Transformer for Skin Lesion Classification
- Publication Year :
- 2023
-
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.<br />Comment: Accepted at ICANN 2023
Details
- Database :
- arXiv
- Publication Type :
- Report
- Accession number :
- edsarx.2308.08669
- Document Type :
- Working Paper