【胶囊网络】完美复现Hinton论文99.23%

我之前在2024-07-15的时候实现过一版胶囊网络，但是当时无论我怎么训练，都没办法达到Hinton论文里的99.23%（MNIST扩展数据集上）：

前两天心血来潮，又认真读了一下Hinton论文，严酷按照论文要求进行复现：终极达到了Hinton里的性能上限，同时训练速度也比以往那版要快大概五六倍。
1. 导入数据集

1. import torch
2. import torchvision
3. import torchvision.transforms as transforms
4. import matplotlib.pyplot as plt
5. import numpy as np
7. # 定义数据预处理
8. transform = transforms.Compose([
9.     transforms.ToTensor(),
10.     transforms.Normalize((0.1307,), (0.3081,))
11. ])
13. # 下载并加载MNIST训练数据集
14. trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
15. testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)
17. def show_image(image, label):
18.     plt.imshow(image, cmap='gray')
19.     plt.title(f'Label: {label}')
20.     plt.show()
22. # 显示一个训练样本
23. show_image(trainset[0][0][0], trainset[0][1])

复制代码

2. 模子代码

1. import torch
2. import torch.nn as nn
3. import torch.nn.functional as F
4. from tqdm import tqdm
6. def squash(s):
7.     norm = torch.norm(s, dim=-1, keepdim=True)
8.     s_squared_norm = norm ** 2
9.     return (s_squared_norm / (1 + s_squared_norm)) * (s / norm)
11. def routing(u_hat, num_iteratiobns):
12.     # u_hat: (B,N,M,D)
13.     batch_size, num_capsules_i, num_capsules_j, n_dim = u_hat.size()
14.     b_ij = torch.zeros(batch_size, num_capsules_i, num_capsules_j, 1).to(u_hat.device) # (B,N,M,1)
16.     for _ in range(num_iteratiobns):
17.         c_ij = F.softmax(b_ij, dim=1) # (B,N,M,D)
18.         s_j = torch.sum(c_ij * u_hat, dim=1, keepdim=True) # (B,1,M,D)
19.         v_j = squash(s_j) # (B,1,M,D)
20.         b_ij += torch.sum(u_hat * v_j, dim=-1, keepdim=True) # (B,N,M,1)
21.         
22.     return v_j.squeeze(1) # (B,M,D)
24. class PrimaryCaps(nn.Module):
25.     def __init__(self, in_channels=256, out_channels=32, capsule_dim=8, kernel_size=9, stride=2):
26.         super(PrimaryCaps, self).__init__()
27.         self.capsule_dim = capsule_dim
28.         self.out_channels = out_channels
29.         # 使用卷积层来生成初级胶囊的输入（激活初始胶囊向量）
30.         self.conv2 = nn.Conv2d(in_channels, out_channels * capsule_dim, kernel_size=kernel_size, stride=stride)
31.         
32.     def forward(self, x):
33.         # x: (B, C, H, W)
34.         B = x.size(0)
35.         # 进行卷积操作，并将输出调整为胶囊向量的形式，然后整合所有胶囊向量
36.         x = self.conv2(x).permute(0, 2, 3, 1).reshape(B, -1, self.capsule_dim).contiguous() # (B, N, D)
37.         # 对每个胶囊的输出向量应用squash函数
38.         x = squash(x)
39.         return x
41. class DigitCaps(nn.Module):
42.     def __init__(self, num_capsules=10, num_route_nodes=1152, in_channels=8, out_channels=16, num_iterations=3):
43.         """
44.         :param in_channels: 输入胶囊的维度
45.         :param out_channels: 输出胶囊的维度
46.         :param num_capsules: 输出的胶囊数量，对应数字类别数（通常为 10）
47.         :param num_iterations: 动态路由的迭代次数
48.         :param W: 权重矩阵，用于将输入胶囊映射到输出胶囊
49.         """
50.         super(DigitCaps, self).__init__()
51.         self.num_capsules = num_capsules
52.         self.num_iterations = num_iterations
53.         self.W = nn.Parameter(torch.randn(1, num_route_nodes, num_capsules, in_channels, out_channels))
54.         # 也可以所以胶囊共享一个Wj，性能并不会比前者差多少。
55.         # self.W = nn.Parameter(torch.randn(1, 1, num_capsules, in_channels, out_channels))
56.    
57.     def forward(self, x):
58.         # x: (B, N, D1)
59.         # 计算预测向量
60.         x = x.unsqueeze(-2).unsqueeze(-2) # (B, N, 1, 1, D1)
61.         u_hat = torch.matmul(x, self.W).squeeze(-2) # (B, N, M, D2)
62.         # 进行动态路由
63.         v = routing(u_hat, self.num_iterations) # (B, M, D2)
64.         # 返回输出胶囊向量的长度
65.         v = torch.norm(v, dim=-1) # (B, M)
66.         return v
68. class CapesuleNet(nn.Module):
69.     def __init__(self, num_classes=10):
70.         super(CapesuleNet, self).__init__()
71.         self.conv1 = nn.Conv2d(1, 256, kernel_size=9, stride=1) # 灰度图只有一个原始维度
72.         self.primary_capsules = PrimaryCaps()
73.         self.digit_capsules = DigitCaps(num_capsules=num_classes)
75.     def forward(self, x):
76.         x = F.relu(self.conv1(x))
77.         x = self.primary_capsules(x)
78.         x = self.digit_capsules(x)
79.         return x
80.    
81. class MarginLoss(nn.Module):
82.     def __init__(self, m_plus=0.9, m_minus=0.1, lambd=0.5):
83.         super(MarginLoss, self).__init__()
84.         self.m_plus = m_plus
85.         self.m_minus = m_minus
86.         self.lambd = lambd
88.     def forward(self, v, target):
89.         """
90.         v: 形状为 (batch_size, num_classes)，v_k表示第k个数字胶囊的实例化向量的长度
91.         target: 形状为 (batch_size, num_classes)，one-hot编码的目标标签
92.         """
93.         target = torch.eye(num_classes)[target.detach().cpu()].to(device) # one-hot编码
94.         left = torch.clamp(self.m_plus - v, min=0) ** 2
95.         right = torch.clamp(v - self.m_minus, min=0) ** 2
96.         loss = target * left + self.lambd * (1 - target) * right
97.         return torch.mean(torch.sum(loss, dim=1))

复制代码

3. 训练代码

1. num_epochs=10
2. num_classes=10
3. batch_size=128
4. device = 'cuda' if torch.cuda.is_available() else 'cpu'
6. model = CapesuleNet().to(device)
7. margin_loss = MarginLoss()
8. # margin_loss = nn.CrossEntropyLoss()
9. optimizer = torch.optim.AdamW(model.parameters())
11. train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size, shuffle=True, num_workers=2)
12. test_loader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=2)
14. @torch.no_grad()
15. def estimate_acc():
16.     model.eval()
17.     acc = {}
18.     loader = test_loader
19.     correct = 0
20.     total = 0
21.     for images, labels in loader:
22.         images = images.to(device)
23.         labels = labels.to(device)
24.         outputs = model(images)
25.         _, predicted = torch.max(outputs.data, 1)
26.         total += labels.size(0)
27.         correct += (predicted == labels).sum().item()
28.     acc.update({'val': (correct / total) * 100})
29.     model.train()
30.     return acc
32. for epoch in range(num_epochs):
33.     correct = 0
34.     total = 0
35.     with tqdm(total=len(train_loader), desc="epoch %d" % epoch) as pbar:
36.         for images, labels in train_loader:
37.             images = images.to(device)
38.             labels = labels.to(device)
39.             optimizer.zero_grad()
40.             outputs = model(images)
41.             loss = margin_loss(outputs, labels)
42.             loss.backward()
43.             optimizer.step()
44.             predicted = torch.argmax(outputs.data, dim=1)
45.             total += labels.size(0)
46.             correct += (predicted == labels).sum().item()
48.             # 更新进度条
49.             pbar.set_postfix({
50.                 'Loss': f'{loss.item():.4f}',
51.                 'Accuracy': f'{(correct / total) * 100:.4f}%(train)',
52.             })
53.             pbar.update(1)
54.     acc = estimate_acc()
55.     print(f"Accuracy: {acc['val']:.4f}%(val)")

复制代码

可以看到神经网络迅速地收敛，并在第7个epoch达到了论文里的性能上限99.23%！

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。