TL; DR: прогнозирование на 2D-срезах; UNet с разделенными свертками, плотно соединенным расширенным узким местом и выпадением; ансамбль модели, обученной с модифицированными потерями Smooth L1, и модель, обученная в структуре GAN для передачи стилей. Код доступен на https://github.com/sanch7/OpenKBP-Solution.

Прежде всего, я невероятно благодарен своим товарищам по команде Джехи Чуну и Сен Яну. У меня был большой опыт работы с ними. Несмотря на то, что мы объединились очень поздно в соревновании, я многому научился у них, особенно у Джехи, который привнес в команду опыт в предметной области. Сказав это, я остановлюсь только на моей части нашего окончательного решения в этом посте. Я также благодарен организаторам за проведение этого интересного соревнования и особый привет Аарону за помощь мне в моем первом начинании в области лучевой терапии и планирования, основанного на знаниях.

Обзор задачи

Из официального описания OpenKBP - 2020 AAPM Grand Challenge,

Целью OpenKBP Challenge является продвижение объективных и последовательных сравнений методов прогнозирования доз для планирования, основанного на знаниях (KBP). Участники испытания будут использовать большой набор данных для обучения, тестирования и сравнения своих методов прогнозирования, используя набор стандартизованных показателей, с методами других участников.

Организаторы предоставляют трехмерную компьютерную томографию головы и шеи пациентов с диагнозом рака ротоглотки (головы и шеи). Для этих сканирований также предусмотрены маски сегментации для органов, подверженных риску (OAR) и плановых целевых объемов (PTV). Вместе они называются структурными масками (SM). Кроме того, также предоставляется двоичная маска вокселов, где может присутствовать ненулевая доза, называемая маской возможной дозы (PDM). Наряду с этим также указаны размеры вокселей, которые немного отличаются для каждого пациента.

Для тренировочного набора из 200 сканирований также предусмотрены предписанные дозы лучевой терапии с модуляцией интенсивности (IMRT). Задача состоит в том, чтобы узнать распределение доз по данным сканирований и доз и спрогнозировать их для тестового набора, состоящего всего из 140 сканирований. Соревнование использует два показателя для оценки: показатель дозы и показатель DVH. Их описание также можно найти на сайте конкурса и на странице github.

Окончательное описание решения

Данные

Проблемы трехмерной медицинской визуализации можно смоделировать либо для двухмерных срезов, либо для полного трехмерного объема. В своем решении я выбрал прогнозирование на двухмерных срезах, но понял, что окружающий трехмерный контекст чрезвычайно важен, и поэтому 4 смежных слайда с каждой стороны прогнозируемого среза добавляются в качестве дополнительных каналов для большего контекста. Предполагается, что сеть изучит пространственные отношения между этими каналами, что сделает задачу 2.5D.

Помимо компьютерной томографии, я включил SM и PDM для всех срезов. Следовательно, окончательный вход в сеть состоит из 12 x 9 = 108 каналов, а выход сети представляет собой прогноз дозы для одного канала для среднего слоя.

На поздних этапах, после формирования нашей команды, я добавил еще 20 каналов данных на срезы, которые были созданы на основе опыта Джехи в предметной области и описаны ниже:

- электронная плотность

- карта расстояний от ПТВ

- конусообразные псевдолучевые линии, представляющие положения и углы луча

- радиологическая глубина (информация о затухании) для лучей

Я решил не нормализовать компьютерную томографию или мишени, просто основываясь на результатах экспериментов. Используются дополнения данных Shift-Scale-Rotate и Flip.

Сетевая архитектура

Сеть представляет собой модифицированный 2D-UNet. Во-первых, все слои свертки заменяются разделяемыми свертками, чтобы сделать сеть более быстрой и эффективной. Во-вторых, плотно связанный блок из 4 расширенных сверточных слоев был вставлен в узкую часть U-Net. Эти слои имеют коэффициенты расширения 1, 2, 4 и 8, что помогает извлекать больше пространственной информации, чем это возможно при обычных свертках. Плотные связи помогают поддерживать эффективную передачу информации и градиентный поток в сети.

Также используется выпадение, при этом скорость отбрасывания линейно увеличивается с 0,375 на первом уровне кодера до 0,75 на последнем уровне, постоянная 0,75 во всем узком месте блока, а затем линейно уменьшается с 0,75 до 0,375 в декодере. Наконец, включен BatchNorm, который помогает стабилизировать тренировку.

Модель 1: обычное обучение под наблюдением с функцией потерь, как описано ниже -

Убыток состоит из двух частей. Первый компонент - это SmoothL1 Loss между прогнозом сети и целевой дозой. Потери вычисляются только в вокселях, которые включены в PDM.

Для второго компонента это также SmoothL1 Loss между сетевым прогнозом и целевой дозой, но вычисляется только в вокселях, которые включены в SM. Потери, полученные для каждого из 10 каналов, усредняются, и полученные потери SM умножаются на объем размера вокселя.

Затем рассчитываются окончательные потери, умножая второй компонент на коэффициент 3 и прибавляя его к первому компоненту.

Модель 2: перенос стиля GAN

Основная идея состоит в том, что мы считаем данные (КТ-сканирование) одним стилем, а выходные данные (дозу) - другим стилем. Затем для передачи между этими стилями используется GAN. Подробнее читайте в этой статье.

Детали обучения

Обе сети обучаются с использованием оптимизатора Rectified Adam со скоростью обучения 0,01 для 60 эпох с использованием расписания Fit 1 Cycle.

Постобработка

Используется левый-правый флип TTA.

Результат

Model1 набрала ~ 2,63 балла в таблице лидеров. Использование псевдо-маркировки с обучением преподаватель-ученик на тестовом наборе дает незначительное улучшение ~ 0,02. Model2 набрала ~ 2,60. Добавление 20 дополнительных каналов данных, описанных выше, улучшило оценку Model2 до 2,45, а наша финальная совокупность всех трех членов команды набрала ~ 2,41.

Путешествие в развитие

Я описал свою часть нашего окончательного решения выше. Если бы вас интересовало только это, можете перестать читать сейчас. Текст ниже содержит мою историю достижения этого решения. Как я уже упоминал, это был мой первый визит в мир компьютерной томографии, IMRT и KBP. Это был огромный опыт обучения, и я надеюсь, что описание итеративного процесса может быть интересно некоторым читателям. Как сказал Ральф Уолдо Эмерсон: «Это не пункт назначения, это путешествие».

Введение и обзор литературы

С этим соревнованием я столкнулся в середине марта. Он начался месяцем ранее, оставалось еще полтора месяца (изначально соревнования должны были завершиться 30 апреля). Задача казалась очень интересной, и мне хотелось иметь опыт моделирования компьютерной томографии. Эту проблему можно классифицировать как трехмерную регрессию, и найти какую-либо основную литературу по этой теме было сложно. Единственная литература, которую я смог найти, относилась к области KBP. Глубокое обучение в области KBP началось совсем недавно, но оно послужило полезным началом.

Организаторы любезно предоставили стартовый код для полного конвейера на основе 3D UNet. Код был в тензорном потоке, поэтому первым шагом для меня было преобразовать его в PyTorch. На мой взгляд, PyTorch - это гораздо более гибкий фреймворк для проведения исследований и быстрых экспериментов, которые часто требуются в соревнованиях по машинному обучению.

В результате моего исследования было обнаружено, что сообщество KBP разделено на 2D и 3D сети, оставляя достаточно места для экспериментов и разработок на обоих фронтах. Лично я хотел изучить 3D-моделирование. У меня также была гипотеза, что будет полезен полный трехмерный контекст, и поэтому я начал свои эксперименты с трехмерным конвейером.

Фон

Моя общая философия экспериментирования состоит в том, чтобы сначала сформировать гипотезу, а затем провести эксперименты, чтобы принять или отклонить ее. Это отличается от альтернативной философии, которая проводит эксперименты, а затем придумывает объяснения для описания наблюдений. На ограниченном оборудовании я обнаружил, что первое помогает сократить пространство поиска для экспериментов. Тем не менее, иногда стоит провести случайный эксперимент с интуицией и выяснить, что это работает :)

Это также хорошее место для обсуждения оборудования, которое я использовал, поскольку оно часто играет очень важную роль в глубоком обучении. Я работал на двух персональных машинах, оба с процессорами AMD, графическими процессорами Nvidia и 32 ГБ оперативной памяти. Первый имеет 6 ядер и 1080 Ti, а другой - 16 ядер и 2080 Ti. Помимо этого, я также приобрел подписку Colab Pro, которая стоит 10 долларов в месяц, и это отличное соотношение цены и качества. Обычно он предоставляет 4-ядерную виртуальную машину с графическим процессором P100, который может работать не более 24 часов. На заключительных этапах я также использовал GCP для нескольких модельных тренингов.

Что касается программного обеспечения, мне нравится PyTorch за его гибкость и питонический интерфейс. Совсем недавно я использовал Fastai, созданный на основе PyTorch, для обучения моих моделей. Мне нравится его чистый интерфейс, реализация политики обучения за один цикл и поддержка обучения FP16. Я использую веса и смещения для отслеживания своих экспериментов, которые также имеют аккуратную интеграцию с Fastai.

Для параллельного проведения экспериментов на разных машинах я использую GitHub. Это работает так - во-первых, я гарантирую, что все параметры в процессе обучения можно изменить из одного файла конфигурации. Если я хочу изменить обучающий код, я обычно сохраняю исходный код для воспроизводимости эксперимента и добавляю еще один параметр в файл конфигурации, который позволяет использовать новый код. При такой настройке довольно легко управлять запуском на разных машинах. Я могу просто начать эксперимент на первой машине, выполнить фиксацию и нажать, а затем задействовать вторую машину, внести все необходимые изменения, а затем запустить второй эксперимент и снова выполнить фиксацию. Это гарантирует, что вы можете, например, провести 5 раз конкретного эксперимента на одной машине, исследуя другие методы на других машинах, не беспокоясь об изменениях в коде.

3D моделирование

Возвращаясь к конкурсу, я сделал свою первую заявку в конце марта, которая была с простым 3D ResNet18 без понижающей дискретизации. Он набрал ~ 8,7 балла по метрике "Доза". Замена его на 3D UNet с кодировщиком ResNet18 улучшила оценку до ~ 7,2, а кодировщик Resnet34 дал оценку ~ 7,04. Дальнейшее увеличение размера кодировщика не улучшило производительность, они также медленно обучались и использовали много памяти графического процессора.

Следующее, что я решил настроить, - это оптимизатор. Я экспериментировал с SGD, Ranger (который дал мне значительно лучшие результаты по другим задачам), Adam и RAdam. Оказалось, что это не имело большого значения, и я придерживался RAdam до конца соревнований.

Следующее улучшение, как и ожидалось, произошло за счет увеличения объема данных. Чтобы смоделировать различные размеры вокселей, я добавил увеличение масштаба и стандартный Flip Left Right. Это повысило производительность до ~ 6,8. Простое изменение размера пакета с 1 до 4 (благодаря обучению Fastai FP16, которое помогло подобрать больший размер пакета на GPU) улучшило его до ~ 6,46.

Следующее, что нужно настроить, - это функция потерь. До этого момента я использовал потерю L1, которая соответствовала метрике конкуренции, которая также использовала расстояние L1. Первое изменение заключалось в использовании асимметричных потерь, которые не ухудшали модель, если выходная доза в вокселе, принадлежащем PTV, была больше целевой. Точно так же модель не подвергалась штрафу, если выходная доза в вокселе, принадлежащем SM, была меньше целевой. Оказалось, что даже несмотря на то, что эта потеря - это то, что мы хотим на практике, она не подходила для конкуренции, поскольку показатель конкуренции наказывал бы модели за «лучше», чем эталонная доза. Так что эта потеря была отброшена.

Поскольку между тренировкой и валидацией была большая разница, до сих пор я использовал снижение веса на 0,1. Но при дальнейшем экспериментировании стало ясно, что это мешает модели соответствовать данным и, следовательно, снижает оценку валидации. Удаление снижения веса повысило производительность до 6,23. Кроме того, добавление LR Flip TTA дало небольшое улучшение и подняло оценку до 6,17.

Это момент, когда я впервые почувствовал насыщение от выступления. Не было очевидного параметра оптимизации, который можно было бы точно настроить, кроме самой модели, но я был ограничен в памяти графического процессора. После кучи разочаровывающих экспериментов, которые включали такие вещи, как увеличение данных трехмерного вращения (которое потребляло много ресурсов ЦП и делало обучение очень медленным), с использованием модуля Res2Net и других модификаций сети, я отказался от трехмерного подхода. .

Переход на 2D

Переход на 2D сразу показался хорошим ходом. Но сначала я хочу поговорить о моей локальной настройке проверки. Метрики в этом соревновании было немного сложно оценить. Они работают не для независимых срезов, а для полного пациента. Это не было проблемой с 3D-подходом, но с 2D, при использовании любого размера пакета, отличного от 128, показатели будут неточными. Сначала я попытался установить размер пакета загрузчика данных проверки равным 128, независимо от размера пакета обучения. Но по мере увеличения размера модели я больше не мог помещать партии размером 128 на GPU. Очевидным решением было накопление пакетов и оценка показателей только тогда, когда было накоплено 128 срезов (которые соответствуют полному пациенту). Fastai не поддерживает это из коробки, поэтому реализация этого дня была веселой, когда мне пришлось повозиться с исходным кодом Fastai и, в свою очередь, много узнать о дизайне библиотеки. Это определенно того стоило, и, на мой взгляд, создание надежного локального конвейера валидации - это самое важное, что нужно делать при глубоком обучении. В конце концов, вы не можете улучшить то, что не можете измерить.

Моя первая попытка работы с срезами заключалась в использовании одного среза КТ для прогнозирования дозы. Он набрал ~ 3,77 балла по метрике дозы, с теми же настройками, что и 3D-модель, размером пакета 128, с увеличением данных ShiftScaleRotate и Flip, плавной потерей L1 и политикой обучения за один цикл для 60 эпох. Время обучения модели сокращено с 30 до 6 часов. С одной стороны, это было невероятное улучшение для меня, когда я достиг лучшей оценки дозы ~ 6,17 с 3D-моделью. Внезапно я переместился в топ-15, и достойный финиш казался возможным. Но с другой стороны, я потратил около месяца на работу с 3D-моделями (большая часть из которых была просто тренировкой), а теперь до соревнований оставалось всего 15 дней.

Я начал заниматься постобработкой модели, которую только что обучил. Значительное улучшение произошло из-за того, что в этой модели срезы не имели информации о смежных срезах. Таким образом, простой сглаживающий фильтр по выходным дозам улучшил оценку до ~ 3,63. Отметив это наблюдение, я продолжил настаивать на производительности одной модели.

Поскольку между моей производительностью набора для обучения и проверки все еще была большая разница, первое изменение, которое пришло в голову, было добавление дополнительной регуляризации. Я добавил в сеть отсев, и это сразу показалось мне хорошим выбором. Последний график выпадения, который я использовал, представлял собой выпадение 0,375 в первом слое кодировщика, которое линейно увеличивается до 0,75 до уровня узкого места, постоянное выпадение 0,75 в узком месте и линейно уменьшающееся выпадение обратно до 0,375 в последнем слое. декодера. Это улучшило мою оценку до 3,70.

Затем, снова надеясь уменьшить переобучение на обучающем наборе, я добавил дополнительные данные в виде аддитивного и мультипликативного шума. Я также внес небольшие изменения в архитектуру, например, используя новую функцию активации Mish или упорядочив Conv - ›ReLU -› BN вместо обычного Conv - ›BN -› ReLU (который, согласно соответствующей публикации, имел значительную влияние на их производительность). Все эти эксперименты не привели к дальнейшему улучшению результатов.

Чтобы проверить гипотезу о важности пространственного контекста, я добавил во входные данные соседние слайды. Добавление 1 среза в каждом направлении к входному срезу значительно улучшило версию 3.22. Однако позже выяснилось, что этот скачок произошел не только из-за пространственного контекста. Чтобы уместить данные с еще двумя срезами (еще 24 канала, включая PDM и SM) в графическом процессоре, мне пришлось уменьшить размер пакета до 32. Оказалось, что уменьшение размера пакета было очень эффективным регуляризатором в этом случае, поскольку оно добавили в обучение больше стохастичности. Это было дополнительно подтверждено уменьшением размера пакета до 8, что было компромиссом между временем обучения (время обучения почти прямо пропорционально обратной величине размера пакета) и производительностью модели. В конечном итоге я решил включить по 4 среза с каждой стороны, то есть всего 9 срезов в сэмпл. Это привело к результату 2,96 в таблице лидеров.

Следующее, что я изменил, - это функция потерь, чтобы придать больший вес вокселям, принадлежащим любой из структурных масок. Это было сделано потому, что метрика DVH фокусируется только на вокселях, принадлежащих либо OAR, либо PTV, и поэтому гипотеза заключалась в том, что придание большего веса этим вокселам улучшит оценку DVH. Я назвал второй компонент потерь SM Loss, который также был Smooth L1, и присвоил ему вес 3. Добавление этого не только улучшило оценку DVH, но, что удивительно, также улучшило оценку дозы до 2,90.

Следующие неудачные эксперименты включали случайное кадрирование, добавление позиции среза к входным данным, то есть индекса среза на компьютерной томографии, с гипотезой, что предоставление этой информации могло бы лучше помочь сети определить геометрию. По аналогичной гипотезе я добавил индексы x и y перед выполнением преобразований, что также не помогло. Я также попытался переключить ось CT-срезов с головы - туловища на лево - право и грудь - назад, но они даже не приблизились к характеристикам исходной модели. Повторная выборка загрузчика данных путем придания большего веса выборке фрагментам с большим количеством ненулевых вокселей также не дала никаких улучшений.

Затем я вернулся к идее увеличения данных, и на этот раз я специально изучил медицинские дополнения. Я пробовала эластичные и диффеоморфные аугментации, но они тоже никак не улучшили результат.

Продление крайнего срока конкурса

К настоящему времени мы приближались к крайнему сроку конкурса, и я был уверен, что попаду в топ-10. Но как только я собирался заморозить свою единственную модель и создать пятикратный ансамбль, организаторы решили продлить конкурс еще на 4 недели из-за того, что многие участвующие исследовательские группы не смогли получить доступ к своему оборудованию из-за закрытия университетов. Итак, эксперименты возобновились ...

Большинство следующих нескольких экспериментов с различными стратегиями обучения и другими простыми настройками не дало никаких улучшений. Следующим крупным экспериментом, на который я возлагал большие надежды, была передискретизация всех сканированных изображений до однородного размера вокселя. Во-первых, я передискретизировал все входные данные до однородного размера вокселя 4,5 x 4,5 мм в направлении нарезки и сохранил размер вокселя по умолчанию в направлении нарезки. Обучение и проверка прошли очень хорошо. Но, к сожалению, когда я попытался сделать прогнозы на тестовом наборе, было трудно восстановить прогнозы для исходной выборки из передискретизированного ввода. Такие методы, как линейная или кубическая интерполяция для восстановления исходной выборки, были неоптимальными и работали хуже, чем исходная модель. После очередного неудачного эксперимента с размером вокселя 2,25 x 2,25 я отказался от этого многообещающего подхода.

До этого момента я не исследовал различные архитектуры моделей, поэтому я начал изучать более свежие модели, такие как DeepLab V3, FPN и PSPNet. Но ни одна из этих новых архитектур не работала лучше, чем UNet.

Разделимые свертки и объединение

Почти для всех задач классификации я по умолчанию выбираю архитектуру сети Xception. В частности, я поклонник раздельной свертки, которая не только сокращает время обучения и размер модели, но также неявно создает эффект регуляризации как следствие менее обучаемых параметров. Поэтому я решил попробовать Separable Convolutions и, наконец, совершил прорыв. Эта новая модель получила 2,80 балла. Более того, теперь у меня было больше места на графическом процессоре, и я решил увеличить ширину сети в два раза по сравнению с исходной, что еще больше улучшило оценку до 2,76. На этом этапе я решил составить с этой моделью ансамбль в пять складок. Ансамбль набрал 2,63 балла в таблице лидеров, что поставило меня на 4 позицию в таблице лидеров!

Дополнительные эксперименты с архитектурой модели, такие как дальнейшее увеличение ширины, изменение количества расширенных блоков в узком месте, изменение размера ядра и добавление сжатия и возбуждения, не помогли.

Здесь я решил использовать псевдометки для набора проверки. Добавление псевдо-меток и создание 5-кратного ансамбля немного улучшили результат до 2,62. Двигаясь дальше, я решил провести обучение учителя и ученика вместо псевдо-ярлыков на тестовом наборе, с предыдущими лучшими моделями, действующими в качестве учителей. Пятикратный ансамбль из этого также немного улучшил оценку до 2,61.

Теперь до завершения фазы валидации оставалось 5 дней. Мне было нечего терять, поэтому я решил инвестировать в совершенно новую парадигму обучения, используя GAN в структуре передачи стиля, которая была основана на этой работе. Удивительно, но это сразу дало мне небольшое улучшение, а пятикратный ансамбль улучшил результат до 2,60, и теперь я стал третьим в таблице лидеров.

Поскольку до завершения фазы проверки оставалось два дня, я, Джехи и Сен собрались вместе, чтобы сформировать команду. Как оказалось, Джехи использовал трехмерный подход, и поэтому было очевидно, что объединение наших моделей улучшит производительность. Простая комбинация моей модели и модели Джехи (оценка 2,61) улучшила нашу оценку до 2,45.

Кроме того, я получил специально разработанные данные Джехи для конкретных областей, которые включали такую ​​информацию, как электронная плотность, формы псевдолуча, радиологическая глубина (информация о затухании) и карты расстояний от PTV. Простое добавление этих данных к моей модели улучшило мою производительность в 5 раз до 2,45. Джехи и Сен также улучшили свои модели на основе обратной связи, которую мы дали друг другу, и наш окончательный ансамбль набрал 2,417 балла на проверочном наборе. Мы оценили набор скрытых тестов на наших моделях и представили наши прогнозы.

В конце соревнования мы оказались на второй позиции в общедоступной тестовой выборке, где лидеры набрали 2,309, а группа, занявшая третье место, - 2,499. Учитывая эту производительность, мы ожидали, что мы займем второе место, но в скрытом тестовом наборе якобы произошла некоторая встряска, и мы оказались на третьей позиции в частном тестовом наборе.

Заключение

Оглядываясь назад, когда я решил участвовать, я бы никогда не ожидал, что я финиширую на третьем месте в конкурсе, особенно без предварительного опыта моделирования компьютерной томографии или лучевой терапии и с участием многих радиологических исследовательских групп в соревновании. Моя цель на конкурсе состояла в том, чтобы узнать как можно больше о глубоком обучении с медицинскими данными, и я очень доволен богатым опытом обучения, который я получил за последние два с половиной месяца.

Ваши мысли и отзывы приветствуются!