Методы объяснения моделей — практичный инструмент для понимания модели, оценки её точности и стабильности. Однако, часто можно столкнуться с ситуацией, когда фреймворк, в котором метод реализован, просто не "дружит" с реализацией модели. В этом туториале хочу подробно показать CAM (class activation map) для объяснения моделей зрения.

Почему CAM?

Class Activation Maps (CAM) — базовый инструмент для визуализации того, какие области изображения наиболее важны для модели при принятии решения. Он позволяет понять:

Какие признаки извлекает модель на разных слоях свертки;
Какие зоны изображения вносят вклад в прогноз конкретного класса;

Кроме того, с практической точки зрения:

Его почти всегда можно реализовать руками (если у вас есть задача классификации и сверточный слой);
Его можно использовать практически в любой задаче классификации с использованием сверточных нейронных сетей.

В процессе туториала я буду использовать YOLO NAS S из библиотеки super-gradients, фреймворк PyTorch и набор данных, включающий изображения с различными животными (коровами, овечками и другими).

В общем, прошу к чтению!

А зовут меня Сабрина. Сейчас я заканчиваю бакалавриат, пишу дипломную работу по интерпретируемости и намеренна идти дальше. Саму область explainable AI я исследую уже 3 года, и просто безумно её люблю. Рада видеть вас среди читателей!

План туториала:

Подготовка данных и загрузка модели YOLO NAS S.
Извлечение карт активации классов (CAM).
Анализ работы модели: что извлекает backbone, что делает head.
Выводы и анализ результатов.

1. Модель, данные и фреймворк

Используем предобученную YOLO NAS S, реализованную в библиотеке super-gradients. Как пишут разработчики — бибилотека предоставляет готовые к использованию state-of-the-art (SOTA) модели, а также инструменты для обучения и развертывания моделей. Преимущества библиотеки (если кратко) описаны такие:

Поддержка различных задач компьютерного зрения: классификация, детекция объектов, сегментация и другие.
Готовые модели: доступны предобученные модели, включая YOLO-NAS и YOLO-NAS-POSE.

Загрузка простая. Процесс инференса — тоже.

from super_gradients.training import models
from super_gradients.common.object_names import Models

model = models.get(Models.YOLO_NAS_S, pretrained_weights="coco")
model.eval();

Данные выложены на GitHub. Вы можете скачать их по ссылке (архив) или следовать туториалу в ноутбуке.

2. CAM: извлечение и теория

Class Activation Maps (CAM) — это базовый метод в explainable AI, который позволяет визуализировать, какие области изображения наиболее важны для модели при принятии решения. Построение таких карт активаций основано на идее сверточных нейронных сетей — если мы начнем разворачивать CNN, то обнаружим, что полносвязный слой, отвечающий за классификацию, содержит столько наборов весов, сколько классов прогнозирует модель — например, для модели, обученной на Imagenet, мы будем иметь 1000 наборов.

Из этого факта и строится идея карты активации класса — значения карт признаков, получаемых на предпоследнем сверточном слое, домножаются на веса определенного класса. А именно построение CAM проводится по формуле:

CAM_k = w^c_kA^k

где:

— карта активации модели для 𝑘-го класса на последнем сверточном слое (осуществляющем классификацию);
— веса, соответствующие 𝑘−му классу, взятые из выходного слоя классификатора;

Обычно, CAM строится по последнему слою CNN, слою перед Global Average Pooling. Но так как у нас иная архитектура, слой выбираем иным образом.

Выбор слоя

Архитектура моделей YOLO NAS (You Only Look Once) включает в себя части backbone, head (головы) и neck (шеи).

Задача backbone — извлечение признаков и особенностей изображения.
Задача neck — это передача признаков, полученных из backbone (основной части сети, которая извлекает особенности из изображения), в head.
Задача head — "принять решение" о том, какой объект присутствует на изображении (классификация), а также вычислить координаты и размеры ограничивающих рамок (локализация).

Чтобы посмотреть всю архитектуру сети, достаточно выполнить команду: model. Чтобы обратиться к какой либо части — указать её название через .:

model.backbone
model.neck
model.heads

Примечание 1:

Сами CAM по определнию строится именно на головах, так как из них можно извлечь веса, которые добавляются к прогнозу и отвечают за конкретный класс. Для backbone просто визуализируем активации (activations maps).

Примечание 2:

У каждой головы есть блок cls_convs — отвечающий за свертки классификации и cls_pred — за прогноз. С них будем извлекать свертки и веса класса соответственно.

Подытожим — при построении карт посмотрим:

какие карты сверток модель получает на выходе из различных этапов backbone, чтобы посмотреть на признаки
как признаки отвечают за прогноз в heads — головах. На них и построим карты активации классов (CAM).

Извлечение карт активаций

Для извлечения карт используется Hook— метод PyTorch, который позволяет "подключаться" к слоям сети для извлечения информации (например, выходных активаций или градиентов).

Основные возможности Hook:

Перехват выходов слоя при прямом проходе (forward hook).
Возможность изменять или анализировать градиенты при обратном проходе (backward hook).

# Простой класс Hook
class Hook():
    def __init__(self, m):
      self.hook = m.register_forward_hook(self.hook_func) #цепляемся за forward проход слоя m, переданного при инициализации

    def hook_func(self, module, input, output):
     # self.features = ((output.cpu()).data).numpy()
      self.features = output.detach().cpu().numpy()  #получаем output слоя и преобразуем в numpy массив

    def remove(self):
      self.hook.remove()

После написания, hook нужно прикрепить — то есть в него нужно просто передать слой. Для этого, в свою очередь, слой нужно из модели достать. Наиболее удобно доставать слои, используя атрибуты модели.

У модели YOLO NAS S можно получить доступ к её основным компонентам через атрибуты model.backbone, model.neck, model.heads. И далее можно детализировать доступ, указывая конкретные подкомпоненты через .или извлекая их через ._modules и обращаясь как к значениям словаря.

Пример для stage1:

stage1_layer = model.backbone.stage1.blocks._modules['bottlenecks']._modules['1'].cv2._modules['branch_1x1']

Пример для head1:

final_layer_head3 = model.heads.head3.cls_convs._modules['0']._modules['seq'][0]

Кода слой извлечен, достаточно передать его в Hook:

act_maps_stage1 = Hook(stage1_layer)

Извлечение карт

Для получения карт, достаточно пропустить изображение через модель. В контексте super-gradients важно при инференсе указать параметр fuse_model=False , иначе теряется информативность карт.

act_maps_stage1 = Hook(stage1_layer)
act_maps_stage2 = Hook(stage2_layer)
act_maps_stage3 = Hook(stage3_layer)
act_maps_stage4 = Hook(stage4_layer)

result = model.predict(img, fuse_model=False, conf=0.8)

act_maps_stage1.remove()
act_maps_stage2.remove()
act_maps_stage3.remove()
act_maps_stage4.remove()

Удаление хуков после вызова предотвращает накопление памяти и снижает вычислительную нагрузку.

3. Анализ

Для примера я взяла мишку на достаточно неприятном фоне. Пропустила его через модель и сняла карты активации со stages и принятие решений о классификации из heads. Результаты:

Оригинальный медведь (распознанный как медведь).

Интерполированные на оригинальное изображение карты активации классов.

Что здесь видно:

С увеличением номера stage модель начинает анализировать более крупные области изображения. Это закономерно, так как каждый последующий stage обрабатывает выход предыдущего, сворачивая его в более высокоуровневые признаки.
Медианные и средние значения активаций для stages различаются по силе и структуре. Однако на уровне голов они стабилизируются. Это может указывать на то, что фильтры свертки по-разному реагируют на признаки на ранних этапах, постепенно формируя обобщенные высокоуровневые представления на финальных этапах.
Локализация усиливается с увеличением номера головы. Это соответствует архитектуре модели: голова 1 обрабатывает признаки из stage1 и stage2, тогда как головы 2 и 3 используют выходы более поздних stages, что позволяет им фокусироваться на более специфичных признаках.

Выводы и анализ результатов

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

С теоретической точки зрения это способствует более глубокому пониманию особенностей работы архитектуры, а с практической — помогает оценить стабильность модели и её реакцию на различные виды возмущений. Анализ стабильности модели, в свою очередь, позволяет выявить, насколько она чувствительна к шуму, возникающему, например, из-за помех при передаче данных, а также к другим видам искажений. Это ценно для улучшения надежности моделей в реальных условиях.

Детальный пример влияния шума на прогноз модели, я привела здесь.

Благодарю за внимание к статье! Желаю вам успешных исследований и новых открытий!

Анализ обработки признаков в YOLO NAS S при помощи CAM