Франкенштейн в медицине: как я скрестил ViT и ruGPT-3, чтобы научить ИИ читать рентген на русском

Сейчас из каждого утюга рассказывают про мультимодальные нейросети: GPT-4o смотрит через камеру, Gemini анализирует видео. В медицине тоже есть крутые открытые ИИ-модели для анализа снимков (например, на базе датасетов MIMIC-CXR), но у них всех есть один фатальный недостаток для нашего рынка — они говорят исключительно на английском.

Мне стало интересно: а можно ли на бесплатных мощностях, буквально "на коленке", собрать русскоязычного ИИ-рентгенолога? Спойлер: можно. В этой статье расскажу, как я скрестил Vision Transformer от Google с ruGPT-3 от Сбера, как боролся с датасетами на Kaggle и что из этого вышло. В конце — ссылки на GitHub и рабочее демо.

Архитектура: как пришить глаза к мозгу

Чтобы нейросеть могла посмотреть на снимок и написать текст, нужна архитектура Vision-Language Model. Обучать такого монстра с нуля у меня не было ни ресурсов, ни желания. Поэтому я пошел по пути Hugging Face VisionEncoderDecoderModel.

Идея проста как кирпич:

1. Энкодер (Глаза): Берем предобученный google/vit-base-patch16-224-in21k. Он отлично дробит картинку на патчи и извлекает визуальные фичи (понимает, где ребра, а где легкие).
2. Декодер (Язык): Берем ai-forever/rugpt3small_based_on_gpt2. У нее нет глаз, она умеет только генерировать текст.

Чтобы их сшить, пришлось немного "взломать" конфиг ruGPT-3, принудительно сказав ей: «Теперь ты декодер, и у тебя есть слои кросс-внимания» (is_decoder=True, add_cross_attention=True). Hugging Face заботливо создал новые пустые веса между двумя моделями. Именно эти связи мне и предстояло обучить.

Data Engineering: боль, страдания и Kaggle

Найти 7-10 тысяч рентгеновских снимков с подробными заключениями на русском языке в открытом доступе — задача нереальная.

Поэтому я взял открытый американский датасет Indiana University Chest X-Ray (IU X-Ray). Там есть картинки и тексты от американских врачей.

Прямо на Kaggle я поднял пайплайн машинного перевода на базе Helsinki-NLP/opus-mt-en-ru. Закинул тексты в GPU батчами по 32 штуки и за 10 минут перевел более 7000 медицинских заключений на вполне сносный русский медицинский язык.

Но тут платформа подкинула сюрприз: Kaggle прячет часть файлов в виртуальной файловой системе (снимков 7000, а стандартный скрипт видел только 4). Пришлось писать суровый маппинг с глубоким сканированием (os.walk), отрезать расширения и жестко связывать ID в CSV с реальными путями на диске.

Обучение: выжимаем все соки из бесплатных T4

Обучение проходило на Kaggle (2x NVIDIA T4). Чтобы модель не умерла от нехватки памяти (OOM), а сессия не отвалилась по тайм-ауту, пришлось шаманить:

• Включил Mixed Precision (fp16) — ускорило обучение в 2 раза.
• Настроил Gradient Accumulation — размер батча на видеокарту был всего 4, но виртуально мы накапливали до 16.
• Столкнулся с тем, что Seq2SeqTrainer крашится при попытке сохранить промежуточный чекпоинт мультимодального "франкенштейна". Решение? Выключить промежуточные сохранения (save_strategy="epoch") и молиться, чтобы Kaggle не завис. (Кстати, спасает JS-скрипт в консоли браузера, делающий клик раз в 60 секунд).

На 15 эпох ушло около 2.5 часов.

Что получилось в итоге? (Потрогать руками)

Получилась нейросеть, которая реально понимает, что изображено на рентгене, и сыпет терминами вроде «кальцифицированная гранулема» или «легочная васкулярность». Да, иногда она "галлюцинирует" (датасет в 7к снимков — это капля в море для ML), но базовые вещи вроде чистых легких или пневмоторакса сечет неплохо.

Я завернул модель в Gradio и выложил на Hugging Face Spaces. Можно зайти с телефона или ПК, загрузить любой снимок рентгена из гугла и посмотреть, что она выдаст.

👉 Потыкать лайв-демо тут: Hugging Face Space

👉 Весь код, пайплайны и веса тут: GitHub Репозиторий

Буду рад, если кому-то этот код сэкономит время при создании своих мультимодальных сеток. Залетайте в репу, ставьте звездочки, форкайте. Если есть идеи, как улучшить датасет (может, прогнать переводы через LLM для чистки медицинского сленга) — пишите в комменты!

(Дисклеймер: модель обучена в исследовательских целях за вечер. Не суйте ей свои снимки вместо похода к реальному врачу)