Целью своей работы я поставил сравнение различных подходов при решении конкретной задачи, а именно детекции микроорганизмов. Мне показалось, что это весьма актуальная и интересная тема, которой было бы интересно заниматься.
Распознавание и автоматическая детекция различных микроорганизмов (включая амёб, бактерий, грибков и других) имеют широкое применение в современной медицине и биомедицинских исследованиях в связи с потребностью в анализе большого количества лабораторных данных. Это может быть необходимо для диагностики инфекционных заболеваний, санитарно-гигиенического контроля, исследования микробиомы человека, определения чувствительности к антибиотикам.
В процессе поиска я выбрал датасет Clinical Bacteria DataSet, который содержит изображения бактерий(бацилл и кокков), окрашенных по Граму, полученные с помощью M-ROSE(микробиологическая экспресс-оценка на месте) из образцов нижних дыхательных путей пациентов с легочной инфекцией. Данные были собраны в период с 2018 по 2022 год в больнице общего профиля НОАК Китая(Chinese PLA General Hospital). Представлено 4833 кокков и 6991 бацилла, которые разделены на отрицательные и положительные категории(бактерии отличаются строением и устойчивостью к антибиотикам, положительные менее устойчивы): 3371 грамотрицательных кокков (G-кокков), 1462 грамположительных кокков (G+кокков), 5799 грамотрицательных бацилл (G-бацилл) и 1192 грамположительных бацилл (G+бацилл). Положительные окрашены в синий цвет, а отрицательные в красный. Ground Truth в датасете даны в YOLO формате(к каждому изображению есть label .txt файл)
https://zenodo.org/records/10526360 https://github.com/Quanlab-Bioimage/Clinical-Bacteria-DataSet
В качестве основы для начала исследования я написал архитектуру, напоминающую RetinaNet. В backbone была взята предобученная на ImageNet-1K модель EfficientNet B0, шея - FPN со skip connections с выходами из backbone и decoupled head в качестве головы. Осуществил загрузку данных и предобработку, применил базовые аугментации. Т.к. в исходных данных наблюдается дисбаланс классов, то деление на обучающую и валидационную выборки выполняется через IterativeStratification. Для решения label assignment я реализовал task alignment learning и также использую алгоритм, который работает на основе расчёта метрики IoU между якорями и Ground Truth. Для расчёта общей функции потери я выбрал SmoothL1 для локализационной части, Focal Loss для ошибки по уверенности (confidence score) и CrossEntropyLoss для классификационной части. При обучении я использую шедулер CosineAnnealingLR и оптимизатор Adam. К сожалению в процессе всей работы я допускал ошибку, которую поздно заметил и использовал вместо CrossEntropyLoss для классификационной части BCEWithLogitsLoss. Все изложенные ниже действия были осуществлены с этой ошибкой. Первой попыткой обучения был запуск на 10 эпох при батче размером 64, в результате которого train loss снижался и была получена итоговая метрика в 0,03 mAP. Сразу же после я запустил обучение на 100 эпох при батче 12, чтобы узнать существующий максимум того, что я могу получить при текущей архитектуре и гиперпараметрах. Итогом стало получение метрики mAP в 0,05. backbone в обоих случаях был разморожен лишь частично, обучались только 2 последних блока из нескольких слоёв. Также в эти две попытки я использовал BCE для расчёта лосса по confidence score и не пытался бороться с дисбалансом классов. После второго запуска я вернул размер батча к 64, добавил веса в функцию потерь для классов и заменил BCE на Focal Loss, полностью разморозил backbone и понизил скорость обучения(была 1e-3, стала 1e-4). В итоге за 100 эпох получилось достичь 0,14 mAP. Я заметил, что очень малое число якорей в итоге отбираются к обучению как положительные и решил исправить эту проблему, реализовав в качестве метода label assignment task alignment learning. Этот метод действительно позволил кратно увеличить число якорей в обучении, но прироста по метрике я не обнаружил при попытках обучения на 10 и 20 эпохах и поэтому не стал дальше использовать его. Я решил, что мне следует укрупнить модель, чтобы поднять возможный потолок по метрике. В новой архитектуре я заменил backbone на EfficientNet B7(т.к. надеялся, что крупнейшая модель из семейства кратно улучшит результат), убрал skip connections между backbone и neck, но оставил их в шее между слоями, таким образом реализовав шею в виде hourglass, также я уменьшил итоговую выходную feature map и число якорей соответственно с 40x40 на 20x20. Изначально я опрометчиво планировал обучать EfficientNet B7 полностью разморозив, но скоро понял, что этого у меня не получится из-за колоссальных требований по видеопамяти, поэтому я размораживал лишь 2 последних блока сети. Метрики в итоге были около нуля. Я вернул шею к прежнему виду и skip connections между шеей и backbone, метрики удалось немного поднять от нуля, но значений как ранее не было. Я сделал вывод, что EfficientNet B7 слишком крупная по сравнению с шеей и головой и т.к. я не мог обучать её полностью, то по итогу сеть не могла адаптироваться под новую задачу. Я решил заменить EfficientNet B7 на EfficientNet B2, также я вернул размер итоговой feature map к 40x40, хотя разницы в качестве от количества якорей не заметил. EfficientNet B2 я стал размораживать полностью, получил 0,0175 mAP за 10 эпох и 0,07 mAP за 100 эпох. Вероятно возможный потолок по метрике с B2 выше, чем с B0, но модель требует больше времени для обучения. Таким образом, я не смог добиться действительно высоких значений mAP и качественно решить задачу, хотя многое попробовал для этого, но получил опыт работы с детектором. В дальнейшем я мог бы предпринять ещё попытки для улучшения результата. Также в процессе экспериментов я забыл установить одно зерно для всех случайных генераций, в том числе для инициализации весов модели, теперь я это исправил. В результате исправления функции потерь и добавления новых аугментаций удалось через дообучение модели на прежних весах достичь mAP в 0.165 на модели с backbone EfficientNet B0. Для вывода лучше всего использовать следующие пороги или приблизительно такие: score_threshold = 0.3, nms_threshold = 0.2. Также я допустил ещё одну неточность, в методе validate необходимо было использовать либо очень маленький confidence threshold, либо нулевой. При score_threshold=0.01 итоговое значение метрики чуть точнее и равно 0.168.
https://drive.google.com/drive/folders/1TQElKKzMS3nzmOaW3D2cy7qaXpaiOwnb?usp=sharing