Киберпанк, который мы заслужили: как я на коленке собрал ИИ, читающий рентген лёгких на русском

Пока все развлекаются тем, что заставляют ChatGPT писать за них дипломы, а Midjourney — рисовать эльфиек, я решил сделать что-то реально полезное.
В мире больших корпораций уже давно существуют медицинские нейросети. 
Пока все развлекаются тем, что заставляют ChatGPT писать за них дипломы, а Midjourney — рисовать эльфиек, я решил сделать что-то реально полезное. В мире больших корпораций уже давно существуют медицинские нейросети. 

Но если вы зайдете на GitHub или Hugging Face, то обнаружите обидную вещь: почти все open-source модели для медицины понимают только английский язык.

Мне стало интересно: а можно ли дома, не имея бюджета OpenAI, собрать своего собственного русскоязычного кибер-врача? Спойлер: да. И в конце поста вы сможете сами подкинуть ему рентгеновский снимок и посмотреть, что он вам выдаст.

Анатомия Франкенштейна: как сшить нейросети

Чтобы ИИ мог посмотреть на картинку и выдать осмысленный текст, нужна мультимодальная архитектура (Vision-Language). Писать её с нуля — самоубийство, поэтому я пошел по пути генной инженерии и взял две готовые модели:

  1. Глаза (Энкодер): Vision Transformer (ViT) от Google. Эта штука отлично рубит картинку на куски и понимает визуальные паттерны. Она знает, где на фото ребра, а где подозрительное темное пятно.
  2. Мозг и Рот (Декодер): ruGPT-3 от Сбера. У нее нет глаз, зато она отлично генерирует русский текст.

Я загнал их в архитектуру VisionEncoderDecoder от Hugging Face. Грубо говоря, я вскрыл конфиг текстовой модели, принудительно активировал ей слои cross-attention (чтобы она могла принимать сигналы от "глаз") и сшил их вместе.

Фарм данных и прожарка GPU
Машину собрали, теперь её нужно учить. Где взять 7 тысяч рентгенов грудной клетки с подробными заключениями врачей на русском языке? Нигде.
Фарм данных и прожарка GPU Машину собрали, теперь её нужно учить. Где взять 7 тысяч рентгенов грудной клетки с подробными заключениями врачей на русском языке? Нигде.

Поэтому я взял открытую базу клиники Индианского университета (США), где лежат реальные снимки и диагнозы на английском. Прямо в облаке поднял нейросеть-переводчик (Helsinki-NLP) и заставил её перевести весь этот архив на наш великий и могучий.

Само обучение я запустил на Kaggle. Если кто не знает, это платформа, где бесплатно дают попользоваться серверами с видеокартами. Мне достались 2x NVIDIA T4.Чтобы не поймать Out of Memory (нехватку видеопамяти), пришлось включать смешанную точность (fp16) и градиентное накопление. Карточки пыхтели на 100% несколько часов подряд (15 эпох обучения), пока модель не начала наконец-то связывать белые пятна на снимках со словами «пневмоторакс» или «кальцинаты».

Потрогать своими руками 🧪

Я выкатил готовую модель на сервера Hugging Face Spaces. Работает прямо в браузере, никаких VPN, регистраций и смс.

Квест для читателей:

  1. Идем в гугл-картинки, ищем «рентген легких норма» или любую патологию.
  2. Качаем картинку, закидываем в мою нейросеть.
  3. Ждем секунд 10-15 (сервера бесплатные, процессоры слабенькие, проявите терпение).
  4. Делимся в комментах, какой диагноз она вам сгенерировала! (Можете попробовать загрузить туда кота или скриншот из игры, посмотрим, как она сломается).

(ВАЖНЫЙ ДИСКЛЕЙМЕР: Я программист. Модель — это пет-проект, обученный за вечер. Она может галлюцинировать и находить болезни там, где их нет. Пожалуйста, не используйте её для реального самолечения! Идите к нормальным врачам).

Для тех, кому нужен код (Open Source)

Весь код, пайплайны подготовки данных и скрипты обучения я выложил на GitHub. Если вы изучаете Machine Learning, там есть полезные куски кода: от обхода таймаутов Kaggle до кастомных PyTorch датасетов.

👉 Ссылка на репозиторий: https://github.com/livadies-collab/Multimodal-XRay-Analyzer-RU

Если пост зайдет, в следующий раз попробую обучить модель искать переломы на рентгенах конечностей. Задавайте вопросы в комментах, буду рад ответить!

5
1
28 комментариев