理论学习

胶囊结构

胶囊可以看成一种向量化的神经元。对于单个神经元而言，目前的深度网络中流动的数据均为标量。例如多层感知机的某一个神经元，其输入为若干个标量，输出为一个标量（不考虑批处理）；而对于胶囊而言，每个神经元输入为若干个向量，输出为一个向量（不考虑批处理）。前向传播如下所示：

其中$I_i$为第i个输入（向量），$W_i$为第i个权值（矩阵），$U_i$为中间变量（向量），由输入和权值叉乘获得。$c_i$为路由权值（标量），需要注意的是该标量是前向传播过程中决定（使用动态路由算法）的，不是通过反向传播优化的参数。Squash为一种激活函数。前向传播使用公式表示如下所示：

$U_i = W_i^T \times I_i$ $S = \sum \limits_{i = 0}^n c_i \cdot U_i$ $Result = Squash(S) = \cfrac{||S||^2}{1+||S||^2} \cdot \cfrac{S}{||S||}$

由以上可以看出，胶囊结构中流动的数据类型为向量，其激活函数Squash输入一个向量，输出一个向量。

动态路由算法

动态路由算法适用于确定胶囊结构中$c_i$的算法，其算法伪代码如下所示：

首先其输入为$U_{j|i}$为本层的中间变量，其中i为这一层胶囊数量，j为下一层胶囊数量，最终获得的胶囊的输出$v_j$，其步骤描述如下：

初始化：初始化一个临时变量b，为一个$i \times j$的全为0的矩阵
获取这一步的连接权值c：$c_i = softmax(b_i)$，将临时变量b通过softmax，保证$c_i$的各分量和为1
获取这一步的加权和结果S：$sj = \sum_i c{ij}u_{j|i}$，按这一步连接权值计算加权和
非线性激活：$v_j = squash(s_j)$，经过非线性激活函数，获取这一步的胶囊输出
迭代临时变量：$b{ij} = b{ij} + u{i|j} \cdot v{j}$，所这一步的输出与中间变量方向相近，增加临时变量b，即增加权值；若这一步输出与中间变量方向相反，减小临时变量b，即减小权值。
若已经迭代到指定次数，输出$v_j$，否侧跳到步骤2

同时，对于迭代次数j，论文中表示过多的迭代会导致过拟合，实践中建议使用3次迭代。

输出与代价函数

输出层胶囊的输出为向量，该向量的长度即为概率。也就是说，前向传播的结果为输出最长向量的输出胶囊所代表的结果。反向传播时，也需要考虑网络的输出为向量而不是标量，因此原论文中了如下的代价函数（每个输出的代价函数，代价函数为所有输出代价函数的和$L = \sum\limits_{c=0}^n L_c$）

$L_c = T_c max(0,m^+ - ||V_c||)^2 + \lambda (1 - T_c)max(0,||v_c|| - m^-) ^ 2$

其中，$T_c$为标量，当分类结果为c时$T_c = 1$，否则$T_c = 0$；$\lambda$为固定值（一般为0.5），用于保证数值稳定性；$m^+$和$m^-$也为固定值：

对于$T_c = 1$的输出胶囊，当输出向量大于$m^+$时，代价函数为0，否则不为0
对于$T_c = 0$的输出胶囊，当输出向量小于$m^-$时，代价函数为0，否则不为0

整体架构

原论文中使举了一个识别MNIST手写数字数据集的例子，网络架构如下图所示：

第一层为普通的卷积层，使用9*9卷积，输出通道数为256，输出数据尺寸为20*20*256
第二层为卷积层，该卷积层由平行的32个卷积层组成，每个卷积层对应向量数据中的一个向量。每个卷积层均为9*9*256*8（输入channel为256，输出channel为8）。因此输出为6*6*32*8，即窗口大小为6*6，输出channel为32，每个数据为8个分量的向量。
第三层为胶囊层，行为类似于全连接层。输入为6*6*32=1152个8分量输入向量，输出为10个16分量的向量，对应的有1152*10个权值，每个权值为8*16的矩阵，最终输出为10个16分量的向量
最终输出10个16分量的向量，最终的分类结果是向量长度最大的输出。

代码阅读（PyTorch）

本次代码阅读并不关心具体的实现方式，主要阅读CapsNet的实现思路

前胶囊层（卷积层）

class PrimaryCaps(nn.Module):
    def __init__(self, num_capsules=8, in_channels=256, out_channels=32, kernel_size=9):
        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]
        u = torch.stack(u, dim=1)
        u = u.view(x.size(0), 32 * 6 * 6, -1)
        return self.squash(u)
    
    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

重点关注forward前向传播部分：

def forward(self, x):
	u = [capsule(x) for capsule in self.capsules]
	u = torch.stack(u, dim=1)
	u = u.view(x.size(0), 32 * 6 * 6, -1)
	return self.squash(u)

self.capsules为num_capsules个[in_channels,out_channels,kernel_size,kernel_size]的卷积层，对应上文所述的第二层卷积层的操作。注意该部分的输出直接被变为[batch size,1152,8]的形式，且通过squash激活函数挤压输出向量

胶囊层

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):
        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)

        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)
            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_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

获得中间向量

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)

这一部分计算中间向量$U_i$

动态路由

for iteration in range(num_iterations):
	c_ij = F.softmax(b_ij)
	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_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)

动态路由的结构中：

第1行计算了softmax函数的结果，对用临时变量b
第5行计算加权和
第6行计算当前迭代次数的输出
第9和10行更新临时向量的值

代价函数

def margin_loss(self, x, labels, size_average=True):
	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)
	loss = labels * left + 0.5 * (1.0 - labels) * right
	loss = loss.sum(dim=1).mean()
	return loss

该函数为代价函数，分别实现了两种情况下($T_c = 0,T_c = 1$)的代价函数。

参考资料

代码来自higgsfield’s github

文字资料参考weakish翻译的Max Pechyonkin的博客：

此外还参考：