Создание синтетических наборов данных ранений с помощью генеративных состязательных сетей

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

Система, называемая WG²AN, является совместным проектом между Баттенским колледжем инженерии и технологий и компанией еКare, специализирующейся на применении методов машинного обучения для измерения и идентификации ранений.

Генеративная состязательная сеть обучается на 100-4000 помеченных стереоскопических хронических изображениях ранений, предоставленных еКare, включая анонимные фотографии типов травм, вызванных давлением, операциями, лимфоваскулярными инцидентами, диабетом и ожогами. Исходный материал варьировался в размере от 1224×1224 до 2160×2160, все снятое при доступном свете врачами.

Чтобы приспособиться к доступному латентному пространству в архитектуре модели, изображения были изменены в размере до 512×512 и извлечены из их фонов. Чтобы изучить эффект размера набора данных, тестовые запуски были реализованы на партиях из 100, 250, 500, 1000, 2000 и 4000 изображений.

Source: https://ietresearch.onlinelibrary.wiley.com/doi/pdfdirect/10.1049/tje2.12033

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

The architecture of WG2GAN. Source: https://ietresearch.onlinelibrary.wiley.com/doi/pdfdirect/10.1049/tje2.12033

WG²AN работает на PyTorch на относительно скромной потребительской установке, с 8 ГБ видеопамяти на GTX 1080 GPU. Обучение заняло от 4 до 58 часов в диапазоне размеров наборов данных от 100 до 4000 изображений и в диапазоне эпох, на партии размером 64 в качестве компромисса между точностью и производительностью. Оптимизатор Adam используется для первой половины обучения при скорости обучения 0,0002, и завершается с линейно убывающей скоростью обучения до достижения нулевой потери.

Above left, segmentation applied to the wound area. Above center, image of the actual wound; above right, a synthetic wound of a type that can be generalized in a dataset, based on the original source. Below, the original wound, and, right, a synthesis of the wound generated by WG2GAN.

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

eKare Wound models

Исследователи отметили проблему, которая часто возникает на начальном этапе обучения, когда набор данных довольно разнообразен и веса случайны – модель требует много времени (75 эпох), чтобы “успокоиться”:

Где данные варьируются, и модели GAN, и модели кодировщика/декодировщика испытывают трудности в получении обобщения на ранних этапах, как мы видим, подтверждено на графике обучения WG²GAN, который отслеживает хронологию обучения с момента начала до нулевой потери.

Необходимо проявлять осторожность, чтобы обеспечить, что процесс обучения не фиксируется на особенностях или характеристиках какого-либо одного итерации или эпохи, а продолжает обобщать до полезной средней потери без получения результатов, которые чрезмерно абстрагируют исходный материал. В случае с WG²GAN это рискует создать неограниченные, совершенно “фиктивные” раны, сконкатенированные среди слишком широкого диапазона не связанных типов ранений, а не производить точный диапазон вариаций внутри определенного типа раны.

Контроль объема в наборе данных машинного обучения

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

Хотя меньшие наборы данных могут достичь высокореалистичных изображений за меньшее время, диапазон изображений и типов ранений, генерируемых ими, будет обязательно более ограничен. Существует хрупкий баланс в режимах обучения GAN и кодировщика/декодировщика между объемом и разнообразием входных данных, верностью произведенных изображений и реализмом произведенных изображений – вопросы объема и веса, которые, безусловно, не ограничиваются медицинской синтезой изображений.

Несбалансированные классы в медицинских наборах данных

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

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