Как я собрал модель для определения стадии диабетической ретинопатии в бесплатном Colab за 3 минуты

Этот проект опубликован как открытый ноутбук и разбирает полный цикл: от загрузки датасета до интерпретации результатов через Grad-CAM (карта внимания, которая показывает, на какой участок снимка среагировала сеть). Главный урок оказался не в архитектуре, а в данных: синтетические клинические показатели содержали подсказку с правильным ответом, и модель научилась читать именно её, а не снимок. Ниже весь путь по шагам, включая способы обнаружить такую утечку.

Что понадобится?

• Google-аккаунт и доступ к Google Colab (бесплатного тарифа с GPU Tesla T4 достаточно)
• Библиотеки Python: torch, timm, medmnist, albumentations, opencv-python, numpy, matplotlib
• Датасет RetinaMNIST из пакета MedMNIST (скачивается одной командой, 1 600 снимков глазного дна, пять стадий по шкале ICDR)
• Примерно 30 минут на первый проход и около 3 минут непосредственно на обучение

Пошаговая инструкция

1. Загрузите датасет.

RetinaMNIST содержит 1 600 снимков из соревнования DeepDRiD, разделённых на обучающую (1 080), валидационную (120) и тестовую (400) выборки. Снимки приведены к размеру 224 на 224 пикселей. Пять классов, от DR 0 (нет признаков ретинопатии) до DR 4 (пролиферативная стадия), распределены неравномерно: в тесте 174 снимка DR 0 и всего 20 снимков DR 4.

from medmnist import RetinaMNIST

train_data = RetinaMNIST(split="train", download=True, size=224)
val_data   = RetinaMNIST(split="val",   download=True, size=224)
test_data  = RetinaMNIST(split="test",  download=True, size=224)

2. Примените препроцессинг по методу Бена Грэма и CLAHE.

Фундус-камеры (приборы для съёмки глазного дна) дают неравномерную засветку. Чтобы убрать плавный градиент яркости и подчеркнуть сосуды, геморрагии и экссудаты, из снимка вычитается его размытая копия по формуле:

I_enh = 4 * I - 4 * G_sigma(I) + 128

Этот приём использовал победитель соревнования Kaggle Diabetic Retinopathy Detection. После него применяется CLAHE (адаптивное выравнивание гистограммы, контрастность подтягивается локально, а не по всему кадру разом) по каналу яркости в цветовом пространстве LAB:

def apply_ben_graham_preprocessing(image_np, size=224, sigmaX=10):
    img = cv2.resize(image_np, (size, size))
    mask = np.zeros((size, size), dtype=np.uint8)
    cv2.circle(mask, (size // 2, size // 2), int(size * 0.46), 255, -1)
    blurred = cv2.GaussianBlur(img, (0, 0), sigmaX)
    enhanced = cv2.addWeighted(img, 4, blurred, -4, 128)
    lab = cv2.cvtColor(enhanced, cv2.COLOR_RGB2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    l = clahe.apply(l)
    enhanced = cv2.cvtColor(cv2.merge((l, a, b)), cv2.COLOR_LAB2RGB)
    return cv2.bitwise_and(enhanced, enhanced, mask=mask)

На стадиях DR 2 и выше после фильтра на сером фоне отчётливо видны тёмные кровоизлияния и светлые очаги экссудации.

3. Постройте двухпоточную сеть.

Визуальный энкодер (кодировщик изображений) efficientnet_b2 из библиотеки timm извлекает из снимка вектор из 1 408 признаков. Табличная ветка принимает четыре нормализованных показателя пациента: HbA1c (гликированный гемоглобин), стаж диабета, систолическое давление, возраст, и переводит их в 32-мерный вектор. Оба вектора склеиваются (1 440 признаков) и поступают в классификатор:

class RetinaTwoStreamNet(nn.Module):
    def __init__(self, backbone_name="efficientnet_b2",
                 num_classes=5, tabular_dim=4):
        super().__init__()
        self.backbone = timm.create_model(
            backbone_name, pretrained=True,
            num_classes=0, global_pool="")
        vision_dim = self.backbone.num_features  # 1408
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.tabular_net = nn.Sequential(
            nn.Linear(tabular_dim, 32),
            nn.BatchNorm1d(32), nn.SiLU(), nn.Dropout(0.2),
            nn.Linear(32, 32), nn.SiLU())
        self.classifier = nn.Sequential(
            nn.Linear(vision_dim + 32, 256),
            nn.BatchNorm1d(256), nn.SiLU(), nn.Dropout(0.35),
            nn.Linear(256, num_classes))

    def forward(self, img_tensor, tab_tensor):
        v_emb = self.pool(self.backbone(img_tensor)).flatten(1)
        t_emb = self.tabular_net(tab_tensor)
        return self.classifier(torch.cat([v_emb, t_emb], dim=1))

4. Настройте обучение.

Ключевые параметры: веса классов обратно пропорциональны частоте в обучающей выборке, label_smoothing=0.05, оптимизатор AdamW с раздельными шагами для бэкбона (3e-4) и головы с табличной веткой (1.2e-3), косинусное расписание, смешанная точность AMP FP16. Аугментации через Albumentations: повороты до 180 градусов, флипы, случайные изменения яркости и контраста. 12 эпох на Tesla T4 заняли 184.9 секунды.

5. Оцените результат и проверьте утечку.

Двухпоточная модель на тесте показала QWK 0.834, Macro AUC 0.859, accuracy 0.633. Выглядит неплохо, но именно здесь начинается самое ценное: проверка, не содержат ли входные данные подсказку.

Три контрольных эксперимента:

• Обнуление табличного входа. QWK падает с 0.83 до 0.45. Четыре числа давали модели больше информации, чем 1 408 визуальных признаков. Для реальной клиники это абсурд: HbA1c и стаж связаны с риском, но текущую стадию определяет снимок.
• Регрессия без снимков (только четыре числа, ни одного пикселя). Она почти догоняет двухпоточную сеть по QWK и обгоняет по AUC.
• Модель только по снимкам (без табличных данных). QWK 0.713, это и есть честная метрика визуального распознавания.

6. Поймите природу утечки.

Клинических данных в RetinaMNIST нет. Автор сгенерировал их по формуле, в которую входила истинная стадия болезни:

hba1c    = np.clip(5.9  + stage * 0.85 + rng.normal(0, 1.15), 4.5, 14.5)
duration = np.clip(2.5  + stage * 3.1  + rng.normal(0, 3.8),  0.0, 35.0)
sys_bp   = np.clip(122  + stage * 6.5  + rng.normal(0, 12.0), 90.0, 200.0)
age      = np.clip(46   + stage * 3.5  + rng.normal(0, 9.5),  20.0, 88.0)

Каждый показатель линейно зависит от стадии с небольшим шумом. Сеть научилась читать эту подсказку вместо того, чтобы анализировать снимок. В реальном проекте утечка (data leakage) выглядит так же: признак, который кажется клиническим, на самом деле закодировал ответ.

Что делать с этим прямо сейчас?

Автору Дзена и копирайтеру. Этот кейс, готовый сюжет для контента о подводных камнях нейросетей в медицине. Тезис «модель показала 0.83» без разбора утечки данных, типичная ошибка поверхностных обзоров. Покажите читателям, почему красивая цифра может быть ложной.

Маркетологу. Если ваш подрядчик демонстрирует метрики модели без контрольных экспериментов, задайте вопрос: «Что будет, если убрать один из входов?». Этот простой тест экономит бюджет.

Разработчику и исследователю в РФ. Пайплайн полностью воспроизводим в бесплатном Colab. RetinaMNIST доступен без VPN. Для работы с реальными клиническими данными в России потребуются датасеты с деперсонализированными медицинскими записями, но сама архитектура и метод проверки на утечку переносятся на любой набор данных.

Главный вывод: когда модель показывает подозрительно хороший результат, проблема почти всегда в данных, а не в архитектуре. Три контрольных эксперимента из этого проекта занимают минуты, но спасают от публикации ложных метрик. Если вы собираете собственную модель для определения стадии диабетической ретинопатии в бесплатном Colab или для любой другой задачи, начните с вопроса: «Не подсказал ли я модели ответ?».