使用pytorch的胶囊网络代码及注释

以下是一个简单的 PyTorch 胶囊网络（Capsule Network）的实现代码，带有注释，可以帮助您了解其工作原理。这个代码基于 Hinton 的论文《Dynamic Routing Between Capsules》。

import torch
import torch.nn.functional as F

class CapsuleLayer(torch.nn.Module):
    """
    胶囊层
    """
    def __init__(self, num_capsules, num_route_nodes, in_channels, out_channels):
        super(CapsuleLayer, self).__init__()
        self.num_route_nodes = num_route_nodes
        self.num_capsules = num_capsules
        self.in_channels = in_channels
        self.out_channels = out_channels

        # 胶囊的权重矩阵
        self.W = torch.nn.Parameter(torch.randn(num_capsules, num_route_nodes, in_channels, out_channels))

    def forward(self, x):
        # x的shape: [batch_size, num_route_nodes, in_channels]
        batch_size = x.size(0)

        # x的形状需要扩展为 [batch_size, num_route_nodes, 1, in_channels, 1]
        x = torch.stack([x] * self.num_capsules, dim=2).unsqueeze(4)

        # 复制W以匹配x的shape
        W = torch.cat([self.W] * batch_size, dim=0)

        # 胶囊的预测值
        u_hat = torch.matmul(W, x)

        # 初始化b_ij，shape为 [batch_size, num_capsules, num_route_nodes, 1, 1]
        b_ij = torch.zeros(1, self.num_capsules, self.num_route_nodes, 1, 1)

        # 迭代动态路由
        num_iterations = 3
        for iteration in range(num_iterations):
            # 计算c_ij
            c_ij = F.softmax(b_ij, dim=1)

            # 计算s_j = sum(c_ij*u_hat)
            s_j = (c_ij * u_hat).sum(dim=2, keepdim=True)

            # 计算v_j
            v_j = self.squash(s_j)

            # 更新b_ij
            if iteration != num_iterations - 1:
                a_ij = (u_hat * v_j).sum(dim=3, keepdim=True)
                b_ij = b_ij + a_ij

        # 返回胶囊的输出
        return v_j.squeeze(3)

    def squash(self, s):
        """
        胶囊的激活函数
        """
        squared_norm = (s ** 2).sum(dim=3, keepdim=True)
        scale = squared_norm / (1 + squared_norm)
        return scale * s / torch.sqrt(squared_norm)

class CapsuleNet(torch.nn.Module):
    """
    胶囊网络
    """
    def __init__(self):
        super(CapsuleNet, self).__init__()
        self.conv1 = torch.nn.Conv2d(in_channels=1, out_channels=256, kernel_size=9, stride=1)
        self.primary_capsules = CapsuleLayer(num_capsules=8, num_route_nodes=-1, in_channels=256, out_channels=32)
        self.digit_capsules = CapsuleLayer(num_capsules=10, num_route_nodes=32, in_channels=8, out_channels=16)
        self.decoder = torch.nn.Sequential(
            torch.nn.Linear(16 * 10, 512),
            torch.nn.ReLU(inplace=True),
            torch.nn.Linear(512, 1024),
            torch.nn.ReLU(inplace=True),
            torch.nn.Linear(1024, 784),
            torch.nn.Sigmoid()
        )

    def forward(self, x):
        x = self.conv1(x)
        x = F.relu(x)
        x = self.primary_capsules(x)
        x = self.digit_capsules(x)
        x = x.view(x.size(0), -1)
        classes = (x ** 2).sum(dim=1) ** 0.5
        classes = F.softmax(classes, dim=0)
        _, max_length_indices = classes.max(dim=1)
        masked = torch.autograd.Variable(torch.sparse.torch.eye(10))
        if torch.cuda.is_available():
            masked = masked.cuda()
        masked = masked.index_select(dim=0, index=max_length_indices.squeeze(1).data)
        reconstructions = self.decoder((x * masked[:, :, None]).view(x.size(0), -1))
        return classes, reconstructions

这个代码实现了一个简单的 Capsule Network，包括一个卷积层、两个胶囊层和一个解码器。在前向传播过程中，输入首先通过卷积层和激活函数进行处理，然后通过两个胶囊层，其中第一个胶囊层生成8个胶囊，第二个胶囊层生成10个胶囊，每个胶囊的输出都是一个16维向量。最后，输出通过解码器进行重构，并返回类别概率和重构结果。