AC慢学屋
首发于AC慢学屋
慢学NLP / Capsule Net 胶囊网络

慢学NLP / Capsule Net 胶囊网络

新年好,由于最近打了那个 Kaggle QIQC,高分 Kernel 几乎都用到了 Capsule Net,借这个机会学习一番,顺便做一点记录。

Capsule 胶囊网络由 Hinton 提出后名噪一时,概念可以说并不复杂,但是却从实验结果和解释性上都有很好的突破。这里参考的视频均来自于 Github 还有李宏毅课程上的截图,视频来源 [李宏毅-胶囊网络]。

Capsule Neural 相较于传统神经网络的区别在于,传统 Neuron 每一个 node 输出为一个激活后的具体数值,而经过 Capsule 输出后得到的则是一个向量,乍一看感觉好好输出个数字,为什么要麻麻烦烦输出一个向量。其实这关乎于一个重点就是神经网络状态的表征,输出向量可以更丰富的表达节点提取的特征,甚至也可以其他降低网络层参数数目的目的。因此对于同一个特征,原本 neuron 的时候我们可能需要多个 nodes 来识别,而现在我们只需要一个 vector,用 vector 中的不同维度来记录同一个特征的不同属性。

这里举个图像的例子,和李宏毅视频中的鸟嘴一致。Capsule 中的神经元的激活情况表示了图像中存在的特定实体的各种性质。这些性质可以包含很多种不同的参数,例如姿势(位置,大小,方向)、变形、速度、反射率,色彩、纹理等等。而输入输出向量的长度表示了某个实体出现的概率,所以它的值必须在 0 到 1 之间。而这些所有的特征都可以用一个向量来涵盖,而下游任务则直接对他取模就可以了。

Capsule Net 和传统 Neuron 的区别

总体流程图为

总体概念图(来源于李宏毅视频)

其中有一个squashing的挤压函数,也就是在模为1的前提下进行缩放,最终得到输出向量。

\large u ^ { 1 } = W ^ { 1 } v ^ { 1 } \quad u ^ { 2 } = W ^ { 2 } v ^ { 2 }\\ s = c_1u^1 + c_2u^2\\ v = \operatorname{Squash} (s)

Squash function函数\large v = \frac { \| s \| ^ { 2 } } { 1 + \| s \| ^ { 2 } } \frac { s } { \| s \| }\\

W1 和 W2 需要学习得到,而 c1 和 c2 却不是学习到的,而是动态路由迭代计算得到的。可以类比于 Pooling 运算,比如 max-pooling 在学习时是不知道那个 neuron 需要 pooling 的,而是经过神经元输出后,套上 Pooling 层,我们才知道哪个是最大的,这是一个online的过程。

Dynamic Routing 示意图

Capsule Dynamic 算法:

 \large \begin{array}  { l } { b _ { 1 } ^ { 0 } = 0 , b _ { 2 } ^ { 0 } = 0 , b _ { 3 } ^ { 0 } = 0 } \\ { For \ r = 1 \text { to } \mathrm { T } \text { do } } \\ \quad { c _ { 1 } ^ { r } , c _ { 2 } ^ { r } , c _ { 3 } ^ { r } = \operatorname { softmax } \left( b _ { 1 } ^ { r - 1 } , b _ { 2 } ^ { r - 1 } , b _ { 3 } ^ { r - 1 } \right) } \\  \quad { s ^ { r } = c _ { 1 } u ^ { 1 } + c _ { 2 } u ^ { 2 } + c _ { 3 } u ^ { 3 } } \\  \quad { a ^ { r } = \operatorname { Squash } \left( s ^ { r } \right) } \\  \quad { b _ { i } ^ { r } = b _ { i } ^ { r - 1 } + a ^ { r } \cdot u ^ { i } }  \end{array}\\

首先初始化 b_i ,对 b_i 取 softmax 表示得到 c_i ,加权求和得到 s^r ,经过挤压后得到 a^r ,这个 a^r 相当于一个改变因子将它与原始输入 u_i 相乘更改最早初始的 b_i 得到更新后的 b _ { i } ^ { r - 1 }

这个迭代的过程,有点类似于排除离群点的概念一旦计算的 a^r 距离 u_1, u_2 比较接近的话,对应的下一轮的 b_i, c_i 也都会变大,对结果施加的影响也就会变大,反之,相关性不高的,对结果施加的影响也会越来越小。

把整个过程可视化之后得到

对C的迭代训练可视化(李宏毅视频截图)

类似于RNN的训练方式,c1和c2就好比于 RNN 中的 hidden layer。

Capsule Network 之所以work的一些解释:

Invariance 和 Equivariance 示意图

传统的 Neural network(如CNN)经过卷积和池化运算后,可以得到一样的特征,也就保证了它学习之后的不变性(Invariance),可以忽略方向上的差别,但是也就只能做到不变性,这样在下游继续接Dense层之后可以归到同一类别。用一句简单话来总结这类Neural Network就是:I don't know the difference.

而Capsule Network则不同,我们在经过Capsule运算后,也可以识别并记录这些差别,能真正做到同变性(Equivariance),用以保真和记录更多信息,而不真正做出反应,因为从Capsule输出的向量我们做取模运算便可以得到其对应的label,取模之后是一样的,但是在取模之前是可以允许向量有差别的,一定程度上也提高了网络的鲁棒性。用一句话总结:I know the difference, but I don't react to it.


实操:

拿两份代码,稍微深入看一下。

第一份是基于CNN和图像来使用卷积的作为 Capsule 输入的,我这里只贴核心的网络搭建,感觉模块分的很清楚(higgsfield/Capsule-Network-Tutorial

  1. 卷积层higgsfield/Capsule-Network-Tutorial卷积层
class ConvLayer(nn.Module):
    def __init__(self, in_channels=1, out_channels=256, kernel_size=9):
        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))

2. 核心 Capsule 层

class PrimaryCaps(nn.Module):
    def __init__(self, num_capsules=8, in_channels=256, out_channels=32, kernel_size=9):
        super(PrimaryCaps, self).__init__()
        # 堆叠 num_capsules capsule net
        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

3. Capsule 连接层,动态路由,输出 V。

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

        # 动态路由3次迭代
        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

4. Decode层,这里加了reconstruction。

class Decoder(nn.Module):
    def __init__(self):
        super(Decoder, self).__init__()
        
        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)
        
        _, max_length_indices = classes.max(dim=1)
        masked = Variable(torch.sparse.torch.eye(10))
        if USE_CUDA:
            masked = masked.cuda()
        masked = masked.index_select(dim=0, index=max_length_indices.squeeze(1).data)
        
        reconstructions = self.reconstraction_layers((x * masked[:, :, None, None]).view(x.size(0), -1))
        reconstructions = reconstructions.view(-1, 1, 28, 28)
        return reconstructions, masked

5. 整个 Capsule Net 网络

class CapsNet(nn.Module):
    def __init__(self):
        super(CapsNet, self).__init__()
        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 = 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):
        return self.margin_loss(x, target) + self.reconstruction_loss(data, reconstructions)
    
    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
    
    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

整个过程就结束了,其实没有很难,但是中间的维度变化以及动态路由还是有不少的巧思。


第二份是在 Kaggle QIQC 时候看到的一个基于 GRU 编码的 Capsule Net,中间用到了numpy的 einsum() 函数来简化张量乘法求和操作。

  1. 定义 Capsule 网络层,参数主要有 num_capsule 和 dim_capsule。这里对是否共享权值进行了分情况的讨论,这块内容可以参考苏神的笔记(揭开迷雾,来一顿美味的Capsule盛宴 - 科学空间|Scientific Spaces)。
class Caps_Layer(nn.Module):
    def __init__(self, input_dim_capsule=GRU_LEN * 2, num_capsule=NUM_CAPSULE, dim_capsule=DIM_CAPSULE,
                 routings=ROUTINGS, kernel_size=(9, 1), share_weights=True,
                 activation='default', **kwargs):
        super(Caps_Layer, self).__init__(**kwargs)

        self.num_capsule = num_capsule
        self.dim_capsule = dim_capsule
        self.routings = routings
        self.kernel_size = kernel_size
        self.share_weights = share_weights
        if activation == 'default':
            self.activation = self.squash
        else:
            self.activation = nn.ReLU(inplace=True)

        if self.share_weights:
            self.W = nn.Parameter(
                nn.init.xavier_normal_(t.empty(1, input_dim_capsule, self.num_capsule * self.dim_capsule)))
        else:
            self.W = nn.Parameter(
                t.randn(BATCH_SIZE, input_dim_capsule, self.num_capsule * self.dim_capsule))

    def forward(self, x):
        if self.share_weights:
            u_hat_vecs = t.matmul(x, self.W)
        else:
            print('add later')

        batch_size = x.size(0)
        input_num_capsule = x.size(1)
        u_hat_vecs = u_hat_vecs.view((batch_size, input_num_capsule,
                                      self.num_capsule, self.dim_capsule))
        u_hat_vecs = u_hat_vecs.permute(0, 2, 1, 3)
        b = t.zeros_like(u_hat_vecs[:, :, :, 0])

        for i in range(self.routings):
            b = b.permute(0, 2, 1)
            c = F.softmax(b, dim=2)
            c = c.permute(0, 2, 1)
            b = b.permute(0, 2, 1)
            outputs = self.activation(t.einsum('bij,bijk->bik', (c, u_hat_vecs)))  # batch matrix multiplication
            # outputs shape (batch_size, num_capsule, dim_capsule)
            if i < self.routings - 1:
                b = t.einsum('bik,bijk->bij', (outputs, u_hat_vecs))  # batch matrix multiplication
        return outputs  # (batch_size, num_capsule, dim_capsule)

    # text version of squash, slight different from original one
    def squash(self, x, axis=-1):
        s_squared_norm = (x ** 2).sum(axis, keepdim=True)
        scale = t.sqrt(s_squared_norm + T_EPSILON)
        return x / scale


class Capsule_Main(nn.Module):
    def __init__(self, embedding_matrix=None, vocab_size=None):
        super(Capsule_Main, self).__init__()
        self.embed_layer = Embed_Layer(embedding_matrix, vocab_size)
        self.gru_layer = GRU_Layer()
        self.gru_layer.init_weights()
        self.caps_layer = Caps_Layer()
        self.dense_layer = Dense_Layer()

    def forward(self, content):
        content1 = self.embed_layer(content)
        content2, _ = self.gru_layer(content1)
        content3 = self.caps_layer(content2)
        output = self.dense_layer(content3)
        return output

写的比较牛逼的,用permute来处理维度变化,用einsum来处理张量操作。值得学习。

for i in range(self.routings):
    b = b.permute(0, 2, 1)
    c = F.softmax(b, dim=2)
    c = c.permute(0, 2, 1)
    b = b.permute(0, 2, 1)
    outputs = self.activation(t.einsum('bij,bijk->bik', (c, u_hat_vecs)))  # batch matrix multiplication
    # outputs shape (batch_size, num_capsule, dim_capsule)
    if i < self.routings - 1:
        b = t.einsum('bik,bijk->bij', (outputs, u_hat_vecs))  # batch matrix multiplication
return outputs  # (batch_size, num_capsule, dim_capsule)

这里 bik,bijk->bij 表示k维相乘后相加,压缩掉了那一维度,对应由 c u 加权求和得到 s 。这块可能需要花点时间理解一下,尤其对应矩阵乘法的三种形式 inner product,element-wise product 和 outer product(这个可能不太常用),掌握之后开发效率事半功倍,尤其对于我这种对维度计算比较头疼的人来说。

参考链接:

张涛:pytorch实现capsule

论智:胶囊网络(Capsule Networks)学习资源汇总

揭开迷雾,来一顿美味的Capsule盛宴 - 科学空间|Scientific Spaces

higgsfield/Capsule-Network-Tutorial

github.com/binzhouchn/c

A basic introduction to NumPy's einsum

Einstein Summation in Numpy

编辑于 2019-02-11

文章被以下专栏收录