天天看點

TensorFlow2實作空間自适應歸一化(Spatial Adaptive Normalization, SPADE)空間自适應歸一化(Spatial Adaptive Normalization, SPADE)在殘差網絡中應用SPADE

空間自适應歸一化(Spatial Adaptive Normalization, SPADE)

空間自适應歸一化 (Spatial Adaptive Normalization, SPADE) 是 GauGAN 中的主要創新點,其用于語義分割圖的層歸一化,為了更好的解釋 SPADE,需要首先了解GauGAN的網絡輸入——語義分割圖。

使用獨熱編碼标記分割蒙版

考慮訓練 GauGAN 所用的 Facades 資料集。其中,分割圖在 RGB 圖像中被編碼為不同的顔色,如下圖所示。例如,一堵牆以藍色表示,柱子以紅色表示。這種表示在視覺上讓我們易于了解,但對神經網絡的學習并沒有幫助,這是因為對于 GAN 而言,顔色沒有語義。

TensorFlow2實作空間自适應歸一化(Spatial Adaptive Normalization, SPADE)空間自适應歸一化(Spatial Adaptive Normalization, SPADE)在殘差網絡中應用SPADE

顔色在顔色空間中更接近并不意味着它們在語義上也接近。例如,我們可以用淺綠色表示草,用深綠色表示飛機,即使分割圖的色相接近,它們的語義也不相關。

是以,我們應該使用類标簽,而不是使用顔色來标記像素。但是,這仍然不能解決問題,因為類别标簽是随機配置設定的數字,并且它們也沒有語義。是以,一種更好的方法是在該像素中存在對象時使用标簽為1的分割蒙版,否則使用标簽為 0 的分割蒙版。換句話說,我們将分割圖中的标簽獨熱編碼為形狀 (H, W, number of classes) 的分割蒙版。

在 JPEG 編碼中,在壓縮過程中會删除一些對視覺效果不太重要的視覺資訊。即使結果像素應該屬于同一類并且看起來是相同的顔色,它們也可能具有不同的值。是以,我們無法将 JPEG 圖像中的顔色映射到類。為了解決這個問題,我們需要使用未壓縮的圖檔格式 BMP。在圖像加載和預進行中,我們将加載檔案,并将它們從 BMP 轉換為獨熱編碼的分割蒙版。

有時,TensorFlow 的基本圖像預處理 API 無法執行一些複雜的任務,是以我們需要使用其他 Python 庫,tf.py_function 允許我們在 TensorFlow 訓練流程中運作通用 Python 函數:

def load(image_file):
    def load_data(image_file):
        jpg_file = image_file.numpy().decode('utf-8')
        bmp_file = jpg_file.replace('.jpg', '.bmp')
        png_file = jpg_file.replace('.jpg', '.png')
        image = np.array(Image.open(jpg_file))/127.5-1
        map = np.array(Image.open(png_file))/127.5-1
        labels = np.array(Image.open(bmp_file), dtype=np.uint8)
        h,w,_ = image.shape
        n_class = 12
        mask = np.zeros((h,w,n_class),dtype=np.float32)
        for i in range(n_class):
            one_hot[labels==i,i] = 1
        return map, image, mask
    [mask, image, label] = tf.py_function(load_data, [image_file], [tf.float32, tf.float32, tf.float32])      

了解了獨熱編碼的語義分割掩碼的格式後,我們将使用 TensorFlow2 實作 SPADE。

實作SPADE

執行個體歸一化已在圖像生成中非常流行,但是它往往會削弱分割蒙版的語義:假設輸入圖像僅包含一個分割标簽;例如,假設整個圖像都是天空,由于輸入具有統一的值,是以輸出在通過卷積層後也将具有統一的值。

執行個體歸一化為每個通道計算跨次元 (H, W) 的平均值。是以,該通道的均值将是相同的統一值,并且用均值減去後的歸一化激活将變為零。顯然,語義已經丢失,這是一個十分極端的示例,但是邏輯是相似的,我們可以看到分割掩碼随着其面積的增大而失去了其語義含義。

為了解決這個問題,SPADE 規範化了由分割蒙版限定的局部區域,而不是整個蒙版。下圖顯示了 SPADE 的體系結構:

TensorFlow2實作空間自适應歸一化(Spatial Adaptive Normalization, SPADE)空間自适應歸一化(Spatial Adaptive Normalization, SPADE)在殘差網絡中應用SPADE

在批歸一化中,計算跨次元 (N, H, W) 的通道的均值和标準差,對于 SPADE 來說是相同的。差別在于,每個通道的 γ 和 β 不再是标量值,而是二維向量,形狀為 (H, W)。換句話說,對于每個從語義分割圖中獲悉的激活,都有一個 γ 和 β 值。是以,歸一化被不同地應用于不同的分割區域。這兩個參數是通過使用兩個卷積層來學習的,如下圖所示:

TensorFlow2實作空間自适應歸一化(Spatial Adaptive Normalization, SPADE)空間自适應歸一化(Spatial Adaptive Normalization, SPADE)在殘差網絡中應用SPADE

SPADE 不僅應用于網絡輸入階段,同樣還應用于内部層。現在,我們可以為使用 TensorFlow2 的自定義層實作 SPADE。

将首先在 __init__ 構造函數中定義卷積層,如下所示:

class SPADE(layers.Layer):
    def __init__(self, filters, epsilon=1e-5):
        super(SPADE, self).__init__()
        self.epsilon = epsilon
        self.conv = layers.Conv2D(128, 3, padding='same', activation='relu')
        self.conv_gamma = layers.Conv2D(filters, 3, padding='same')
        self.conv_beta = layers.Conv2D(filters, 3, padding='same')      

接下來,獲得激活圖尺寸,以在以後調整大小時使用:

def build(self, input_shape):
        self.resize_shape = input_shape[1:3]      

最後,在 call() 中将層和操作連接配接在一起,如下所示:

def call(self, input_tensor, raw_mask):
        mask = tf.image.resize(raw_mask, self.resize_shape, method='nearest')
        x = self.conv(mask)
        gamma = self.conv_gamma(x)
        beta = self.conv_beta(x)
        mean, var = tf.nn.moments(input_tensor, axes=(0,1,2), keepdims=True)
        std = tf.sqrt(var+self.epsilon)
        normalized = (input_tensor - mean) / std
        output = gamma * normalized + beta
        return output      

接下來,我們将研究如何利用 SPADE。

在殘差網絡中應用SPADE

最後,将研究如何将 SPADE 插入殘差塊中:

TensorFlow2實作空間自适應歸一化(Spatial Adaptive Normalization, SPADE)空間自适應歸一化(Spatial Adaptive Normalization, SPADE)在殘差網絡中應用SPADE

SPADE 殘差塊中的基本建構塊是 SPADE-ReLU-Conv 層。每個 SPADE 都接受兩個輸入——上一層的激活和語義分割圖。

與标準殘差塊一樣,有兩個卷積 ReLU 層和一個跳躍路徑。隻要殘差塊之前和之後的通道數發生變化,就需要進行學習跳躍連接配接連接配接。發生這種情況時,前向路徑中兩個 SPADE 的輸入處的激活圖将具有不同的尺寸。但是,我們已經在 SPADE 塊中内置了調整大小的功能。以下是用于 SPADE 殘差塊建構所需圖層的代碼:

class Resblock(layers.Layer):
    def __init__(self, filters):
        super(Resblock, self).__init__()
        self.filters = filters
   
    def build(self, input_shape):
        input_filter = input_shape[-1]
        self.spade_1 = SPADE(input_filter)
        self.spade_2 = SPADE(self.filters)
        self.conv_1 = layers.Conv2D(self.filters, 3, padding='same')
        self.conv_2 = layers.Conv2D(self.filters, 3, padding='same')
        self.learned_skip = False
        if self.filters != input_filter:
            self.learned_skip = True
            self.spade_3 = SPADE(input_filter)
            self.conv_3 = layers.Conv2D(self.filters, 3, padding='same')      

最後,在 call() 中将各層連接配接起來:

def call(self, input_tensor, mask):
        x = self.spade_1(input_tensor, mask)
        x = self.conv_1(tf.nn.leaky_relu(x, 0.2))
        x = self.spade_2(x, mask)
        x = self.conv_2(tf.nn.leaky_relu(x, 0.2))
        if self.learned_skip:
            skip = self.spade_3(input_tensor, mask)
            skip = self.conv_3(tf.nn.leaky_relu(skip, 0.2))
        else:
            skip = input_tensor
        output = skip + x
        return output