Unet-boosted classifier – мультизадачная архитектура для малых выборок на примере классификации МРТ снимков головного мозга

Проблема обучения глубоких нейронных сетей на малых выборках особенно актуальна для медицинских задач. В работе рассматривается влияние попиксельной разметки значимых объектов на изображении, в дополнении к истинной метке класса, на качество решения задачи классификации. Для достижения лучших результатов классификации на малых выборках предлагается мультизадачная архитектура Unet-boosted classifier (UBC), обучаемая одновременно для решения задач классификации и семантической сегментации. В качестве исследуемого набора данных используются МРТ-снимки пациентов c доброкачественной глиомой и глиобластомой, взятые из открытого набора данных BraTS 2019. В качестве входа рассматривается один горизонтальный срез МРТ-изображения, содержащий глиому (всего 380 кадров в обучающей выборке), в качестве выхода – вероятность глиобластомы. В качестве базового решения используется ResNet34, обученный без аугментаций с функцией потерь на основе взаимной энтропии (CrossEntropyLoss). В качестве альтернативного решения используется UBC-ResNet34 – тот же ResNet34 усиленный декодером, построенным по принципу U-Net, и предсказывающим положение глиомы. В качестве дополнительной функции потерь используется сглаженный коэффициент Соренсена-Дайса (DiceLoss). Результаты на тестовой выборке: доля правильных ответов (accuracy) для базовой модели составила 0.71, для альтернативной – 0.81, коэффициент Дайса (Dice score) при этом составил 0.77. Таким образом, глубокую модель можно качественно обучить даже на небольшом наборе данных, используя предложенную архитектуру и добавив в разметку информацию о пораженных тканях в виде семантической маски. Предлагаемый подход потенциально может быть полезен и в любых других задачах классификации изображений с ограниченным набором данных.