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教授的论文。有兴趣的同学可以去看看，都是很好的说明材料。