SkinDistilViT:用于皮肤病变分类的轻量级视觉Transformer
计算机视觉与模式识别
2023-08-21 v1 机器学习
摘要
皮肤癌若及早发现是可治疗的疾病。我们提供了一种面向生产环境的皮肤癌分类解决方案,通过在专家标注的黑色素瘤医学图像上训练视觉Transformer,在黑色素瘤识别上达到与人类相当的性能。由于推理成本(时间和内存方面)在实践中十分重要,我们采用知识蒸馏来获得一个模型,该模型以极小的代价保留了教师模型98.33%的均衡多类准确率。在内存方面,我们的模型比教师模型小49.60%。在时间方面,我们的解决方案在GPU上快69.25%,在CPU上快97.96%。通过在Transformer的每一层添加分类头并采用级联蒸馏过程,我们将基础模型的均衡多类准确率提高了2.1%,同时生成了一系列大小不同但性能相当的模型。我们在 https://github.com/Longman-Stan/SkinDistilVit 提供了代码。
引用
@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}
}
备注
Accepted at ICANN 2023