Обзор архитектур image-to-image translation

Обзор архитектур image-to-image translation

https://t.me/data_analysis_ml

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

В скоуп наших проектов попадают как системы, управляющие (или частично управляющие) технологическим процессом (например, флотация или плавка), так и системы промышленного машинного зрения, которые по сути представляют из себя одну из разновидностей датчиков.

В этой статье я расскажу про основные архитектуры генеративных сетей для задачи перевода изображения из одного домена в другой (image-to-image translation). В конце расскажу, для чего именно мы применяем синтетические данные и приведу примеры изображений, которых нам удалось достичь. Но перед погружением в данную тему рекомендую ознакомиться с тем, что такое свёрточная сеть, U-Net и генеративная сеть. Если же Вы готовы, поехали.

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

Слева - конвейер с дроблёным файнштейном, справа - семантическая маска

Существующие архитектуры можно разделить на 2 типа: supervised и unsupervised. Для supervised подхода входная маска и целевое изображение должны соответствовать друг другу, как на картинке выше. Для unsupervised такого условия не требуется. В обоих случаях мы хотим сгенерировать изображение таким образом, чтобы сохранилось содержимое маски (расположение объектов), при этом стиль сгенерированного изображения совпадал с целевым.

Supervised image-to-image translation

pix2pix

[Статья, реализация]


Обучение pix2pix

Про архитектуру

Генератор представляет собой U-Net подобную архитектуру, на вход которой подаётся бинарная маска. Дискриминатор - это классификатор типа PatchGAN, на вход которому подаются бинарная маска и целевое изображение.

Как и в задаче сегментации, в архитектуре генератора нужны skip connections из энкодер блоков в декодер блоки. Но есть и некоторые отличия от архитектур сегментации.

Например, при обучении ГАНов для уменьшения размерности рекомендуется заменить операции max pooling на свёртки со страйдом. А в качестве финальной активации используется гиперболический тангенс, т.к. сгенерированные изображение мы затем подаём на вход дискриминатору, т.е. мат. ожидание выходного изображения должно быть нулевым, а дисперсия единичной.

PatchGAN - это свёрточный классификатор, который классифицирует каждый патч (реальный или фейковый). Реализовано это следующим образом.

NLayerDiscriminator


В дискриминаторе ReLU активации заменены на LeakyRELU, удалены все fully connected слои, а на выходе нет никакой активации, т.е. выходные значения ничем не ограничены. Теоретически это может привести к gradient explosion, однако в данной конфигурации это работает стабильно.
Про обучение и борьбу с затуханием градиентов
Модели обучаются в генеративно-состязательном режиме. Т.е. цель дискриминатора - научиться отличать оригинальные изображения от сгенерированных, цель генератора - научиться генерировать такие изображения, что дискриминатор не сможет их отличить от реальных.

\min\limits_{G} \max\limits_{D} V(D, G)=\mathbb{E}_{x\sim p_{data}(x)}[logD(x)] + \mathbb{E}_{z\sim p_{z}(z)}[log(1-D(G(z)))].
Такой способ обучения называется minimax game. В этой игре дискриминатор минимизирует кросс-энтропию между целевым классом и предсказанным распределением а генератор наоборот максимизирует кросс-энтропию. Это приводит к тому, что, когда дискриминатор слишком уверенно отклоняет сгенерированные изображения, градиенты генератора затухают.

Один из способов для борьбы с затуханием градиентов генератора - non-saturating game, который был найден эвристическим путём:
Вместо минимизации

\mathbb{E}_{z\sim p_{z}(z)}[log(1-D(G(z)))]
для генератора максимизируем

\mathbb{E}_{z}logD(G(z))
Авторы pix2pix и используют non-saturating game.

Также они добавляют к лосс-функции (L2 loss даёт более размытые результаты):

\mathcal{L}_{L1}(G)=   \mathbb{E}_{x, y,z}[\|y - G(x, z)\|_{1}].
И L1 и L2 loss между реальным и сгенерированным изображения хорошо улавливают низкочастотные детали, но не высокочастотные детали.

Для увеличения фокуса на высокочастотных деталях, картинку делят на патчи, как это и сделано в PatchGAN.
Архитектура pix2pix является базовой для supervised image-to-image translation. Pix2pixHD и SPADE, которые мы рассмотрим дальше, улучшают её идеи.
pix2pixHD
[Статья, реализация]
Данная архитектура, как следует из названия, предназначена для генерации изображений более высокого разрешения, чем её предшественница.
Про архитектуру
Можно выделить следующие модификации:

Multi-scale discriminators
Исходная картинка уменьшается в 2 и 4 раза - получаем 3 изображения разного масштаба. На каждый масштаб создаётся свой дискриминатор PatchGAN, которые имеют одинаковую структуру. Реализация.

MultiscaleDiscriminator

Progressive generator
Сперва сеть G1 обучается на низком разрешении, потом добавляются слои G2 и модель дообучается на высоком разрешении.
Про обучение
Feature matching + perceptual loss
Во время обучения из дискриминатора извлекаются промежуточные признаки с разных слоёв и считается L1 loss между признаками реального изображения и сгенерированного. Реализация.
Perceptual loss похож на Feature matching, но вместо дискриминатора используется предобученная на Imagenet модель, например, VGG16 (во время обучения GANов она не обучается). Реализация.
Обе лосс-функции стабилизируют обучение, т.к. генератор должен производить правдоподобные статистики на разных масштабах изображения.
Lsgan
Вместо cross-entropy в minimax game считаем L2 loss [lsgan].

Про instance-wise генерацию. Boundary map
Ко входной семантической маски добавляется boundary map, которая содержит instance-wise информацию. Это делается для того, чтобы разные объекты одного класса не сливались в один объект, как на рисунке выше.
Про instance-wise генерацию. Feature encoder network    
Feature Encoder Network представляет собой encoder-decoder архитектуру (без skip-connections), выход которой подаётся на слой instance-wise average pooling. Новый слой усредняет признаки для каждого инстанса по boundary map, полученная маска добавляется к семантической маске.
Таким образом, во время инференса для одной и той же семантической маски мы можем сгенерировать разные изображения, подавая на вход сети E, случайные изображения из нашего набора (instance-wise average pooling усредняет признаки, в соответствии с семантической маской).
 Энкодер E обучается совместно с генератором и дискриминатором.
SPADE
[Статья, реализация]

Про архитектуру
Multi-sclae discriminators
Как и в pix2pixHD используются Multi-scale discriminators.

Feature matching + perceptual loss
Также как и в pix2pixHD.

SPADE generator
Генератор сильно отличается от двух предыдущих архитектур. На вход генератору подаётся либо случайный вектор (если используем конфигурацию с VAE) либо маска уменьшенной размерности.
В конфигурации с автоэнкодером вводится дополнительная сеть - вариационный энкодер, на выходе которого и получаем случайный вектор (входной вектор для генератора). Энкодер и генератор, таким образом, формируют VAE.
Архитектура состоит из несколько SPADE Res блоков, внутри которых применяется новый слой нормализации - SPADE.
Как видно на рисунке маска прогоняется через свёртки для получения параметров гамма и бетта (В BatchNorm2d либо InstanceNorm2d параметр affine=False) [реализация].
Когда мы применяем InstanceNorm2d на однотонную маску, практически вся информация теряется (см. рисунок ниже). Блок SPADE это исключает.
Spectral Norm
Авторы также добавляют спектральную нормализацию во все архитектуры - генератор, дискриминатор и энкодер. Например, для свёрток в SPADEResnetBlock и слоёв нормализации Batch/InstanceNorm.

\textbf W_{SN}=\frac{\textbf W}{\sigma(\textbf W)}, \sigma(\textbf W) = \max\limits_{\textbf h: \textbf h\ne0} \frac {\|\textbf W \textbf h\|_{2}}{\|\textbf h\|_{2}}.
Спектральная норма (сигма) по сути измеряет, как сильно матрица W может растянуть любой ненулевой вектор h, и равна максимальному сингулярному значению матрицы W [proof].

Спектральная нормализация стабилизирует процесс обучения ГАНов за счёт того, что делает константу Липшица L = 1 для всех слоев.
Про обучение
Обучается SPADE таким же образом, как и pix2pixHD, только вместо LSGAN используется Hinge loss.

Unsupervised image-to-image translation
Cycle-GAN
[Cтатья, реализация]

Про архитектуру
Архитектуры генератора и дискриминатора в Cycle-GAN такие же, как и в pix2pix. Меняется только постановка задачи и процедура обучения.

Про обучение
Мы хотим перевести изображения из домена X в домен Y, сохранив при этом содержимое изображения. И наоборот, из домена Y в X. Т.е. решаем задачу style transfer.
Создаём два генератора G : X -> Y, F: Y -> X и два дискриминатора Dx, Dy. Каждую пару генератора и дискриминатора обучаем, например, через minimax game. И добавляем ещё две лосс-функции cycle consistency losses, необходимые для сохранения содержимого изображения.

\mathcal{L}_{cyc}(G, F)=   \mathbb{E}_{x\sim p_{data}(x)}[\|F(G(x)) - x\|_{1}] + \mathbb{E}_{y\sim p_{data}(y)}[\|G(F(y)) - y\|_{1}].
UGATIT
[Статья, реализация]

Про архитектуру
UGATIT архитектура похожа на Cycle GAN. Она состоит из двух генераторов (по одному на каждый домен) и 4 дискриминаторов (по 2 на каждый домен).
Дискриминаторы одного домена - это два PatchGAN разной глубины (таким образом, реализуется multi-scale дискриминатор).
Генератор представляет собой Resnet-based архитектуру (без skip connections).
Внутри Resnet блоки с AdaLIN (adaptive layer instance normalization)

AdaLIN(a,\gamma, \beta) = \gamma \cdot  (\rho\cdot \hat a_{I} + (1 - \rho)\cdot \hat a_{L}) + \beta,
\hat a_{I}=\frac{a - \mu_{I}}{\sqrt {\sigma^2_{I} + \epsilon}}, \hat a_{L}=\frac{a - \mu_{L}}{\sqrt {\sigma^2_{L} + \epsilon}},
\rho \leftarrow clip_{[0,1]} (\rho-\tau\ \Delta\rho),
где μI , μL and σI , σL - это channel-wise и layer-wise мат. ожидание и стандартное отклонений активаций соответственно, γ и β - параметры, генерируемые fully connected слоем, τ - это learning rate и ∆ρ - градиент, определённый оптимайзером. Значения ρ ограничены отрезком [0, 1].

Про обучение
Внутри дискриминатора есть дополнительная классификационная голова. Ей на вход подаются 2D карты признаков, по ним с помощью avg_pooling и fully connected слоёв считаются веса и умножаются на соответствующие карты признаков.
Во время инференса по этим картам строятся class activation map (CAM; attention feature map; heatmap)
Полученный 2d тензор и есть выход дополнительной классификационной головы. Получаем + 4 дополнительных лосса для обучения генератора:
G_ad_cam_loss_GA = self.MSE_loss(fake_GA_cam_logit, torch.ones_like(fake_GA_cam_logit).to(self.device))
G_ad_cam_loss_LA = self.MSE_loss(fake_LA_cam_logit, torch.ones_like(fake_LA_cam_logit).to(self.device))

G_ad_cam_loss_GB = self.MSE_loss(fake_GB_cam_logit, torch.ones_like(fake_GB_cam_logit).to(self.device))
G_ad_cam_loss_LB = self.MSE_loss(fake_LB_cam_logit, torch.ones_like(fake_LB_cam_logit).to(self.device))

И дискриминатора:
D_ad_cam_loss_GA = self.MSE_loss(real_GA_cam_logit, torch.ones_like(real_GA_cam_logit).to(self.device)) + self.MSE_loss(fake_GA_cam_logit, torch.zeros_like(fake_GA_cam_logit).to(self.device))
D_ad_cam_loss_LA = self.MSE_loss(real_LA_cam_logit, torch.ones_like(real_LA_cam_logit).to(self.device)) + self.MSE_loss(fake_LA_cam_logit, torch.zeros_like(fake_LA_cam_logit).to(self.device))

D_ad_cam_loss_GB = self.MSE_loss(real_GB_cam_logit, torch.ones_like(real_GB_cam_logit).to(self.device)) + self.MSE_loss(fake_GB_cam_logit, torch.zeros_like(fake_GB_cam_logit).to(self.device))
D_ad_cam_loss_LB = self.MSE_loss(real_LB_cam_logit, torch.ones_like(real_LB_cam_logit).to(self.device)) + self.MSE_loss(fake_LB_cam_logit, torch.zeros_like(fake_LB_cam_logit).to(self.device))

Основные лоссы (minimax game):
G_ad_loss_GA = self.MSE_loss(fake_GA_logit, torch.ones_like(fake_GA_logit).to(self.device))
G_ad_loss_LA = self.MSE_loss(fake_LA_logit, torch.ones_like(fake_LA_logit).to(self.device))

G_ad_loss_GB = self.MSE_loss(fake_GB_logit, torch.ones_like(fake_GB_logit).to(self.device))
G_ad_loss_LB = self.MSE_loss(fake_LB_logit, torch.ones_like(fake_LB_logit).to(self.device))

D_ad_loss_GA = self.MSE_loss(real_GA_logit, torch.ones_like(real_GA_logit).to(self.device)) + self.MSE_loss(fake_GA_logit, torch.zeros_like(fake_GA_logit).to(self.device))
D_ad_loss_LA = self.MSE_loss(real_LA_logit, torch.ones_like(real_LA_logit).to(self.device)) + self.MSE_loss(fake_LA_logit, torch.zeros_like(fake_LA_logit).to(self.device))

D_ad_loss_GB = self.MSE_loss(real_GB_logit, torch.ones_like(real_GB_logit).to(self.device)) + self.MSE_loss(fake_GB_logit, torch.zeros_like(fake_GB_logit).to(self.device))
D_ad_loss_LB = self.MSE_loss(real_LB_logit, torch.ones_like(real_LB_logit).to(self.device)) + self.MSE_loss(fake_LB_logit, torch.zeros_like(fake_LB_logit).to(self.device))

Также как и в cycle gan добавляются cycle consistency losses
L^{s\rightarrow t}_{cycle}=\mathbb{E}_{x \sim X_{s}}[|x-G_{t \rightarrow s}(G_{s \rightarrow t}(x))|_{1}].
И identity loss (для сохранения содержимого):

L^{s\rightarrow t}_{identity}=\mathbb{E}_{x \sim X_{t}}[|x-(G_{s \rightarrow t}(x))|_{1}].
fake_A2A, fake_A2A_cam_logit, _ = self.genB2A(real_A)
fake_B2B, fake_B2B_cam_logit, _ = self.genA2B(real_B)

G_identity_loss_A = self.L1_loss(fake_A2A, real_A)
G_identity_loss_B = self.L1_loss(fake_B2B, real_B)

В генераторы, как и в дискриминаторы, добавляются дополнительные классификационные головы, которая предсказывают принадлежит ли изображение домену А или домену B, в зависимости от генератора.
G_cam_loss_A = self.BCE_loss(fake_B2A_cam_logit, torch.ones_like(fake_B2A_cam_logit).to(self.device)) + self.BCE_loss(fake_A2A_cam_logit, torch.zeros_like(fake_A2A_cam_logit).to(self.device))
G_cam_loss_B = self.BCE_loss(fake_A2B_cam_logit, torch.ones_like(fake_A2B_cam_logit).to(self.device)) + self.BCE_loss(fake_B2B_cam_logit, torch.zeros_like(fake_B2B_cam_logit).to(self.device))

Финальный лосс:
\min\limits_{G} \max\limits_{D} \lambda_{1}L_{lsgan}+\lambda_{2}L_{cycle} + \lambda_{3}L_{identity} + \lambda_{4}L_{cam}
MUNIT
[Статья, реализация]

Про архитектуру
Данная архитектура значительно отличается от двух предыдущих. В отличие от них, MUNIT кодирует исходное изображение на содержимое и стиль при помощи Content and Style Encoder.
Style Encoder состоит из strided свёрток, global average pooling и fully connected слоя. Instance Normalization в Style Encoder не используется, т.к. он удаляет оригинальное мат. ожидание и дисперсию признаков, которые содержат информацию о стиле.
Content encoder состоит из тех же свёрток и residual blocks.
Декодер по содержимому и стилю восстанавливает исходное изображение. Он состоит из свёрток и интерполяции (nearest-neighbor interpolation) для повышения размерности. Внутри residual blocks используется Adaptive Instance Normalization.

AdaIN(z,\gamma,\beta) = \gamma (\frac{z-\mu(z)}{\sigma(z)}) + \beta,
где z - это активация предыдущего свёрточного слоя, μ и σ - это channel-wise мат. ожидание и стандартное отклонение активаций, γ и β - это параметры, генерируемые fully connected слоем.

Дискриминатор - multi-scale дискриминатор с LSGAN лоссом (в minimax game)
Про обучение
Также как и в pix2pixHD используется perceptual loss. Перед вычислением дистанции к эмбэддингам применяется Instance Normalization, чтобы удалить информацию о стиле.

MUNIT обучается одновременно и как автоэнкодер и как GAN. Для преобразования изображения из домена А в B и наоборот они пропускаются через style и content энкодер. У каждого домена свои энкодеры. Затем, чтобы получить изображение из домена B(с2 - содержимое; s2 - стиль), но с содержимым изображения из домена А (с1 - содержимое; s1 - стиль), мы заменяем s1 на s2 и пару (с1, s2) пропускаем через декодер для домена B. Аналогично для преобразования из B в А.
Для обучения как обычно используется adversarial loss (minimax game), reconstruction loss (cycle consistency losses) для изображений, как и в двух предыдущих моделях, и reconstruction loss для содержимого и стиля.

\min\limits_{E_{1}, E_{2}, G_{1}, G_{2}} \max\limits_{D_{1}, D_{2}} \mathcal{L}(E_{1}, E_{2}, G_{1}, G_{2}, D_{1}, D_{2}) =  \mathcal{L}^{x_{1}}_{GAN} + \mathcal{L}^{x_{2}}_{GAN} + \lambda_{x}( \mathcal{L}^{x_{1}}_{recon} +  \mathcal{L}^{x_{2}}_{recon}) + \\ \lambda_{c}( \mathcal{L}^{c_{1}}_{recon} +  \mathcal{L}^{c_{2}}_{recon}) + \lambda_{x}( \mathcal{L}^{x_{1}}_{recon} +  \mathcal{L}^{x_{2}}_{recon})
 \mathcal{L}^{x_{1}}_{recon} =\mathbb{E}_{x_{1} \sim p(x_{1})}[\| G_{1}(E^{c}_{1}(x_{1}), E^{s}_{1}(x_{1})) - x_{1} \|_{1}]
 \mathcal{L}^{c_{1}}_{recon} =\mathbb{E}_{c_{1} \sim p(c_{1}), s_{2} \sim q(s_{2}))}[\| E^c_{2}(G_{2}(c1, s2)) - c_{1} \|_{1}]
 \mathcal{L}^{s_{2}}_{recon} =\mathbb{E}_{c_{1} \sim p(c_{1}), s_{2} \sim q(s_{2}))}[\| E^s_{2}(G_{2}(c1, s2)) - s_{2} \|_{1}]
 \mathcal{L}^{x_{2}}_{GAN} =\mathbb{E}_{c_{1} \sim p(c_{1}), s_{2} \sim q(s_{2})}[log(1-D_{2}(G_{2}(c_{1}, s_{2})))] + \mathbb{E}_{x_{2} \sim p(x_{2})}[logD_{2}(x_{2})].
Что получилось у нас
И так мы разобрали основные, на мой взгляд, архитектуры в image-to-image translation, а теперь мне бы хотелось рассказать, зачем мы это применяли у себя и как это нам помогло.

Для чего применяем
Нередко у нас появляется задача - обучить Unet/LinkNet модель, чтобы точно находить объекты и определять их размеры. Разметка таких объектов крайне кропотливое занятие, а новые домены могут сильно отличаться от тех, с которыми мы уже умеем работать.

Как вариант, можно применять Domain Adaptation методы (DA), но иногда домены отличаются довольно сильно, например, руда и пузыри. В таком случае DA не поможет. Поэтому мы захотели научиться генерировать синтетические изображения по бинарной маске, имея всего лишь небольшое количество размеченных изображений (~ 100).
Генерация масок
До сих пор мы ничего не говорили про генерацию самих масок, хотя это тоже довольно важный этап для получения разнообразных и качественных изображений.
Мы попробовали обучить ProgressiveGAN, но получилось плохо. В масках было мало разнообразия.
Наилучшие результаты получились с масками, которые генерировали сегментационные модели, и с рандомно расставленными по пустому изображению масками объектов (из небольшого набора).
Генерация изображений по маскам
Мы попробовали все приведённые в статье архитектуры и наилучшие результаты показали SPADE (среди supervised архитектур) и UGATIT (среди unsupervised).

Таким образом, мы сгенерировали набор синтетических данных (порядка 10 000 изображений), преодобучили UNet модель на них, а затем дообучили на реальных данных. Это дало ~1-2% прироста в точности по метрикам precision, recall (для объектов, не для пикселей) по сравнению с моделью, обученной только на реальных данных.
Заключение
Генеративные сети можно применять для предобучения моделей либо в качестве аугментации. Сгенерированные, таким образом, пары - маска и картинка - лучше соответствуют друг другу, чем пара - реальное изображения и маска, сгенерированная сегментационной моделью. Но совсем без разметки не обойтись.
Качество unsupervised методов сильно ниже supervised. Последние же хоть и генерируют реалистичные изображения, но недостаточно разнообразные.
https://t.me/machinelearning_interview - вопросы с собеседований ml
источнник



          

          

            

              
              
            

            

              
              

            

          

          

            

              
            

          

          
        

      

          

      
Report content on this page

    

    

      

    

      

        

          
Report Page

          

            
            
            
            
            
            
          

          

            
          

          

            Please submit your DMCA takedown request to dmca@telegram.org