Dynamic Routing Between Capsules(胶囊网络与源码解析)

"我要指出 CNN 的问题，告诉你们为什么 CNN 是垃圾。"

伟大先驱Hinton发表在NIPS2017的首次正式提出Capsules，但其实Hinton早在2011年就有类似的想法了，原因是他认为CNN不够好，有很多的缺点，其中最突出的就是：CNN 不使用坐标系，当我们人类观察东西的时候，只要看到一个形状，我们就会给它假定一个坐标系。但CNN不具备这种旋转组合的能力。

• 不变性 (invariance) 。指不随一些变换来识别一个物体，具体变换包括平移 (translation)，旋转 (rotation)，视角 (viewpoint)，放缩 (scale) 等。不变性通常在物体识别上是好事，因为不管一个房子怎么平移、2D旋转、3D旋转和放缩，我们都可以识别出它是房子。但如果我们的任务比物体识别稍微困难一点，比如我想知道它平移了多少个单位，旋转了多少度，放缩了百分之多少，那么不变性远远不够，这时需要的是同变性。
• 同变性(equivariance) 。不变性是表示不随变换 (transformation) 变化，而同变性就是表示的变换等价于变换的表示。

著名的就是将人脸的眼睛和鼻子互换之后这种很不合理的图仍然被识别成person咯，由于不变性，对于模型来说图片中某些部分不管怎么变动是一样的，即便它是不合理的。所以我们想要的不仅仅是有眼睛，鼻子，嘴巴这些组件就能够成为person了，而且需要这三者能够有一定的空间相互关系才行，这是因为CNN在pooling的时候，会丢失掉一些特征间的关系（如max pooling，不管感受野里面哪一块数据最大，结果都是一样的），虽然对于识别整张图像做分类之类没有很大的影响，但在需要更精确得到位置的场景下就无能为力了，所以在这种情况下pooling是没有用的（只能给出这个卷积核区域中的最大值，但是并不知道是从什么方向来的）。

所以能否把这个普通的值，变成方向呢？？即向量不仅仅只是向量了，应该变成能够表示这种空间关系的张量！带着方向感，所以某种程度上，胶囊网络叫做向量神经元 (vector neuron) 甚至张量神经元 (tensor neuron) 更贴切。为了解决这个问题，Hinton提出使用一个叫做“routing-by-agreement”的过程。这意味着，较为底层的特征（手、眼睛、嘴巴等）将只被传送到与之匹配的高层。如果，底层特征包含的是类似于眼睛或者嘴巴的特征，它将传递到“面部”的高层，如果底层特征包含的是类似手指、手掌等特征，它将传递到“手”的高层。

「胶囊」的出发点是，在神经网络中建立更多的结构，然后希望这些新增的结构可以帮助模型更好的泛化。即capsule想要学习到的就是object part组件之间的方向性，以更好的保留对象的姿态信息。如上图里面的v1，v2两个组件通过一些空间关系最后组成一个entity，其中的v1组件的空间位置仍然是用向量，如v1的向左和向右对于向量的变化虽然是某一维的对立如是-1，1，但我们想要的结果是 ：

• 向量不同，但表达的事物是一致的，不影响分类
• 通过控制这一维，可以得到特定旋转的组件
• 并通过多个这样的组件，组合成为最后的结果

如何做到？Dynamic Routing Between Capsules。

Dynamic Routing Between Capsules

所谓的一个胶囊Capsule其实就是一种试图在给定位置上预测是否有特定对象存在，以及其实例化参数的方法。激活向量的方向是由这些对象的实例化参数编码而成，然后用每个激活向量的模长代表着预测出的要找的东西确实存在的概率。如上图是控制方向的capsule（即一个组件），蓝色箭头和黑色箭头，如果要组合成一个帆船，两个箭头有多种旋转的组合，最匹配当前image的就是activations中最长的两个箭头。当然除了旋转角度还可以有很多其他种，比如对象的大小，如何延伸，如何扭曲等等，即很多Capsules去组合更复杂的结果。

首先比较一下传统神经网络和胶囊的区别：

• Neuron：output a value。每个Neuron只做一件事，识别某个特定的区域，静态平移不变的。
• Capsule：output a vector。而capsule想要识别某一类某一组件的特征表达，动态可组合的。
  

在传统神经网络里，一个神经元一般会进行如下的标量操作：

• 输入标量的标量加权，$W_i{x}_i$。
• 对加权后的标量求和，$\sum W_i{x}_i$。
• 对和进行非线性变换生成新标量，Sigmoid，ReLU等等。

而在胶囊网络里面：

• 输入向量与权重矩阵的矩阵乘法，泛函空间映射，$u_i=W_i{v}_i$。因为对于向量的矩阵变化是一种空间变化关系，通过W可以编码图像中低级特征和高级特征之间非常重要的空间关系。
• 加权输入向量，$c_i{u}_i$。这些权重决定当前胶囊将其输出发送到哪个更高级的胶囊（这个c就和神经网络里面的W很类似了，确定每个胶囊的权重），通过动态路由（dynamic routing）的过程完成的。
• 对加权后的向量求和，$s=\sum c_i{u}_i$。 这一点没什么差别。
• 非线性化使用squashing函数生成新向量。该函数将向量进行压缩使得它的最大长度为1，最小长度为0，同时保持其方向不变。看图中的公式可以发现，在向量维度很长的时候它近乎于1，很短的时候近乎于0，这和Sigmoid函数很像，也不是线性变换的。

通过对比，最大的不同应该就是vector to vector（tensor to tensor）了，以及处理这种输入的动态路由。动态路由的流程示意图如下。

先看算法流程如下：

• 对于底层的两个胶囊$v^1$，$v^2$，先通过$W$进行空间映射得到$u^1$，$u^2$.
• 然后利用$b_1^0$，$b_2^0$计算权重$c_1^1$，$c_2^1$，加和得到$s^1$
• $s^1$使用squashing激活就得到了$a^1$
• 利用$a^1$更新$b_1^1$，$b_2^1$，即$b_i^r=b_i^{r-1}+a^r\cdot u^i$
• 通过$b_1^1$，$b_2^1$重复第二步，如图中的红线
• 就得到下一次的$c_1^2$，$c_2^2$，同样加和得到$s^2$
• $s^2$ 也squashing得到$a^2$，同理得到$a^3$
• $a^3$即是更新后的新向量$v$

这个过程看起来有些奇怪，但实际上不就是对每个底层胶囊计算一个softmax权重之后，再用类似RNN的循环操作不断的调整每个胶囊对结果的权重，这也就是“动态”的含义了。这里权重是基于一次结果的a与每个u的点积之后再softmax得到的（很像Attention），然后不断每次的调整。如果预测向量a与上层胶囊的u具有较大的内积（内积描述两个向量的相关度！），则存在自顶向下的反馈（相关度更高就应该加强这个方向的权重），所以就具有增加上层胶囊耦合系数，减小其他胶囊耦合系数的效果（softmax重新分配权重）。

直观来讲这样的处理，是第一层的胶囊试图预测下一层胶囊将要输出什么样的结果，所以不断更新b然后更新c的实质是学习到了部分与整体的关系（part-whole）（部分与整体这一点体现在，整体的结果与每个part胶囊的内积决定，这也就达成了不仅仅是学到图像中有什么组件，还想学到它们更为细致的关系），并自动地将学到的知识推广到不同的新场景中。

这种类似RNN串联的思路，使每次节点得到结果后，需要再回过头比较当前结果的平均方向（sum的整体）和各个节点的建议方向（部分）是否一致，即动态路由以决定重要程度，再通过激活函数表示能够学习到底层节点间的位置空间关系。其实这个更新和knn聚类很像，等于说离群的点就几乎没可能。

Capsules Network：
完整的结构如下图。
首先由普通的Conv得到向量，然后直接reshape成多个（32个），即得到它们的“部分”，将这些part输出到caps胶囊网络中，用上面的方法进行更新和计算。对于minist任务来说，对每个数字的模式各设置一个cap胶囊，即图中10维的DigitCaps，最后使用向量模长代表概率（这也是squashing的结果需要设计在0-1之间）。

实际上最后还会有一个3层的FC做图像重构，然后loss由两部分得到$$loss=los{s}_{margin}+\alpha los{s}_{reconstruction}$$margin loss就是强制使capsule之间的差别越来越大（如预测1的胶囊要和2的胶囊越大越好，$||L_2||$），而采用reconstruction重建损失的好处是可以强制网络保留所有重建图像所需要的信息（其实强烈怀疑capsule之所以work，这个重建loss作用很大）。

pytorch版本的实现代码如下：

import numpy as np

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from torch.optim import Adam
from torchvision import datasets, transforms

USE_CUDA = True #gpu


#载入Mnist数据
class Mnist:
    def __init__(self, batch_size):
        dataset_transform = transforms.Compose([
                       transforms.ToTensor(),
                       transforms.Normalize((0.1307,), (0.3081,))
                   ])

        train_dataset = datasets.MNIST('../data', train=True, download=True, transform=dataset_transform)
        test_dataset = datasets.MNIST('../data', train=False, download=True, transform=dataset_transform)
        
        self.train_loader  = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
        self.test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=True)


#按模型框架图，先由普通卷积层再PrimaryCaps，最后映射到DigitCaps。
#普通卷积层
class ConvLayer(nn.Module):
    def __init__(self, in_channels=1, out_channels=256, kernel_size=9):
        #9x9卷积，256个通道，输出的大小是20x20x256
        #大一些的感受野能在层数较少的情况下得到更多的信息
        super(ConvLayer, self).__init__()

        self.conv = nn.Conv2d(in_channels=in_channels,
                               out_channels=out_channels,
                               kernel_size=kernel_size,
                               stride=1
                             )

    def forward(self, x):
        return F.relu(self.conv(x))

#Primarycaps卷积
class PrimaryCaps(nn.Module):
    def __init__(self, num_capsules=8, in_channels=256, out_channels=32, kernel_size=9):
        #32个平行的卷积，每个数据为8个分量
        super(PrimaryCaps, self).__init__()

        self.capsules = nn.ModuleList([
            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=2, padding=0) 
                          for _ in range(num_capsules)])
    
    def forward(self, x):
        u = [capsule(x) for capsule in self.capsules]#num_capsules个卷积层
        u = torch.stack(u, dim=1)
        u = u.view(x.size(0), 32 * 6 * 6, -1)#窗口大小是6x6
        return self.squash(u)#squash激活函数挤压向量
    
    def squash(self, input_tensor):
        #实现0-1的压缩，同时保持其方向不变
        squared_norm = (input_tensor ** 2).sum(-1, keepdim=True)
        output_tensor = squared_norm *  input_tensor / ((1. + squared_norm) * torch.sqrt(squared_norm))
        return output_tensor

#DigitCaps胶囊层，最后输出为10个16分量的向量，分类结果是向量长度最大的输出
class DigitCaps(nn.Module):
    def __init__(self, num_capsules=10, num_routes=32 * 6 * 6, in_channels=8, out_channels=16):
        super(DigitCaps, self).__init__()

        self.in_channels = in_channels
        self.num_routes = num_routes
        self.num_capsules = num_capsules

        self.W = nn.Parameter(torch.randn(1, num_routes, num_capsules, out_channels, in_channels))

    def forward(self, x):
        #先计算中间向量u
        batch_size = x.size(0)
        x = torch.stack([x] * self.num_capsules, dim=2).unsqueeze(4)

        W = torch.cat([self.W] * batch_size, dim=0)
        u_hat = torch.matmul(W, x)#输入x通过W进行空间映射，编码图像中低级和高级特征之间的空间关系

        #b的初始化为0
        b_ij = Variable(torch.zeros(1, self.num_routes, self.num_capsules, 1))
        if USE_CUDA:
            b_ij = b_ij.cuda()

        #动态路由
        num_iterations = 3
        for iteration in range(num_iterations):
            c_ij = F.softmax(b_ij) #用b计算softmax的权重c
            c_ij = torch.cat([c_ij] * batch_size, dim=0).unsqueeze(4)

            s_j = (c_ij * u_hat).sum(dim=1, keepdim=True)#加权和
            v_j = self.squash(s_j)#当前迭代的输出
            
            if iteration < num_iterations - 1: #更新a和b
                a_ij = torch.matmul(u_hat.transpose(3, 4), torch.cat([v_j] * self.num_routes, dim=1))
                b_ij = b_ij + a_ij.squeeze(4).mean(dim=0, keepdim=True)

        return v_j.squeeze(1)#最后的输出
    
    def squash(self, input_tensor):
        squared_norm = (input_tensor ** 2).sum(-1, keepdim=True)
        output_tensor = squared_norm *  input_tensor / ((1. + squared_norm) * torch.sqrt(squared_norm))
        return output_tensor


#重构函数，强制网络保留所有重建图像所需要的信息
class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()

        #从DigitCaps的16x10开始重建完整图片
        self.reconstraction_layers = nn.Sequential(
            nn.Linear(16 * 10, 512),
            nn.ReLU(inplace=True),
            nn.Linear(512, 1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, 784),
            nn.Sigmoid()
        )
        
    def forward(self, x, data):
        classes = torch.sqrt((x ** 2).sum(2))
        classes = F.softmax(classes)#最后输出的x向量最长的为最后的结果
        
        _, max_length_indices = classes.max(dim=1)
        masked = Variable(torch.sparse.torch.eye(10))#one-hot
        if USE_CUDA:
            masked = masked.cuda()
        masked = masked.index_select(dim=0, index=max_length_indices.squeeze(1).data)
        #3层FC做重建
        reconstructions = self.reconstraction_layers((x * masked[:, :, None, None]).view(x.size(0), -1))
        reconstructions = reconstructions.view(-1, 1, 28, 28)
        
        return reconstructions, masked

#CapsNet完整的流程
class CapsNet(nn.Module):
    def __init__(self):
        super(CapsNet, self).__init__()
        #由四个class组成
        self.conv_layer = ConvLayer()
        self.primary_capsules = PrimaryCaps()
        self.digit_capsules = DigitCaps()
        self.decoder = Decoder()
        
        self.mse_loss = nn.MSELoss()#均方差
        
    def forward(self, data):
        #三层的胶囊网络结构得到output
        output = self.digit_capsules(self.primary_capsules(self.conv_layer(data)))
        #用输出重建图像
        reconstructions, masked = self.decoder(output, data)
        return output, reconstructions, masked
    
    def loss(self, data, x, target, reconstructions):
        #完整的loss由margin和reconstruction两部分组成
        return self.margin_loss(x, target) + self.reconstruction_loss(data, reconstructions)
    
    def margin_loss(self, x, labels, size_average=True):
        #margin loss强制使capsule之间（如预测1和预测2）的差别越来越大
        batch_size = x.size(0)

        v_c = torch.sqrt((x**2).sum(dim=2, keepdim=True))#长度表示某个类别的概率

        #上边界和下边界
        left = F.relu(0.9 - v_c).view(batch_size, -1)
        right = F.relu(v_c - 0.1).view(batch_size, -1)

        #惩罚偏离边缘（错位的分类对边缘0.1 or 0.9的距离）
        #如果预测是0.8，label是1，那么loss是0.1很小
        #如果label是0，那么loss的惩罚要算与right的距离，其中0.5是downweight
        loss = labels * left + 0.5 * (1.0 - labels) * right
        loss = loss.sum(dim=1).mean()

        return loss
    
    def reconstruction_loss(self, data, reconstructions):
        loss = self.mse_loss(reconstructions.view(reconstructions.size(0), -1), data.view(reconstructions.size(0), -1))
        return loss * 0.0005 #这个系数为0.0005


capsule_net = CapsNet()
if USE_CUDA:
    capsule_net = capsule_net.cuda()
optimizer = Adam(capsule_net.parameters())#优化器


batch_size = 100
mnist = Mnist(batch_size)#导入数据

n_epochs = 30#周期

#开始训练
for epoch in range(n_epochs):
    capsule_net.train()#调到训练模式
    train_loss = 0
    for batch_id, (data, target) in enumerate(mnist.train_loader):

        target = torch.sparse.torch.eye(10).index_select(dim=0, index=target)#手写数字10维
        data, target = Variable(data), Variable(target)

        if USE_CUDA:#gpu
            data, target = data.cuda(), target.cuda()

        optimizer.zero_grad() #梯度清零
        output, reconstructions, masked = capsule_net(data) #得到模型输出
        loss = capsule_net.loss(data, output, target, reconstructions) #计算loss
        loss.backward() #反向传播
        optimizer.step() #参数更新

        train_loss += loss.data[0]#记录总loss
        
        if batch_id % 100 == 0: #定期打印结果
            print "train accuracy:", sum(np.argmax(masked.data.cpu().numpy(), 1) == 
                                   np.argmax(target.data.cpu().numpy(), 1)) / float(batch_size)
        
    print train_loss / len(mnist.train_loader)
        
    capsule_net.eval()#评估模式
    test_loss = 0
    for batch_id, (data, target) in enumerate(mnist.test_loader):

        target = torch.sparse.torch.eye(10).index_select(dim=0, index=target)
        data, target = Variable(data), Variable(target)

        if USE_CUDA:#gpu
            data, target = data.cuda(), target.cuda()

        output, reconstructions, masked = capsule_net(data)#得到评估结果
        loss = capsule_net.loss(data, output, target, reconstructions)#计算loss

        test_loss += loss.data[0]
        
        if batch_id % 100 == 0: #定期打印结果
            print "test accuracy:", sum(np.argmax(masked.data.cpu().numpy(), 1) == 
                                   np.argmax(target.data.cpu().numpy(), 1)) / float(batch_size)
    
    print test_loss / len(mnist.test_loader)


#可视化
import matplotlib
import matplotlib.pyplot as plt

def plot_images_separately(images):
    "Plot the six MNIST images separately."
    fig = plt.figure()
    for j in xrange(1, 7):
        ax = fig.add_subplot(1, 6, j)
        ax.matshow(images[j-1], cmap = matplotlib.cm.binary)
        plt.xticks(np.array([]))
        plt.yticks(np.array([]))
    plt.show()

plot_images_separately(data[:6,0].data.cpu().numpy())#原结果
plot_images_separately(reconstructions[:6,0].data.cpu().numpy())#重建结果

完整逐行代码解析：：https://github.com/nakaizura/Source-Code-Notebook/tree/master/Capsules

Capsules的重点

• 想学到具体的实体或者关系模式，等变性
• 得到组件方向性，能动态的调整agreement routing
• Margin + reconstruction

Capsules的优缺点
优点主要有3点：

• 胶囊的输出只会被引导至合适的下一层胶囊，这些胶囊将会获得非常干净的输入信号，并且更加精确的特定对象的姿态。
• 通过审视激活的路径，我们能清楚的观察到组件的层次结构。如下图，可以发现某些中间cap确实学习到了一些姿态信息，某些控制粗细，某些与倾斜度，方向有关等。
• routing将会帮助解析拥有大量重叠在一起的对象所构成的拥挤场景或者暧昧不清。即overlapping handing，也是paper中的一个实验，如下下图，能比较好的显示出overlapping 的现象。
  

缺点是：

• 训练非常的慢，最大原因是routing有内部循环。
• 无论是什么类型的对象，在给定位置永远只有一个胶囊，所以它不能检测出非常靠近的同一类型的对象，这种效应叫“挤出效应”，在人类视觉中也有这种现象，所以其实不算是很大的问题。

Capsules的到底学习什么？
下图是李宏毅老师对invariance和equivariance的discusstion了。对于左边的图是普通的CNN经过max pooling的结果，对于不同的输入max的结果都是3，这种情况就容易导致在人眼镜，嘴巴等一些组件明明调换了，CNN却仍然能识别出是person的不合理现象，“i dont know the difference”。因为CNN想要学习到的是某一类整体的通用表示，不管什么样的图，都得到同样的结果就是最好的。

而对于右边capsule，输入的两个“1”的倾斜度是对称的，通过不同的cap得到的v是可以不一样的，这能包含一些姿态信息，但只有保证最后的结果||v||值类似的就行，“i know the difference，but i do not react to it”。

Capsules的本质是什么？
Hopping Memory Networks+ Attention（或 rnn）。毕竟每次都从前一次取值（memory）做调整（attention），而且多个（hopping）。