Архитектура U-Net, представленная в 2015 году, стала настоящим прорывом в области сегментации биомедицинских изображений и впоследствии нашла широкое применение в различных задачах компьютерного зрения. Основной целью U-Net является разделение входного изображения на сегменты — например, отделение клеток от фона — с помощью сверточных нейронных сетей. Уникальность модели заключается в своеобразной структуре, сочетающей сжимающий и расширяющий пути, объединённых между собой скип-соединениями, которые обеспечивают передачу информации с высокого разрешения на соответствующие этапы декодирования. Воссоздание этой архитектуры на основе бумаги 2015 года позволяет не только ознакомиться с её внутренним устройством, но и критикуя оригинальные решения, улучшить понимание и адаптировать модель под современные задачи. Суть архитектуры U-Net заключается в последовательном уменьшении пространственного разрешения входного изображения с помощью сверточных слоев и операций пулинга, что позволяет модели концентрироваться на изучении признаков высокого уровня и расширять каналы представления.
После этого происходит обратный процесс — расширение изображения посредством транспонированных свёрток и объединение соответствующих слоёв с помощью скип-соединений. Это даёт возможность сохранить детальную информацию, важную для точной сегментации, и одновременно учитывать контекст на более высоких уровнях абстракции. Важно отметить, что в публикации был использован входной размер 572 на 572 пикселя, однако распространённые датасеты имеют размеры, например, 512 на 512, что создаёт определённые трудности с выравниванием размеров признаков, особенно из-за отсутствия паддинга в свёртках. Техника центрированного кропа или обрезки обычно используется для подгонки размеров слоёв скип-соединений со слоями декодера, гарантируя совпадение пространственных размеров перед конкатенацией по канальному измерению. Такой подход не прописан в оригинальной статье явно, но является необходимостью при реализации модели без добавления паддинга, что отражает один из основных вызовов при воспроизведении работы исключительно из текста статьи.
Особое внимание уделяется стратегии обработки больших изображений при ограниченных ресурсах видеопамяти. В оригинальном исследовании указано использование «overlap-tiling» — разбиения больших изображений на перекрывающиеся тайлы фиксированного размера, что позволяет вычислять маски для частей изображения по отдельности с обязательной усреднённой агрегацией результирующих предсказаний. При этом каждая маска предсказывается для центральной части тайла, чтобы избежать артефактов на границах. Данный подход требует отражения и дополнения краевых областей изображения для сохранения контекста при обработке угловых тайлов, что является важной деталью для адекватного обучения и применения модели. Еще одним важным элементом реализации стала функция потерь, которая выходит за рамки классической кросс-энтропии.
В частности, в работе используется взвешенный вариант потерь для набора данных EM (электронная микроскопия). Весовые карты создаются с помощью морфологических операций, выделяя границы клеток, что позволяет акцентировать внимание модели на наиболее сложных для сегментации областях — межклеточных границах. Это является формой внедрения индуктивных биасов — в некотором смысле искусственного фокуса на ключевых зонах, что, однако, ставится под сомнение при воспроизведении, поскольку практическая выгода от такого подхода оказывается неочевидной по итогам экспериментов. Для обучения сети в оригинальном исследовании применялся стохастический градиентный спуск с моментумом 0.99 и фиксированной, хотя и не явно обозначенной скоростью обучения.
Использование различных оптимизаторов, таких как Adam, становится одной из ключевых тем для дальнейших экспериментов, выявляя преимущества и недостатки каждого из них в контексте малого объема данных и склонности к переобучению. Данные — это ещё один аспект, требующий особого рассмотрения. Датасет EM содержит всего около 30 высококоррелированных изображений клеток, что подчёркивает необходимость использования продвинутых методов аугментации данных. Именно благодаря таким приёмам, как эластичные деформации, случайные отражения и масштабирования, а также формирование тайлов с различными сдвигами, удаётся создать достаточно разнообразный обучающий набор, минимизируя переобучение и улучшая обобщаемость модели. Важность интересного и разнообразного аугментационного конвейера в задачах с ограниченным объёмом данных трудно переоценить.
Эксперименты, проведённые в ходе воссоздания модели, показывают интересные закономерности. Например, при одинаковом числе обрабатываемых пикселей за один шаг обучения, большая площадь тайлов предпочтительнее и даёт более плавные и стабильные показатели потерь и метрик качества сегментации. Это связано с тем, что меньшие тайлы, хотя и обеспечивают большую вариативность за счёт увеличения количества батчей, вследствие меньшего объёма контекста приводят к более шумным градиентам. Еще одним важным выводом стало то, что взвешенные весовые карты в функции потерь для EM датасета не дают заметного преимущества по метрикам качества, таким как Intersection Over Union (IoU), что ставит под вопрос целесообразность их использования. Возможно, подобный индуктивный биас можно заменить или дополнить современными методами регуляризации и сложными архитектурными решениями.
Использование разных оптимизаторов тоже значительно влияет на процесс обучения. Adam демонстрирует более быстрое сходимость, однако склонен к перенастройке и переобучению, особенно при высоких скоростях обучения. В свою очередь, классический SGD с моментумом оказывается более устойчивым на небольших выборках, хотя обучение занимаем больше времени. Использование адаптивного снижения скорости обучения для Adam позволяет избежать срывов и получать высокое качество сегментации на валидационных данных. Стоит также отметить ограничения, связанные с отсутствием открытого доступа к тестовым аннотациям, что усложняет объективную оценку качества модели по стандартным метрикам.
В результате приходится ориентироваться на результаты на валидационных выборках либо на общепринятые метрики потерь, что предполагает определенную погрешность и сдерживает полностью объективную валидацию. Исследование воспроизводимости U-Net ставит акцент на важности прозрачности описания архитектур и протоколов обучения, а также необходимости доступности репозиториев с исходным кодом и данными. Особенно важной становится внимательность при настройке деталей, таких как размер входных изображений, согласование размеров в скип-соединениях, выбор функции потерь, а также стратегия аугментаций и гиперпараметров обучения. Разбор исторического подхода U-Net показывает, что, несмотря на возможные шероховатости и недосказанности в оригинальной публикации, базовая архитектура оказалась чрезвычайно успешной и вдохновила массу последующих исследований и модификаций. Воссоздание модели без прямого обращения к исходному коду позволяет глубже понять архитектурные решения и способствует развитию критического мышления у исследователей.
В итоге, проект по воссозданию U-Net демонстрирует сложность и многогранность задач, связанных с воспроизведением независимых исследований в области глубокого обучения. Важно помнить, что даже при использовании тщательного переноса параметров и годами проверенных протоколов всегда остаются детали, которые влияют на результат, и именно критический подход и экспериментальное подтверждение позволяют найти лучшие практики. Для специалистов, работающих с биомедицинской сегментацией и смежными областями компьютерного зрения, понимание архитектуры U-Net, её фаз обучения, особенностей обработки входных данных и используемых методов оптимизации — ключ к созданию эффективных, надёжных и адаптивных решений. В преддверии новых вызовов и развития технологий важно уметь не только использовать готовые модели, но и глубоко анализировать и воспроизводить классические архитектуры, что служит фундаментом для создания инноваций и повышения качества диагностики и анализа изображений в различных дисциплинах.
