Тестовое задание для компании AIDiagnostic.
Бейзлайн для обучения сегментационной нейросети для обнаружения плеврального выпота на основе 3D архитектуры Unet. В качестве обучающей выборки представлено всего 10 снимков, что для такой задачи очень мало. Полученна метрика Dice=0.36 на кросс-валидации (обучение на 8 снимках, валидация на двух снимках).
Краткое описание пайплайна:
- Чтение данных из каталогов, указанных в config.py.
- Приведение изображений к одному размеру (обрезка по краям/добавление паддинга) и нормализация значений в тензоре (preprocessing.py).
- Формирование pytorch датасетов (train/val) и даталоадеров (training.py, dataset.py) c применением аугментаций.
- Обучение модели (model.py, training.py) с валидацией после каждой эпохи, с применением "разогрева" learning rate. Методы train_epoch и eval_epoch принимают список метрик, которые рассчитываются и сохраняются в истории.
- В процессе обучения (main.py): вывод номера эпохи, learning rate, текущих метрик, длительности эпохи, сохранение весов лучшей эпохи, вывод и сохранение на диск графика обучения в конце.
Настройки (main.py):
Выбор "масштаба" модели: глубина сети постоянна, количество фильтров в каждом сверточном слое пропорционально значению «масштаба». Это значение следует увеличивать, если данных много, и уменьшать в случае переобучения или для снижения ресурсоемкости сети.
Выбор аугментаций.
Выбор размера батча.
Выбор параметров обучения: максимальный learning rate, количество эпох для разогрева, скорость ументшения learning rate после разогрева, общее количество эпох обучения.