Deeplearning4j 實戰 （19）：基於膠囊網絡（Capsule Network）的手寫體數字識別

Eclipse Deeplearning4j GitChat課程：https://gitbook.cn/gitchat/column/5bfb6741ae0e5f436e35cd9f
Eclipse Deeplearning4j 系列博客：https://blog.csdn.net/wangongxi
Eclipse Deeplearning4j Github：https://github.com/eclipse/deeplearning4j

CapsNet(Capsule Network)最早是在2017年由G.Hinton教授及其團隊提出的。之後，幾乎以每年一更新的頻率推出優化版本。CapsNet引起人們的關注的原因主要在於Hinton教授對於在機器視覺領域大放異彩的卷積神經網絡(CNN)的運行機制和原理有不同看法，並表示CNN存在明顯不同於人類大腦感知機制的缺陷。以圖像爲例，Hinton教授認爲CNN無法感知局部特徵之間存在的相對位置關係。他曾舉例澳大利亞地理形狀以及旋180度轉後人類對其感知的不同來闡述人類的對事物的感知其實是存在座標或者說相對位置關係這一要素的。然後對於CNN來講，對這樣的變換其感知差異的能力較弱，這在某些場景下並不會產生致命影響，甚至可以允許這類情況的發生，但當需要區分的時候CNN很難做到。另外，Hinton教授曾希望放棄BP算法來訓練神經網絡的參數，很大程度上的原因是因爲他認爲人類大腦沒有如此精準的反向更新機制。同時，基於BP算法訓練的神經網絡是否work很大程度上依賴於損失函數的設計。只要損失函數設計得稍有偏差，那麼最終的模型效果將大打折扣。從某種意義上說，Hinton教授似乎希望充分借鑑人腦的感知機制，來對神經網絡系列模型重新探索一條出路，這條出路可能不會是越來越deep的方向，很有可能是一種相對淺層的結構但可以做到和deep模型一樣的效果。對於這些方向性的問題，在文末也會稍加闡述。
Deeplearning4j從1.0.0-beta4版本開始基於2017年CapsNet的論文《Dynamic Routing Between Capsules》實現了PrimaryCaps以及DigitCaps等核心的膠囊網絡結構以及內部的動態路由算法。這篇文章我們嘗試在Mnist手寫體數據集上使用CapsNet來進行分類問題的建模。下面的內容我們將先分析CapsNet模型的結構以及原理，然後基於Deeplearning4j提供的膠囊網絡結構來實現分類的相關邏輯。需要先說明一下的是，本文提到的CapsNet僅針對於2017年發表的膠囊網絡而言，概念上不包括後續的Matrix Capsule以及2019最新的層疊式的膠囊編碼器網絡模型。

CapsNet結構解析

首先需要說下Capsule結構是什麼。Capsule，所謂的膠囊結構在2017年的論文中其實是對傳統神經元的一次升級。回憶一下傳統神經元的輸入輸出，輸入的是一組標量，輸出則是經過加權求和之後（一般經過激活函數的非線性變換）的一個標量。

相對的，膠囊則是向量化版本的神經元結構，它的輸入是一組向量或者說多個向量，輸出是一個向量，內部則是經過和神經元類似的加權求和以及squashing操作（類似於傳統神經元的激活函數）。

這裏需要說明的是，膠囊結構的輸入一般是下層膠囊結構的輸出。換句話說，最初版本的膠囊結構確實很像是向量版的傳統神經元。另外一個比較重要的創新點就是可迭代動態路由算法（iterative dynamic routing process）。我們先來看下算法的流程。

從以上的描述可以看出，路由算法的思想在於通過不斷迭代更新低層膠囊每個輸入向量的耦合係數，從而使得高層的膠囊可以選擇與輸出結果相對一致的特徵，某種意義上是一種特徵選擇的過程。這種思想其實和之前卷積神經網絡中的池化思想是有點相似的。我們舉最大池化的例子來說，這個選擇也是在預測的時候決定的，而且選擇數值最大的那個作爲輸出結果，也就是一種選擇最有代表性特徵的方式。另外，這種方式和attention機制也是有點相似，耦合係數可以看作是每個輸入特徵向量的權重，這些權重因爲softmax的關係可以認爲是一種概率，那麼越接近於1，輸出的結果和該輸入越是接近。下面說下具體迭代的過程。
初始狀態下，每個耦合係數都是0。經過softmax函數後，這些係數將形成等概分佈。接着該膠囊的多個輸出向量分別乘以各自的耦合係數並將所有的這些結果求和，並經過Squashing函數後得到當前迭代輪次下的輸出。隨後，該次迭代出的輸出向量分別和輸入向量做內積，也就是相似度計算，得到一個標量數值用於更新各自的耦合係數。到此，結束本次迭代，並且耦合係數得到一次更新。如果迭代次數大於1次（論文中建議是3次），那麼以上迭代過程將重複多次。以上便是論文中提及的動態路由算法。
這裏簡單提一下非線性函數Squashing。首先看一下它的函數式。

簡單分析下可以看出，對於向量需要先單位化，然後根據該向量自身的模長，如果接近於0，則單位向量乘上一個接近0的係數，反之則該向量最大乘上縮放係數1，也就是幾乎不進行縮放。簡言之，模長越長的向量幾乎不會進行縮放，反之可能會縮小至0。

以上便是2017年膠囊網絡中的一些基本概念，下面我們看下總體的架構。

CapsNet整體架構

在2017年的第一版的CapsNet中，整體的架構其實和傳統的卷積神經網絡很相似，都有類似卷積+池化的操作。首先看下論文中的整體架構截圖。

根據截圖可以看到該網絡結構中至少存在3層網絡，第一層是傳統的卷積層，第二層是PrimaryCaps，第三層是DigitCaps。卷積層用256個9x9的大卷積核（步長爲1）來提取特徵。以Mnist數據集作爲輸出，那麼輸入是28x28x1的張量，經過第一層卷積層的特徵提取後，得到20x20x256的張量輸出並作爲PrimaryCaps的輸入。PrimaryCaps這一層的作用可以看作是膠囊網絡的卷積層。首先這一層有32個膠囊，每一個膠囊將8個9x9x256（步長等於2）的卷積核去過濾輸入的張量數據，最終得到32x8x6x6的輸出張量。下面這張圖從另外一個角度來解釋PrimaryCaps的操作。我們可以理解爲9x9（步長等於2）的256個卷積覈對輸出做卷積操作，暫時得到6x6x256的張量，然後對於這256個矩陣，8個矩陣合爲一組形成一個膠囊的輸出數據，同樣也可以得到6x6x8x32的輸出張量。最後則是DigitCaps這一層。這一層膠囊網絡包含了上面提到的動態路由算法，主要的功能類似於傳統CNN中的最大池化層。我們對照下圖來解釋下。PrimaryCaps的輸出是一個6x6x8x32的張量，我們做一個reshape的操作使其變成1152x8的張量。由於Mnist數據集最終涉及0~9共10個label。所以DigitCaps包含10個膠囊結構。1152x8的張量作爲這10個膠囊的輸入。對於每一個輸入，我們將其做一個線性變換通過8x16的權重矩陣相乘來實現。這樣對於每一個膠囊，我們得到一個1152x16的張量。對於該張量做上述的動態路由算法，最終每個膠囊可以得到一個1x16的向量輸出。由於是10個膠囊，所以是10x16的張量。每一個向量對應於一個label狀態，可以認爲提取了對於該label最合適的特徵。

從上述描述可以看出，2017年版本的CapsNet還是借鑑了很多卷積神經網絡中卷積層和最大池化層的思路。抓取特徵的方式傾向於向量的形式，而且卷積核尺寸都比較大，理論上可以抓取更宏觀的特徵。下面我們嘗試基於Deeplearning4j提供的CapsNet網絡層構建膠囊網絡分類器。

CapsNet建模Mnist數據集

從1.0.0-beta4的版本開始，Deeplearning4j開始提供CapsNet相關的網絡結構。在前文對論文的分析中，可以看到膠囊網絡基本分爲PrimaryCaps和DigitCaps兩大結構，在Deeplearning4j的實現中，org.deeplearning4j.nn.conf.layers.PrimaryCapsules是對於論文中PrimaryCaps結構的實現，而DigitCaps結構則是由org.deeplearning4j.nn.conf.layers.CapsuleLayer進行實現。下面給出pom文件的示例。

<dependency>
	<groupId>org.deeplearning4j</groupId>
	<artifactId>deeplearning4j-core</artifactId>
	<version>${dl4j.version}</version><!-- 1.0.0-beta6 -->
</dependency>	
<dependency>
	<groupId>org.nd4j</groupId>
	<artifactId>nd4j-api</artifactId>
	<version>${nd4j.version}</version><!-- 1.0.0-beta6 -->
</dependency>
<dependency>
	<groupId>org.nd4j</groupId>
	<artifactId>nd4j-native-platform</artifactId>
	<version>${nd4j.version}</version><!-- 1.0.0-beta6 -->
</dependency>

這裏我都用了最新的1.0.0-beta6的版本，當然beta4的版本開始就支持膠囊網絡結構，開發人員也可以用相對老的版本。
PrimaryCapsules的實現其實和傳統卷積層是類似的。我們來簡單分析下源碼。

@Override
public SDVariable defineLayer(SameDiff SD, SDVariable input, Map<String, SDVariable> paramTable, SDVariable mask) {
    Conv2DConfig conf = Conv2DConfig.builder()
            .kH(kernelSize[0]).kW(kernelSize[1])
            .sH(stride[0]).sW(stride[1])
            .pH(padding[0]).pW(padding[1])
            .dH(dilation[0]).dW(dilation[1])
            .isSameMode(convolutionMode == ConvolutionMode.Same)
            .build();

    SDVariable conved;

    if(hasBias){
        conved = SD.cnn.conv2d(input, paramTable.get(WEIGHT_PARAM), paramTable.get(BIAS_PARAM), conf);
    } else {
        conved = SD.cnn.conv2d(input, paramTable.get(WEIGHT_PARAM), conf);
    }

    if(useRelu){
        if(leak == 0) {
            conved = SD.nn.relu(conved, 0);
        } else {
            conved = SD.nn.leakyRelu(conved, leak);
        }
    }

    SDVariable reshaped = conved.reshape(-1, capsules, capsuleDimensions);
    return CapsuleUtils.squash(SD, reshaped, 2);
}

在之前的博客中，我們介紹過Deeplearning4j的自動微分工具SameDiff。SameDiff提供了常見的張量操作算子。那麼PrimaryCapsules就是基於SameDiff進行實現，它的主要實現邏輯是通過覆寫SameDiff的defineLayer方法來實現的。可以看到，首先定義了卷積操作的相關參數，包括卷積核的大小、步長等。接着便對input張量進行傳統的卷積操作。卷積操作結束後，可以通過relu或者leakyRelu進行非線性變換，最後reshape一下張量的並通過論文中提到的squash變換得到最後的輸出。如果按照paper中的參數設置，那麼得到的reshaped對象就是一個1152x8的張量（不考慮minibatch的情況下）。那麼到此，PrimaryCaps的實現其實就結束了。下面看下DigitCaps的實現。

@Override
public SDVariable defineLayer(SameDiff SD, SDVariable input, Map<String, SDVariable> paramTable, SDVariable mask) {

    // input: [mb, inputCapsules, inputCapsuleDimensions]

    // [mb, inputCapsules, 1, inputCapsuleDimensions, 1]
    SDVariable expanded = SD.expandDims(SD.expandDims(input, 2), 4);

    // [mb, inputCapsules, capsules  * capsuleDimensions, inputCapsuleDimensions, 1]
    SDVariable tiled = SD.tile(expanded, 1, 1, capsules * capsuleDimensions, 1, 1);

    // [1, inputCapsules, capsules * capsuleDimensions, inputCapsuleDimensions]
    SDVariable weights = paramTable.get(WEIGHT_PARAM);

    // uHat is the matrix of prediction vectors between two capsules
    // [mb, inputCapsules, capsules, capsuleDimensions, 1]
    SDVariable uHat = weights.times(tiled).sum(true, 3)
            .reshape(-1, inputCapsules, capsules, capsuleDimensions, 1);

    // b is the logits of the routing procedure
    // [mb, inputCapsules, capsules, 1, 1]
    SDVariable b = SD.zerosLike(uHat).get(SDIndex.all(), SDIndex.all(), SDIndex.all(), SDIndex.interval(0, 1), SDIndex.interval(0, 1));

    for(int i = 0 ; i < routings ; i++){

        // c is the coupling coefficient, i.e. the edge weight between the 2 capsules
        // [mb, inputCapsules, capsules, 1, 1]
        SDVariable c = CapsuleUtils.softmax(SD, b, 2, 5);

        // [mb, 1, capsules, capsuleDimensions, 1]
        SDVariable s = c.times(uHat).sum(true, 1);
        if(hasBias){
            s = s.plus(paramTable.get(BIAS_PARAM));
        }

        // v is the per capsule activations.  On the last routing iteration, this is output
        // [mb, 1, capsules, capsuleDimensions, 1]
        SDVariable v = CapsuleUtils.squash(SD, s, 3);

        if(i == routings - 1){
            return SD.squeeze(SD.squeeze(v, 1), 3);
        }

        // [mb, inputCapsules, capsules, capsuleDimensions, 1]
        SDVariable vTiled = SD.tile(v, 1, (int) inputCapsules, 1, 1, 1);

        // [mb, inputCapsules, capsules, 1, 1]
        b = b.plus(uHat.times(vTiled).sum(true, 3));
    }

    return null; // will always return in the loop
}

DigitCaps在Deeplearning4j中的通過CapsuleLayer結構來實現。上面這段邏輯就是CapsuleLayer的主邏輯。這裏的每一步張量shape的變換其實在註釋中都做了詳細的解釋，這裏就不再詳述了，開發人員可以自己研究。這段邏輯大體上分爲兩部分，第一部分是得到一個類似上文截圖的1152x16x10的張量，這裏10代表label的數量也是膠囊的數量，然後就進行動態路由算法，經過幾次迭代後，通過更新代碼中b這個對象來計算耦合係數。那麼到此DigitCaps的功能基本實現了。
最後需要提一下的是CapsuleStrengthLayer這一網絡結構。這一層網絡其實是下層每個膠囊的輸出向量計算L2範數，這樣方便後續接入softmax+log-loss的經典分類結構。這在論文中沒有太多提及，可以作爲方便落地的一種手段。下面我們看下完整的神經網絡結構。

MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
            .seed(123)
            .updater(new Adam())
            .list()
            .layer(new ConvolutionLayer.Builder()
                    .nOut(256)
                    .kernelSize(9, 9)
                    .stride(1, 1)	
                    .build())
            .layer(new PrimaryCapsules.Builder(8, 32)
                    .kernelSize(9, 9)
                    .stride(2, 2)
                    .build())
            .layer(new CapsuleLayer.Builder(10, 16, 3).build())
            .layer(new CapsuleStrengthLayer.Builder().build())
            .layer(new ActivationLayer.Builder(new ActivationSoftmax()).build())
            .layer(new LossLayer.Builder(new LossNegativeLogLikelihood()).build())
            .setInputType(InputType.convolutionalFlat(28, 28, 1))
            .build();

以上的膠囊網絡中參數的設置都遵從了Hinton的論文。由於我們需要對Mnist數據集進行分類建模，所以在最後兩層我們使用了softmax+log-loss的結構。這裏需要提一下的是，路由次數默認是3次，如果需要人爲設置，那麼可以在聲明CapsuleLayer的時候調整成爲你需要的參數。最後經過若干輪次的訓練，可以得到以下的評估結果，總體上比較一般，仍有進一步優化的空間。

========================Evaluation Metrics========================
 # of classes:    10
 Accuracy:        0.9557
 Precision:       0.9553
 Recall:          0.9557
 F1 Score:        0.9554
Precision, recall & F1: macro-averaged (equally weighted avg. of 10 classes)

總結

這次我們圍繞2017年的第一版的膠囊網絡進行分析並結合Deeplearning4j進行建模。應當說膠囊網絡的整體結構和傳統CNN還是非常相像的。無論是PrimaryCaps還是DigitCaps都或多或少可以找到Convolution以及Pooling的影子。在之後的每一年，Hinton教授都推出了新的膠囊網絡結構，18年如果沒記錯的話是Matrix的膠囊網絡，19年則是Stacked的結構。在2020年的AAAI會議上，Hinton教授自己說只有Stacked的結構纔是正確的膠囊網絡結構，之前的都是錯誤的。我覺得錯誤倒也不一定，但不完善是肯定的。從一些報道來看，Hinton教授考慮膠囊結構已經有一段時間，但任何創新一般都會參考已有的成果，所謂站在巨人的肩膀上，這也算是人之常情，畢竟CNN雖然有這樣那樣的缺陷，但畢竟現在一些落地的AI場景，CNN仍然是主力。如果有更加直觀以及可解釋的神經網絡結構可以做CNN同樣做的事情，那麼AI肯定會往前走一大步，也期待後面的諸如膠囊網絡的新神經網絡可以發揮這樣的作用。
需要說明下的是，文中的前兩張截圖來自於李宏毅老師課程中關於CapsNet的課件，其餘的都是來自於Hinton教授的論文。有興趣的同學可以去看看，都是很好的說明材料。