Створення синтетичних наборів даних про рани за допомогою генеративних суперницьких мереж

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

Система, яка називається WG²AN, є спільним проєктом Коледжу інженерії та технологій Баттена та компанії з охорони здоров’я eKare, яка спеціалізується на застосуванні методологій машинного навчання для вимірювання та ідентифікації ран.

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

Для того, щоб розмістити доступний латентний простір у архітектурі тренування моделі, зображення були масштабовані до 512×512 та витягнуті з їхнього фону. Для вивчення впливу розміру набору даних були проведені тестові запуски на партіях з 100, 250, 500, 1000, 2000 та 4000 зображень.

Джерело: https://ietresearch.onlinelibrary.wiley.com/doi/pdfdirect/10.1049/tje2.12033

Зображення вище показує збільшення деталізації та гранулярності залежно від розміру тренувального набору даних та кількості епох, проведених на кожному проході.

Архітектура WG2GAN. Джерело: https://ietresearch.onlinelibrary.wiley.com/doi/pdfdirect/10.1049/tje2.12033

WG²AN працює на PyTorch на відносно лекому споживчому устаткуванні, з 8 ГБ відеопам’яті на GPU GTX 1080. Тренування тривало від 4 до 58 годин над діапазоном розмірів наборів даних від 100 до 4000 зображень, та над діапазоном епох, на розмірі партії 64 як компроміс між точністю та продуктивністю. Оптимізатор Адама використовується для першої половини тренування при швидкості навчання 0,0002, та завершується лінійним спадом швидкості навчання до втрати нуля.

Зверху ліворуч, сегментація, застосована до області рани. Зверху центру, зображення реальної рани; зверху праворуч, синтетична рана типу, який можна узагальнити в наборі даних, на основі оригінального джерела. Знизу, оригінальна рана, та праворуч, синтез рани, згенерований WG2GAN.

У медичних наборах даних, як і в багатьох інших галузях машинного навчання, позначення є невід’ємною瓶. У цьому випадку дослідники використовували напівавтоматичну систему позначення, яка використовує раніше дослідження компанії eKare, яке використовувало реальні моделі ран, створені з пластиліну та приблизно пофарбовані для семантичного контексту.

Моделі ран компанії eKare

Дослідники відзначили проблему, яка часто виникає на початкових етапах тренування, коли набір даних досить різноманітний, а ваги рандомізовані – модель витрачає багато часу (75 епох) на “стабілізацію”:

Коли дані різноманітні, і генеративна суперницька мережа, і моделі кодувача/декодувача мають труднощі з узагальненням на початкових етапах, як ми бачимо, підтверджуючи графік тренування WG²GAN, який відстежує графік тренування від початку до нульової втрати.

Потрібно бути обережним, щоб процес тренування не фіксувався на особливостях чи характеристиках будь-якої однієї ітерації чи епохи, а продовжував узагальнюватися до корисної середньої втрати без生产 результатів, які надмірно абстрагують джерельний матеріал. У випадку WG²GAN це ризикує створити необмежені, цілком “вигадані” рани, з’єднані серед надто широкого діапазону не пов’язаних типів ран, а не виробляти точний діапазон варіацій у рамках певного типу рани.

Контроль обсягу в наборі даних машинного навчання

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

Хоча менші набори даних можуть досягти високої реалістичності зображень за менший час, діапазон зображень і типів ран, згенерованих таким чином, буде обмеженішим. Існує делікатний баланс у тренувальних режимах генеративної суперницької мережі та кодувача/декодувача між об’ємом і різноманітністю вхідних даних, вірогідністю згенерованих зображень та реалістичністю згенерованих зображень – питання обсягу та ваг, які не обмежуються медичною синтезою зображень.

Несбалансованість класів у медичних наборах даних

Загалом, машинне навчання в галузі охорони здоров’я стикається не тільки з дефіцитом наборів даних, але й з несбалансованістю класів, коли важлива інформація про певну хворобу становить таку малу частку свого набору даних, що ризикує бути відхиленою як аутлієр-дані або бути асимільованою під час узагальнення протягом тренування.

Було запропоновано ряд методів для вирішення цієї проблеми, таких як ундер-семплінг або овер-семплінг. Однак проблема часто обходиться шляхом створення наборів даних, специфічних для певної хвороби, які повністю прив’язані до однієї медичної проблеми. Хоча такий підхід є ефективним у кожному окремому випадку, він сприяє культурі балканалізації у сфері досліджень машинного навчання в галузі охорони здоров’я, та, ймовірно, сповільнює загальний прогрес у цій галузі.