在TensorFlow+Keras環境下使用RoI池化一步步實現注意力機制

機器之心發表於2019-05-10

專案地址:https://gist.github.com/Jsevillamol/0daac5a6001843942f91f2a3daea27a7

理解 RoI 池化

RoI 池化的概念由 Ross Girshick 在論文「Fast R-CNN」中提出,RoI 池化是其目標識別工作流程中的一部分。

在 RoI 池化的一般用例中,我們會有一個類似影像的目標,以及用邊界框指定的多個感興趣區域。我們要從每個 RoI 中生成一個嵌入。

例如,在 R-CNN 的設定下,我們有一個影像和一個為影像中可能感興趣的部分生成邊界框的候選機制。接下來,我們要為每一個候選的影像塊生成嵌入:

在TensorFlow+Keras環境下使用RoI池化一步步實現注意力機制

簡單地裁剪每個候選區域是行不通的,因為我們想要將最終得到的嵌入疊加在一起,而候選區域的形狀不一定相同!

因此,我們需要想出一種方法對每個影像塊進行變換,以生成預定義形狀的嵌入。我們要怎麼實現這一點?

計算機視覺領域,使用池化操作是縮小影像形狀的一種標準做法。

最常見的池化操作是「最大池化」。此時,我們將輸入影像劃分成形狀相同的區域(通常是不重疊的),然後通過取每個區域的最大值來得到輸出。

在TensorFlow+Keras環境下使用RoI池化一步步實現注意力機制

最大池化操作將每個區域劃分為若干大小相同的池化區域

這並不能直接解決我們所面臨的問題——形狀不同的影像塊將被劃分成數量不一的形狀相同的區域,產生不同形狀的輸出。

但這為我們提供了一個思路。如果我們把每個感興趣的區域劃分成相同數量的形狀不同的區域,並取每個區域的最大值呢?

在TensorFlow+Keras環境下使用RoI池化一步步實現注意力機制

RoI 的池化操作將所有區域劃分為相同數量的池化區域網格。

這正是 RoI 池化層所做的工作。

使用注意力機制的好處

ROI 池化實現了所謂的「注意力機制」,它讓我們的模型可以專注於輸入的特定特徵。

在目標識別任務的環境下,我們可以將任務工作流程劃分為兩部分(候選區域和區域分類),同時保留端到端的可微架構。

在TensorFlow+Keras環境下使用RoI池化一步步實現注意力機制

展示 RoI 池化層的 Fast R-CNN 架構。圖源:Ross Girshick 的論文《Fast R-CNN》。

RoI 池化是一種泛化能力很強的的注意力工具,可以用於其他任務,比如對影像中預選區域的一次性上下文感知分類。也就是說,它允許我們對同一張影像的不同區域進行一次標記處理。

更一般而言,注意力機制受到了神經科學和視覺刺激研究的啟發(詳見 Desimone 和 Duncan 1995 年發表的論文「Neural Mechanism of Selective Visual Attention」。

如今,對注意力機制的應用已經超越了計算機視覺的範疇,它在序列處理任務中也廣受歡迎。我覺得讀者可以研究一下 Open AI 的注意力模型示例:《Better Language Models and their Implications》,該模型被成功地用於處理各種自然語言理解任務。

RoI 層的型籤

在我們深入研究實現細節之前,我們可以先思考一下 RoI 層的型籤(type signature)。

RoI 層有兩個輸入張量

  • 一批影像。為了同時處理這些,所有影像必須具備相同的形狀。最終得到的 Tensor 形狀為(batch_size,img_width,img_height,n_channels)。

  • 一批候選的感興趣區域(RoIs)。如果我們想將它們堆疊在一個張量中,每張影像中候選區域的數量必須是固定的。由於每個邊界框需要通過 4 個座標來指定,該張量的形狀為(batch_size,n_rois,4)。

RoI 層的輸出應該為:

  • 為每章影像生成的嵌入列表,它編碼了每個 RoI 指定的區域。對應的形狀為(batch_size,n_rois,pooled_width,pooled_height,n_channels)

Keras 程式碼

Keras 讓我們可以通過繼承基本層類來實現自定義層。

「tf.keras」官方文件建議我們為自定義層實現「__init__」、「build」以及「call」方法。然而,由於「build」函式的目的是為層新增權重,而我們要實現的 RoI 層並沒有權重,所以我們並不需要覆蓋該方法。我們還將實現方便的「compute_output_shape」方法。

我們將分別對每個部分進行編碼,然後在最後將它們整合起來。

def __init__(self, pooled_height, pooled_width, **kwargs):
    self.pooled_height = pooled_height
    self.pooled_width = pooled_width
    super(ROIPoolingLayer, self).__init__(**kwargs)

類的 constructor 很容易理解。我們需要指定待生成嵌入的目標高度和寬度。在 constructor 的最後一行中,我們呼叫 parent constructor 來初始化其餘的類屬性。

def compute_output_shape(self, input_shape):
    """ Returns the shape of the ROI Layer output
    """
    feature_map_shape, rois_shape = input_shape
    assert feature_map_shape[0] == rois_shape[0]
    batch_size = feature_map_shape[0]
    n_rois = rois_shape[1]
    n_channels = feature_map_shape[3]
    return (batch_size, n_rois, self.pooled_height, 
            self.pooled_width, n_channels)

「compute_output_shape」是一個很好用的效用函式,它將告訴我們對於特定的輸入來說,RoI 層的輸出是怎樣的。

接下來,我們需要實現「call」方法。「call」函式是 RoI 池化層的邏輯所在。該函式應該將持有 RoI 池化層輸入的兩個張量作為輸入,並輸出帶有嵌入的張量

在實現這個方法之前,我們需要實現一個更簡單的函式,它將把單張影像和單個 RoI 作為輸入,並返回相應的嵌入。

接下來,讓我們一步一步實現它。

@staticmethod (http://twitter.com/staticmethod)
def _pool_roi(feature_map, roi, pooled_height, pooled_width):
  """ Applies ROI Pooling to a single image and a single ROI
  """
# Compute the region of interest        
  feature_map_height = int(feature_map.shape[0])
  feature_map_width  = int(feature_map.shape[1])

  h_start = tf.cast(feature_map_height * roi[0], 'int32')
  w_start = tf.cast(feature_map_width  * roi[1], 'int32')
  h_end   = tf.cast(feature_map_height * roi[2], 'int32')
  w_end   = tf.cast(feature_map_width  * roi[3], 'int32')

  region = feature_map[h_start:h_end, w_start:w_end, :]
...

函式的前六行在計算影像中 RoI 的起始位置和終止位置。

我們規定每個 RoI 的座標應該由 0 到 1 之間的相對數字來指定。具體而言,每個 RoI 由包含四個相對座標(x_min,y_min,x_max,y_max)的四維張量來指定。

我們也可以用絕對座標來指定該 RoI,但是通常而言這樣做效果會較差。因為輸入影像在被傳遞給 RoI 池化層之前會經過一些會改變影像形狀的卷積層,這迫使我們跟蹤影像的形狀是如何改變的,從而對 RoI 邊界框進行適當的放縮。

第七行使用 TensorFlow 提供的超強張量切片語法將圖片直接裁剪到 RoI 上。

...
# Divide the region into non overlapping areas
region_height = h_end - h_start
region_width  = w_end - w_start
h_step = tf.cast(region_height / pooled_height, 'int32')
w_step = tf.cast(region_width  / pooled_width , 'int32')

areas = [[(
           i*h_step, 
           j*w_step, 
           (i+1)*h_step if i+1 < pooled_height else region_height, 
           (j+1)*w_step if j+1 < pooled_width else region_width
          ) 
          for j in range(pooled_width)] 
         for i in range(pooled_height)]
...

在接下來的四行中,我們計算了待池化的 RoI 中每個區域的形狀。

接著,我們建立了一個二維張量陣列,其中每個元件都是一個元組,表示我們將從中取最大值的每個區域的起始座標和終止座標。

生成區域座標網格的程式碼看起來過於複雜,但是請注意,如果我們只是將 RoI 劃分成形狀為(region_height / pooled_height,region_width / pooled_width)的區域,那麼 RoI 的一些畫素就不會落在任何區域內。

我們通過擴充套件右邊和底部的大部分割槽域將預設情況下不會落在任何區域的剩餘畫素囊括進來,從而解決這個問題。這是通過在程式碼中宣告每個邊界框的最大座標來實現的。

該部分最終得到的是一個二維邊界框列表。

...
# Take the maximum of each area and stack the result
def pool_area(x): 
  return tf.math.reduce_max(region[x[0]:x[2],x[1]:x[3],:], axis=[0,1])

pooled_features = tf.stack([[pool_area(x) for x in row] for row in areas])
return pooled_features上面幾行程式碼十分巧妙。我們定義了一個輔助函式「pool_area」,其輸入為我們剛剛建立的元組指定的邊界框,輸出為該區域中每個通道的最大值。

我們使用列表解析式對每個已宣告的區域進行「pool_area」對映

由此,我們得到了一個形狀為(pooled_height,pooled_width,n_channels)的張量,它儲存了單張影像某個 RoI 的池化結果。

接下來,我們將對單張影像的多個 RoI 進行池化。使用一個輔助函式可以很直接地實現這個操作。我們還將使用「tf.map_fn」生成形狀為(n_rois,pooled_height,pooled_width,n_channels)的張量

@staticmethod (http://twitter.com/staticmethod)
def _pool_rois(feature_map, rois, pooled_height, pooled_width):
  """ Applies ROI pooling for a single image and varios ROIs
  """
  def curried_pool_roi(roi): 
    return ROIPoolingLayer._pool_roi(feature_map, roi, 
                                     pooled_height, pooled_width)

  pooled_areas = tf.map_fn(curried_pool_roi, rois, dtype=tf.float32)
  return pooled_areas

最後,我們需要實現 batch 級迭代。如果我們將一個張量系列(如我們的輸入 x)傳遞給「tf.map_fn」,它將會把該輸入壓縮為我們需要的形狀。

def call(self, x):
  """ Maps the input tensor of the ROI layer to its output
  """
  def curried_pool_rois(x): 
    return ROIPoolingLayer._pool_rois(x[0], x[1], 
                                      self.pooled_height, 
                                      self.pooled_width)

  pooled_areas = tf.map_fn(curried_pool_rois, x, dtype=tf.float32)
  return pooled_areas

請注意,每當「tf.map_fn」的預期輸出與輸入的資料型別不匹配時,我們都必須指定「tf.map_fn」的「dtype」引數。一般來說,我們最好儘可能頻繁地指定該引數,從而通過 Tensorflow 計算圖來明確型別是如何變化的。

下面,讓我們將上述內容整合起來:

import tensorflow as tf
from tensorflow.keras.layers import Layer

class ROIPoolingLayer(Layer):
    """ Implements Region Of Interest Max Pooling 
        for channel-first images and relative bounding box coordinates

        # Constructor parameters
            pooled_height, pooled_width (int) -- 
              specify height and width of layer outputs

        Shape of inputs
            [(batch_size, pooled_height, pooled_width, n_channels),
             (batch_size, num_rois, 4)]

        Shape of output
            (batch_size, num_rois, pooled_height, pooled_width, n_channels)

    """
    def __init__(self, pooled_height, pooled_width, **kwargs):
        self.pooled_height = pooled_height
        self.pooled_width = pooled_width

        super(ROIPoolingLayer, self).__init__(**kwargs)

    def compute_output_shape(self, input_shape):
        """ Returns the shape of the ROI Layer output
        """
        feature_map_shape, rois_shape = input_shape
        assert feature_map_shape[0] == rois_shape[0]
        batch_size = feature_map_shape[0]
        n_rois = rois_shape[1]
        n_channels = feature_map_shape[3]
        return (batch_size, n_rois, self.pooled_height, 
                self.pooled_width, n_channels)

    def call(self, x):
        """ Maps the input tensor of the ROI layer to its output

            # Parameters
                x[0] -- Convolutional feature map tensor,
                        shape (batch_size, pooled_height, pooled_width, n_channels)
                x[1] -- Tensor of region of interests from candidate bounding boxes,
                        shape (batch_size, num_rois, 4)
                        Each region of interest is defined by four relative 
                        coordinates (x_min, y_min, x_max, y_max) between 0 and 1
            # Output
                pooled_areas -- Tensor with the pooled region of interest, shape
                    (batch_size, num_rois, pooled_height, pooled_width, n_channels)
        """
        def curried_pool_rois(x): 
          return ROIPoolingLayer._pool_rois(x[0], x[1], 
                                            self.pooled_height, 
                                            self.pooled_width)

        pooled_areas = tf.map_fn(curried_pool_rois, x, dtype=tf.float32)

        return pooled_areas

    @staticmethod
    def _pool_rois(feature_map, rois, pooled_height, pooled_width):
        """ Applies ROI pooling for a single image and varios ROIs
        """
        def curried_pool_roi(roi): 
          return ROIPoolingLayer._pool_roi(feature_map, roi, 
                                           pooled_height, pooled_width)

        pooled_areas = tf.map_fn(curried_pool_roi, rois, dtype=tf.float32)
        return pooled_areas

    @staticmethod
    def _pool_roi(feature_map, roi, pooled_height, pooled_width):
        """ Applies ROI pooling to a single image and a single region of interest
        """

        # Compute the region of interest        
        feature_map_height = int(feature_map.shape[0])
        feature_map_width  = int(feature_map.shape[1])

        h_start = tf.cast(feature_map_height * roi[0], 'int32')
        w_start = tf.cast(feature_map_width  * roi[1], 'int32')
        h_end   = tf.cast(feature_map_height * roi[2], 'int32')
        w_end   = tf.cast(feature_map_width  * roi[3], 'int32')

        region = feature_map[h_start:h_end, w_start:w_end, :]

        # Divide the region into non overlapping areas
        region_height = h_end - h_start
        region_width  = w_end - w_start
        h_step = tf.cast( region_height / pooled_height, 'int32')
        w_step = tf.cast( region_width  / pooled_width , 'int32')

        areas = [[(
                    i*h_step, 
                    j*w_step, 
                    (i+1)*h_step if i+1 < pooled_height else region_height, 
                    (j+1)*w_step if j+1 < pooled_width else region_width
                   ) 
                   for j in range(pooled_width)] 
                  for i in range(pooled_height)]

        # take the maximum of each area and stack the result
        def pool_area(x): 
          return tf.math.reduce_max(region[x[0]:x[2], x[1]:x[3], :], axis=[0,1])

        pooled_features = tf.stack([[pool_area(x) for x in row] for row in areas])
        return pooled_features

接下來,測試一下我們的實現方案!我們將使用一個高度和寬度為 200x100 的單通道影像,使用 7x3 的池化影像塊提取出 2 個 RoI。影像最多可以有 4 個標籤來對區域進行分類。示例特徵圖上的每個畫素都為 1,只有處於(height-1,width-3)位置的一個畫素值為 50。

import numpy as np
# Define parameters
batch_size = 1
img_height = 200
img_width = 100
n_channels = 1
n_rois = 2
pooled_height = 3
pooled_width = 7
# Create feature map input
feature_maps_shape = (batch_size, img_height, img_width, n_channels)
feature_maps_tf = tf.placeholder(tf.float32, shape=feature_maps_shape)
feature_maps_np = np.ones(feature_maps_tf.shape, dtype='float32')
feature_maps_np[0, img_height-1, img_width-3, 0] = 50
print(f"feature_maps_np.shape = {feature_maps_np.shape}")
# Create batch size
roiss_tf = tf.placeholder(tf.float32, shape=(batch_size, n_rois, 4))
roiss_np = np.asarray([[[0.5,0.2,0.7,0.4], [0.0,0.0,1.0,1.0]]], dtype='float32')
print(f"roiss_np.shape = {roiss_np.shape}")
# Create layer
roi_layer = ROIPoolingLayer(pooled_height, pooled_width)
pooled_features = roi_layer([feature_maps_tf, roiss_tf])
print(f"output shape of layer call = {pooled_features.shape}")
# Run tensorflow session
with tf.Session() as session:
    result = session.run(pooled_features, 
                         feed_dict={feature_maps_tf:feature_maps_np,  
                                    roiss_tf:roiss_np})

print(f"result.shape = {result.shape}")
print(f"first  roi embedding=\n{result[0,0,:,:,0]}")ooled_features.shape}")

上面的幾行為該層定義了一個測試輸入,構建了相應的張量並執行了一個 TensorFlow 會話,這樣我們就可以檢查它的輸出。

執行該指令碼將得到如下輸出:

feature_maps_np.shape = (1, 200, 100, 1)
roiss_np.shape = (1, 2, 4)
output shape of layer call = (1, 2, 3, 7, 1)
result.shape = (1, 2, 3, 7, 1)
first  roi embedding=
[[1. 1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1. 1.]
 [1. 1. 1. 1. 1. 1. 1.]]
second roi embedding=
[[ 1.  1.  1.  1.  1.  1.  1.]
 [ 1.  1.  1.  1.  1.  1.  1.]
 [ 1.  1.  1.  1.  1.  1. 50.]]

如上所示,輸出張量的形狀與我們期望的結果相符。除了我們指定為 50 的畫素,最終得到的嵌入都是 1。

我們的實現似乎是有效的。

結語

在本文中,我們瞭解了 RoI 池化層的功能,以及如何使用它來實現注意力機制。此外,我們還學習瞭如何擴充套件 Keras 來實現不帶權重的自定義層,並給出了上述 RoI 池化層的實現。

希望本文對你有所幫助。

原文連結:https://medium.com/xplore-ai/implementing-attention-in-tensorflow-keras-using-roi-pooling-992508b6592b

相關文章