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

選自 Medium

作者：Jaime Sevilla

機器之心編譯

參與：Geek AI、Chita

在本文中，作者解釋了感興趣區域池化（RoI 池化）的基本概念和一般用法，以及如何使用它來實現注意力機制。他一步步給出了在 Keras 和 TensorFlow 環境下使用 RoI 池化的實現。

項目地址：https://gist.github.com/Jsevillamol/0daac5a6001843942f91f2a3daea27a7

理解 RoI 池化

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

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

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

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

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

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

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

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

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

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

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

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

使用注意力機制的好處

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」方法。

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

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

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

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

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

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

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

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

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

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

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

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

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

我們通過擴展右邊和底部的大部分區域將默認情況下不會落在任何區域的剩餘像素囊括進來，從而解決這個問題。這是通過在代碼中聲明每個邊界框的最大坐標來實現的。

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

我們使用列表解析式對每個已聲明的區域進行「pool_area」映射。

由此，我們得到了一個形狀為（pooled_height，pooled_width，n_channels）的張量，它存儲了單張圖像某個 RoI 的池化結果。

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

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

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

下面，讓我們將上述內容整合起來：

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

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

運行該腳本將得到如下輸出：

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

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

結語

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

希望本文對你有所幫助。

本文為機器之心編譯，轉載請聯繫本公眾號獲得授權。

------------------------------------------------

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！