本文有一些公式，由于简书不支持LaTeX公式渲染，公式完整版请移步个人博客
欢迎转载，转载请注明出处（简书地址和个人博客地址均可）
理论学习
胶囊结构
胶囊可以看成一种向量化的神经元。对于单个神经元而言，目前的深度网络中流动的数据均为标量。例如多层感知机的某一个神经元，其输入为若干个标量，输出为一个标量（不考虑批处理）；而对于胶囊而言，每个神经元输入为若干个向量，输出为一个向量（不考虑批处理）。前向传播如下所示：
其中$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：$s_j = \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$)的代价函数。
参考资料
文字资料参考weakish翻译的Max Pechyonkin的博客：
此外还参考：
