小黑胶囊网络笔记:CapsNet公式与代码对应

1.整体框架图

2.代码细节

import torch
from torch import nn

# device设定
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

（1）squash操作

def squash(x,dim = -1):
    squared_norm = (x ** 2).sum(dim = dim,keepdim = True)
    scale = squared_norm / (1 + squared_norm)
    return scale * x / (squared_norm.sqrt() + 1e-8)
# x = torch.ones([5,2])
# squash(x)

（2）ReLU Conv1层

输入:[batch_size,1,28,28]
中间过程:[batch_size,1,28,28]->(256个9x9的卷积核操作&ReLU)->[batch_size,256,20,20]
输出:[batch_size,256,20,20]

本部分代码较简单，在整体架构代码的类中嵌入进去

（3） PrimaryCaps层（32个胶囊）

参数设定: num_conv_units = 32,in_channels = 256,out_channels = 8,kernel_size = 9,stride = 2
输入:[batch_size,256,20,20]
中间过程: [batch_size,256,20,20]->((8x32)个9x9的卷积核,stride = 2)->[batch_size,8x32,6,6]->(reshape&squash)->[batch_size,6x6x32,8]
含义: 生成6x6x32个dim为8的胶囊，作为下一层的输入。
输出: [batch_size,6x6x32,8]

class PrimaryCaps(nn.Module):
    def __init__(self,num_conv_units,in_channels,out_channels,kernel_size,stride):
        super(PrimaryCaps,self).__init__()
        
        self.conv = nn.Conv2d(in_channels = in_channels,
                              out_channels = out_channels * num_conv_units,
                              kernel_size = kernel_size,
                              stride = stride)
        self.out_channels = out_channels
    def forward(self,x):
        out = self.conv(x)    # out:[batch_size,out_channels * num_conv_units,6,6]
        batch_size = out.shape[0]
        # output:[batch_size,out_capsules * height * weight,out_channels]
        return squash(out.contiguous().view(batch_size,-1,self.out_channels),dim = -1)

(4) DigitCaps层（10个胶囊）

动态路由过程:

参数设定: in_dim = 8,in_caps = 32x6x6,out_caps = 10,out_dim = 16,num_routing = 3,W.shape = [16(out_caps),10(in_caps),16(out_dim),8(in_dim)]
输入: [batch_size,6x6x32(in_caps),8(in_dim)]
中间过程: [batch_size,6x6x32,8]->(unsqueeze)->[batch_size,1,6x6x32,8,1]->(W线性变换)->[batch_size,10(out_caps),6x6x32(in_caps),16(out_dim),1]->(squeeze&detach)->[batch_size,10,6x6x32,16]->(num_routing=3的动态路由&squash)->[batch_size,10,16]

class DigitCaps(nn.Module):
    def __init__(self,in_dim,in_caps,out_caps,out_dim,num_routing):
        # in_dim:Dimensionality of each capsule vector
        # in_caps:Number of imput capsules if digits layer
        # out_caps:Number of capsules in the capsule vector
        # out_dim:Dimensionality,of the output capsule vector
        # num_routing:Number of iterations during routing algorithm
        super(DigitCaps,self).__init__()
        self.in_dim = in_dim
        self.in_caps = in_caps
        self.out_caps = out_caps
        self.out_dim = out_dim
        self.num_routing = num_routing
        self.device = device
        # W:[1,out_caps,in_caps,out_dim,in_dim]
        self.W = nn.Parameter(0.01 * torch.randn(1,out_caps,in_caps,out_dim,in_dim),requires_grad = True)
    def forward(self,x):
        batch_size = x.size(0)
        x = x.unsqueeze(1).unsqueeze(4)    # [batch_size,1,in_caps,in_dim,1]
        u_hat = torch.matmul(self.W,x)    # [batch_size,out_caps,in_caps,out_dims,1]
        u_hat = u_hat.squeeze(-1)    # [batch_size,out_caps,in_caps,out_dims]
        temp_u_hat = u_hat.detach()    # [batch_size,out_caps,in_caps,out_dims] 
        b = torch.zeros(batch_size,self.out_caps,self.in_caps,1).to(device)    # [batch_size,out_caps,in_caps,1]
        for route_iter in range(self.num_routing - 1):
            c = b.softmax(dim = 1)    # [batch_size,out_caps,in_caps,1]
            # [batch_size,out_caps,in_caps,1] .* [batch_size,out_caps,in_caps,out_dims]
            # ->[batch_size,out_caps,in_caps,out_dims] ->(sum.dim=2)->[batch_size,out_caps,out_dims]
            s = (c * temp_u_hat).sum(dim = 2)    # [batch_size,out_caps,out_dims]
            v = squash(s)    # [batch_size,out_caps,out_dims]
            # [batch_size,out_caps,in_caps,out_dims] * [batch_size,out_caps,out_dims,1]
            # -> [batch_size,out_caps,in_caps,1]   含义:v与每个incap的相似度
            uv = torch.matmul(temp_u_hat,v.unsqueeze(-1))
            b += uv    # [batch_size,out_caps,in_caps,1]
        c = b.softmax(dim = 1)
        # [batch_size,out_caps,in_caps,1] .* [batch_size,out_caps,in_caps,out_dims]
        # ->sum.dim = 2->[batch_size,out_caps,out_dims]
        s = (c * u_hat).sum(dim = 2)
        v = squash(s)    # [batch_size,out_caps,out_dims]
        return v

(5) reconstruction&整体架构&CapsuleLoss

(i)reconstruction操作

输入:[batch_size,10(out_caps),16(out_dim)]
输出:[batch_size,784]
含义:将特征向量还原为图片

本部分代码较简单，在整体架构代码的类中嵌入进去

(ii)整体架构

class CapsNet(nn.Module):
    
    def __init__(self):
        super(CapsNet,self).__init__()
        # Conv2d layer
        self.conv = nn.Conv2d(1,256,9)
        self.relu = nn.ReLU(inplace = True)
        # Primary capsule
        self.primary_caps = PrimaryCaps(num_conv_units = 32,in_channels = 256,out_channels = 8,kernel_size = 9,stride = 2)
        # Digit capsule
        self.digit_caps = DigitCaps(in_dim = 8,in_caps = 32 * 6 * 6,out_caps = 10,out_dim = 16,num_routing = 3)
        # Reconstruction layer
        self.decoder = 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):
        out = self.relu(self.conv(x))    # [batch_size,in_channels,20,20] 20 = 28 - 9 + 1
        out = self.primary_caps(out)    # [batch_size,out_capsules*height*weight,out_channels]
        out = self.digit_caps(out)    # [batch_size,out_caps,out_dim]
        logits = torch.norm(out,dim = -1)    # [batch_size,out_caps]
        # [batch_size,out_caps]
        pred = torch.eye(10).to(device).index_select(dim = 0,index = torch.argmax(logits,dim = 1))    
        # Reconstruction
        batch_size = out.shape[0]
        # (out * pred.unsqueeze(2)):[batch_size,out_caps,out_dim]->view->[batch_size,out_caps*out_dim]
        # reconstruction:[batch_size,784]
        reconstruction = self.decoder((out * pred.unsqueeze(2)).contiguous().view(batch_size,-1))
        return logits,reconstruction

(iii)CapsuleLoss

class CapsuleLoss(nn.Module):
    def __init__(self,upper_bound = 0.9,lower_bound = 0.1,lmda = 0.5):
        super(CapsuleLoss,self).__init__()
        self.upper = upper_bound
        self.lower = lower_bound
        self.lmda = lmda
        self.reconstruction_loss_scalar = 5e-4
        self.mse = nn.MSELoss(reduction = 'sum')
    def forward(self,images,labels,logits,reconstructions):
        left = (self.upper - logits).relu() ** 2
        right = (logits - self.lower).relu() ** 2
        margin_loss = torch.sum(labels * left) + self.lmda * torch.sum((1 - labels) * right)
        # Reconstruction loss
        reconstruction_loss = self.mse(reconstructions.contiguous().view(images.shape),images)
        return margin_loss + self.reconstruction_loss_scalar * reconstruction_loss

论文链接:https://arxiv.org/abs/2006.04768
代码链接:https://github.com/Riroaki/CapsNet