Распознавание изображений нейросетью и работа с эмбеддингами¶

Этот пример позволит вам сделать первые шаги в проекте, где нужна классификация изображений или где планируется задействовать нейросеть ResNet для извлечения признаков из изображения или матрицы.

Описание задачи¶

Многие из сложнейших задач компьютерного зрения можно решить при помощи "фундаментальных" (foundational) нейросетей, и одним из ярких достижений прошлого десятилетия была сеть ResNet, на вычислениях которой, или даже на их промежуточных результатах, основаны многие современные методы.

Для работы с этой нейросетью лучше всего воспользоваться библиотекой, предоставляющей доступ к солидному парку предобученных нейросетей, Metalhead. Нейросети в этой библиотеке создаются в формате Flux, команды этой библиотеки позволят вам легко работать с топологией моделей. Напоследок нам понадобится библиотека DataAugmentation для упрощения действий над изображениями, хотя все те же операции (масштабирование и растеризация) можно выполнить при помощи обычных матричных сложений и умножений.

Pkg.add( ["Flux", "Metalhead", "DataAugmentation"] )

Распознаем одно изображение¶

Выполним небольшой расчет на предобученной нейросети ResNet, который позволит нам определить класс объекта, представленного на изображении.

using Flux, Metalhead
model = ResNet(18; pretrain = true);

Нейросеть ResNet создавалась для работы с RGB изображениями. Поэтому кроме загрузки самой картинки мы переводим ее в трехканальный формат, избавляясь от альфа-канала, при помощи простой команды RGB, применяемой к каждому пикселю матрицы изображения.

using Images
img = RGB.(load( "dog.png" ))

using DataAugmentation
DATA_MEAN = (0.485, 0.456, 0.406)
DATA_STD = (0.229, 0.224, 0.225)
augmentations = CenterResizeCrop((224, 224)) |> ImageToTensor() |> Normalize(DATA_MEAN, DATA_STD)

data = apply(augmentations, Image(img)) |> itemdata
labels = readlines( "imagenet_classes_ru.txt" )
img_labels = Flux.onecold( model(Flux.unsqueeze(data, 4)), labels )
print( img_labels )

["эскимосская собака"]

Заметим, что мы получили вектор. Впоследствие мы будем брать первый элемент этого вектора.

Итак, если изображение не слишком необычное с точки зрения датасета ImageNet, то нейросеть, которая была обучена на этом датасете, распознает его с предсказуемой достоверностью. Мы задействовали самую небольшую из "хрестоматийных" реализаций ResNet размером в 18 слоёв.

Узнать размер других доступных в библиотеке нейросетей ResNet можно просто вызвать справку по этой команде:

# ?ResNet

Распознаем несколько изображений¶

Чтобы выполнить эту операцию над несколькими изображениями можно просто упаковать вызов нейросети в цикл for. Небольшая нейросеть в 18 слоёв даже на CPU будет выполняться считанные секунды.

using Flux, Metalhead, Images, DataAugmentation

Классифицируем ряд изображений:

# Цепочка предобработки включает стандартных для ImageNet нормализацию
DATA_MEAN = (0.485, 0.456, 0.406)
DATA_STD = (0.229, 0.224, 0.225)
augmentations = CenterResizeCrop((224, 224)) |> ImageToTensor() |> Normalize(DATA_MEAN, DATA_STD)

# Загрузим изображения средствами библиотеки Images
imgs = load.("imgs/" .* ["dog1.png", "dog2.png", "cat1.png", "cat2.png", "chair1.png", "chair2.png", "rocket.png", "airplane.png"])
imgs = [ RGB.(img) for img in imgs ]

# Загрузим предобученную модель из Metalhead
model = ResNet(18; pretrain = true);

# Преобразуем каждое изображение
imgs_data = [ apply(augmentations, Image(img)) |> itemdata for img in imgs ]

# Загрузим перечень классов из текстового файла
labels = readlines("imagenet_classes_ru.txt")

# Отправляем каждое изображение в нейросеть и получаем индекс наиболее вероятной метки
img_labels = [Flux.onecold(model(Flux.unsqueeze(data, 4)), labels)[1] for data in imgs_data]

# Вывод изображений и подписей к ним
println(img_labels)
imgs

["золотистый ретривер", "чихуахуа", "сиамская кошка", "африканский хамелеон", "раскладной стул", "диван-кровать", "снаряд", "авиалайнер"]

Очевидно, в этом списке не все метки подобраны идеально. Как всегда, сказываются ограничения тренировочной выборки и последнее изображение выходит за границы надежного распознавания.

Промежуточное латентное представление¶

Для более сложной работы с нейросетями ResNet можно воспользоваться промежуточным представлением объектов, которые эти нейросети создают в процессе вычисления результата. Один из вариантов состоит в том, чтобы исключить финальный слой и работать только с backbone частью модели (базовая модель?), которую можно использовать в качестве генератора признаков для какой-либо другой задачи.

# Первая часть ResNet, называемая backbone моделью, принимает изображения любого размера
backbone_model = Metalhead.backbone( model );
img_embeddings = Flux.activations(backbone_model, Flux.unsqueeze(data[1], 4))
length( img_embeddings )

5

Почему слоёв не 18? Помимо того что мы работаем с укороченным вариантом ResNet, некоторые слои этой матрицы объединены. Например, Conv и MeanPool образуют один слой типа Chain. Функция activations анализирует только объект Chain верхнего уровня.

Объект какого размера нейросеть вырабатывает на последнем слое backbone части?

size(img_embeddings[end])

(7, 7, 512, 1)

Можно убедиться что перед нами 512 картинок размера 7 на 7. Последняя размерность - это номер батча, но мы отправляли изображения в нейросеть одно за другим, поэтому размерность по 4-му измерению у нас равна 1.

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

# Предположим, img_embeddings[end] имеет размер (7, 7, 512, 1)
img_activations = dropdims(img_embeddings[end], dims=4)  # Удаляем последнюю размерность -> (7, 7, 512)

# Параметры тайлов, Количество тайлов в строке и столбце
tile_height, tile_width, n_tiles_per_row, n_tiles_per_col = 7, 7, 32, 16

# Создаем пустую матрицу для tilemap
tilemap = zeros(Float32, tile_height * n_tiles_per_col, tile_width * n_tiles_per_row)

# Заполняем tilemap
for k in 1:min(512, n_tiles_per_row * n_tiles_per_col)  # На случай, если активаций меньше чем 512
    i = div(k - 1, n_tiles_per_row) + 1  # Номер строки в tilemap
    j = mod(k - 1, n_tiles_per_row) + 1  # Номер столбца в tilemap
    
    # Вычисляем координаты в tilemap
    row_range = (1:tile_height) .+ (i-1)*tile_height
    col_range = (1:tile_width) .+ (j-1)*tile_width
    
    # Вставляем активации
    tilemap[row_range, col_range] = img_activations[:, :, k]
end

# Визуализация
heatmap(tilemap, 
       aspect_ratio=:equal, color=:viridis, cbar=false,
       title="Карта активаций предпоследнего слоя (7×7×512 → 224×112)")

Мы видим множество маленьких "изображений". На самом деле нейросеть преобразовала трехслойное изображение в 512-слойное, попутно сократив исходную картинку до размера 7 на 7.

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

# Добавление слоя mean pooling - Размер пулинга должен соответствовать выходу feature extractor
pool_model = Chain(
    backbone_model,
    AdaptiveMeanPool((1, 1)),
    Flux.flatten
)

# Загрузка и предобработка
imgs = load.("imgs/" .* ["dog1.png", "dog2.png", "cat1.png", "cat2.png", "chair1.png", "chair2.png", "rocket.png", "airplane.png"])
imgs = [ RGB.(img) for img in imgs ]
imgs_data = [ apply(augmentations, Image(img)) |> itemdata for img in imgs ]

# Классификация моделью `model`` (чтобы получить метки классов)
img_labels = [Flux.onecold(model(Flux.unsqueeze(data, 4)), labels)[1] for data in imgs_data]

# Еще одна классификация моделью `pool_model` (просто чтобы получить эмбеддинги)
mlp_activations_512 = [Flux.activations( pool_model, Flux.unsqueeze(data, 4) )[end] for data in imgs_data];
plot(
    ([heatmap(reshape( activation, :, 32 ), cbar=false, title=title) for (activation,title) in zip(mlp_activations_512, img_labels)])...,
    layout=( 1,: ), size=(900,200)
)

Финальное латентное представление¶

Латентным представлением (эмбеддингом) может считаться информация на любом срезе нейросети.

Последний слой ResNet-18 преобразует 512 признаков, которые мы видели на прошлом графике, в 1000 признаков, и наиболее интенсивно выраженный из них характеризует класс объекта.

Распределение меток между другими "пикселями" тоже может что-нибудь нам сказать об уверенности нейросети в своём прогнозе или о классе объекта, но поскольку нейросеть не обучалась классифицировать высокоуровневые концепции (собака/кошка), а училась проставлять конкретные метки, внутриклассовая близость между всеми породами собак и всеми породами кошек может быть невысокой.

DATA_MEAN = (0.485, 0.456, 0.406)
DATA_STD = (0.229, 0.224, 0.225)
augmentations = CenterResizeCrop((224, 224)) |> ImageToTensor() |> Normalize(DATA_MEAN, DATA_STD)

imgs = load.("imgs/" .* ["dog1.png", "dog2.png", "cat1.png", "cat2.png", "chair1.png", "chair2.png", "rocket.png", "airplane.png"])
imgs = [ RGB.(img) for img in imgs ]

model = ResNet(18; pretrain = true);

imgs_data = [ apply(augmentations, Image(img)) |> itemdata for img in imgs ]
labels = readlines("imagenet_classes_ru.txt")
img_labels = [Flux.onecold(model(Flux.unsqueeze(data, 4)), labels)[1] for data in imgs_data]

mlp_activations = [Flux.activations( model.layers, Flux.unsqueeze(data, 4) )[end] for data in imgs_data];
plot(
    ([heatmap(reshape( activation, 20, : ), cbar=false, title=title) for (activation,title) in zip(mlp_activations, img_labels)])...,
    layout=( 1,: ), size=(900,300), titlefont = Plots.font(9)
)

На CPU выполнение предыдущей ячейки занимает 10-15 секунд. Использование батчей могло бы ускорить классификацию.

Напоследок построим корреляционную диаграмму. Если бы мы классифицировали спектры при помощи предобученной нейросети, этот результат мог бы представлять интерес уже до какого либо машинного обучения.

plot( heatmap( cor( hcat(mlp_activations_512...) ), title="512 признаков", c=:viridis, yflip=true, xticks=(1:8, img_labels), yticks=(1:8, img_labels)),
      heatmap( cor( hcat(mlp_activations...) ), title="1000 признаков", c=:viridis, yflip=true, xticks=(1:8, img_labels), yticks=(1:8, img_labels) ),
      size=(1000,500))

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

Вплоть до этого слоя нейросеть тренировалась извлекать признаки, релевантные для классификации на заявленные 1000 классов, так что распределение финального эмбеддинга может быть информативно. Но поскольку конечная цель обучения состояла в максимизации отклика релевантного логита, одного из 1000 (а не в изучении распределения классов), между семантически близкими объектами (разные породы собак) может не быть ожидаемой близости распределений, и мы можем ничего не увидеть на корреляционной диаграмме.

Заключение¶

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