参考资料:
- 阿里云实践:https://developer.aliyun.com/article/581717
动态路由机制讲解:https://www.bilibili.com/video/BV1oW411H7G1/?spm_id_from=333.337.search-card.all.click&vd_source=3b5e1109bdab0d21b23a5c46c4ed667d
Hinton论文:Dynamic Routing Between Capsules
代码:
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import numpy as np
class CapsuleLayer(nn.Module):
def __init__(self, num_capsules, num_route_nodes, in_channels, out_channels, routing_iters=3, batch_size=128):
super(CapsuleLayer, self).__init__()
self.num_route_nodes = num_route_nodes
self.num_capsules = num_capsules
self.routing_iters = routing_iters
self.W = nn.Parameter(torch.randn(1, num_capsules, num_route_nodes, in_channels, out_channels))
# ([1, 10, 1152, 8, 16])
def forward(self, x):
# ([128, 32, 6, 6, 8])
x = x.view(x.size(0), x.size(1) * x.size(2) * x.size(3), x.size(4)) # ([128, 1152, 8])
x = x.unsqueeze(1) # ([128, 1, 1152, 8])
x = x.repeat(1, self.num_capsules, 1, 1) # # ([128, 10, 1152, 8])
x = x.unsqueeze(3) # ([128, 10, 1152, 8, 1])
u_hat = torch.matmul(x, self.W) # ([128, 10, 1152, 1, 16])
u_hat = u_hat.squeeze(3) # ([128, 10, 1152, 16])
b = torch.zeros(x.size(0), self.num_capsules, self.num_route_nodes, 1) # ([128, 10, 1152, 1])
if next(self.parameters()).is_cuda:
b = b.cuda()
for _ in range(self.routing_iters):
c = F.softmax(b, dim=2)
s = (c * u_hat).sum(dim=2, keepdim=True) # ([128, 10, 1, 16])
v = self.squash(s) # ([128, 10, 1, 16])
if _ < self.routing_iters - 1:
b = b + (u_hat * v).sum(dim=-1, keepdim=True)
return v.squeeze(dim=-1)
def squash(self, input_tensor):
squared_norm = (input_tensor ** 2).sum(dim=-1, keepdim=True)
scale = squared_norm / (1 + squared_norm)
output_tensor = scale * input_tensor / torch.sqrt(squared_norm)
return output_tensor
class PrimaryCaps(nn.Module):
def __init__(self, in_channels, out_channels, kernel_size, stride, dim_capsule):
super(PrimaryCaps, self).__init__()
self.conv = nn.Conv2d(in_channels, out_channels * dim_capsule, kernel_size=kernel_size, stride=stride)
self.dim_capsule = dim_capsule
def forward(self, x):
x = self.conv(x)
batch_size = x.size(0)
out_channels = int(x.size(1) / self.dim_capsule)
height = x.size(2)
width = x.size(3)
# Reshape to [batch_size, out_channels, height, width, dim_capsule]
# 使用 view 方法进行形状改变
x = x.view(batch_size, out_channels, height, width, self.dim_capsule)
return x
class CapsNet(nn.Module):
def __init__(self, input_shape, n_class, routings):
super(CapsNet, self).__init__()
self.input_shape = input_shape
self.n_class = n_class
self.routings = routings
self.conv1 = nn.Conv2d(in_channels=input_shape[0], out_channels=256, kernel_size=9, stride=1)
self.primarycaps = PrimaryCaps(dim_capsule=8, in_channels=256, out_channels=32, kernel_size=9, stride=2)
self.digitcaps = CapsuleLayer(num_capsules=n_class, num_route_nodes=32*6*6, in_channels=8, out_channels=16, routing_iters=routings)
self.decoder = nn.Sequential(
nn.Linear(16 * n_class, 512),
nn.ReLU(inplace=True),
nn.Linear(512, 1024),
nn.ReLU(inplace=True),
nn.Linear(1024, np.prod(input_shape)),
nn.Sigmoid()
)
def forward(self, x):
x = F.relu(self.conv1(x)) # ([128, 256, 20, 20])
x = self.primarycaps(x) # ([128, 32, 6, 6, 8])
x = self.digitcaps(x) # ([128, 10, 1, 16])
# Length of output capsules
lengths = x.norm(dim=-1).squeeze(2) # ([128, 10])
# Reconstruction
x = x.view(x.size(0), -1) # ([128, 160])
reconstructions = self.decoder(x)
reconstructions = reconstructions.view(-1, *self.input_shape)
return lengths, reconstructions
# Create CapsNet model
input_shape = (1, 28, 28) # Example for MNIST
n_class = 10 # Number of classes
routings = 3 # Number of routing iterations
model = CapsNet(input_shape, n_class, routings)
print(model)
# ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
# Load MNIST dataset
batch_size = 128
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# Initialize model, optimizer, and loss function
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = CapsNet(input_shape=(1, 28, 28), n_class=10, routings=3).to(device)
optimizer = optim.Adam(model.parameters())
criterion = nn.CrossEntropyLoss()
# Training function
def train(model, train_loader, optimizer, criterion, epoch):
model.train()
train_loss = 0
correct = 0
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
lengths, reconstructions = model(data)
classification_loss = criterion(lengths, target)
reconstructions_loss = F.mse_loss(reconstructions, data)
loss = classification_loss + reconstructions_loss
loss.backward()
optimizer.step()
train_loss += loss.item()
pred = torch.argmax(lengths, dim=1)
correct += pred.eq(target).sum().item()
if batch_idx % 100 == 0:
print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)} '
f'({100. * batch_idx / len(train_loader):.0f}%)]\tLoss(c/r): {loss.item():.6f} ({classification_loss.item():.6f} / {reconstructions_loss.item():.6f})')
train_loss /= len(train_loader.dataset)
accuracy = 100. * correct / len(train_loader.dataset)
print(f'Train Epoch: {epoch}\tAverage loss: {train_loss:.4f}\tAccuracy: {accuracy:.2f}%')
# Testing function
def test(model, test_loader, criterion):
model.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(device), target.to(device)
lengths, reconstructions = model(data)
loss = criterion(lengths, target) + F.mse_loss(reconstructions, data)
test_loss += loss.item()
pred = lengths.argmax(dim=-1)
correct += pred.eq(target).sum().item()
test_loss /= len(test_loader.dataset)
accuracy = 100. * correct / len(test_loader.dataset)
print(f'\nTest set: Average loss: {test_loss:.4f}, Accuracy: {correct}/{len(test_loader.dataset)} '
f'({accuracy:.2f}%)\n')
# Train the model
epochs = 10
for epoch in range(1, epochs + 1):
train(model, train_loader, optimizer, criterion, epoch)
test(model, test_loader, criterion)
这个是去掉了reconstruction的训练效果:
CapsNet(
(conv1): Conv2d(1, 256, kernel_size=(9, 9), stride=(1, 1))
(primarycaps): PrimaryCaps(
(conv): Conv2d(256, 256, kernel_size=(9, 9), stride=(2, 2))
)
(digitcaps): CapsuleLayer()
(decoder): Sequential(
(0): Linear(in_features=160, out_features=512, bias=True)
(1): ReLU(inplace=True)
(2): Linear(in_features=512, out_features=1024, bias=True)
(3): ReLU(inplace=True)
(4): Linear(in_features=1024, out_features=784, bias=True)
(5): Sigmoid()
)
)
Train Epoch: 1 [0/60000 (0%)] Loss: 2.301852
Train Epoch: 1 [12800/60000 (21%)] Loss: 2.302585
Train Epoch: 1 [25600/60000 (43%)] Loss: 2.302585
Train Epoch: 1 [38400/60000 (64%)] Loss: 2.302585
Train Epoch: 1 [51200/60000 (85%)] Loss: 2.302585
Train Epoch: 1 Average loss: 0.0180 Accuracy: 15.21%
Test set: Average loss: 0.0182, Accuracy: 2033/10000 (20.33%)
Train Epoch: 2 [0/60000 (0%)] Loss: 2.302584
Train Epoch: 2 [12800/60000 (21%)] Loss: 2.302584
Train Epoch: 2 [25600/60000 (43%)] Loss: 2.302585
Train Epoch: 2 [38400/60000 (64%)] Loss: 2.302584
Train Epoch: 2 [51200/60000 (85%)] Loss: 2.302584
Train Epoch: 2 Average loss: 0.0180 Accuracy: 16.46%
Test set: Average loss: 0.0182, Accuracy: 1753/10000 (17.53%)
Train Epoch: 3 [0/60000 (0%)] Loss: 2.302584
Train Epoch: 3 [12800/60000 (21%)] Loss: 2.302584
Train Epoch: 3 [25600/60000 (43%)] Loss: 2.302583
Train Epoch: 3 [38400/60000 (64%)] Loss: 2.302582
Train Epoch: 3 [51200/60000 (85%)] Loss: 2.302579
Train Epoch: 3 Average loss: 0.0180 Accuracy: 18.76%
Test set: Average loss: 0.0182, Accuracy: 2707/10000 (27.07%)
Train Epoch: 4 [0/60000 (0%)] Loss: 2.302577
Train Epoch: 4 [12800/60000 (21%)] Loss: 2.302552
Train Epoch: 4 [25600/60000 (43%)] Loss: 2.302583
Train Epoch: 4 [38400/60000 (64%)] Loss: 2.302581
Train Epoch: 4 [51200/60000 (85%)] Loss: 2.302576
Train Epoch: 4 Average loss: 0.0180 Accuracy: 21.94%
Test set: Average loss: 0.0182, Accuracy: 1935/10000 (19.35%)
Train Epoch: 5 [0/60000 (0%)] Loss: 2.302576
Train Epoch: 5 [12800/60000 (21%)] Loss: 2.302568
Train Epoch: 5 [25600/60000 (43%)] Loss: 2.302579
Train Epoch: 5 [38400/60000 (64%)] Loss: 2.302577
Train Epoch: 5 [51200/60000 (85%)] Loss: 2.302578
Train Epoch: 5 Average loss: 0.0180 Accuracy: 23.75%
Test set: Average loss: 0.0182, Accuracy: 2909/10000 (29.09%)
Train Epoch: 6 [0/60000 (0%)] Loss: 2.302570
Train Epoch: 6 [12800/60000 (21%)] Loss: 2.302553
Train Epoch: 6 [25600/60000 (43%)] Loss: 2.302286
Train Epoch: 6 [38400/60000 (64%)] Loss: 1.515582
Train Epoch: 6 [51200/60000 (85%)] Loss: 1.519818
Train Epoch: 6 Average loss: 0.0147 Accuracy: 66.09%
Test set: Average loss: 0.0119, Accuracy: 9556/10000 (95.56%)
Train Epoch: 7 [0/60000 (0%)] Loss: 1.531795
Train Epoch: 7 [12800/60000 (21%)] Loss: 1.491780
Train Epoch: 7 [25600/60000 (43%)] Loss: 1.485394
Train Epoch: 7 [38400/60000 (64%)] Loss: 1.485206
Train Epoch: 7 [51200/60000 (85%)] Loss: 1.483377
Train Epoch: 7 Average loss: 0.0117 Accuracy: 97.54%
Test set: Average loss: 0.0118, Accuracy: 9801/10000 (98.01%)
Train Epoch: 8 [0/60000 (0%)] Loss: 1.492681
Train Epoch: 8 [12800/60000 (21%)] Loss: 1.511164
Train Epoch: 8 [25600/60000 (43%)] Loss: 1.471053
Train Epoch: 8 [38400/60000 (64%)] Loss: 1.472341
Train Epoch: 8 [51200/60000 (85%)] Loss: 1.496702
Train Epoch: 8 Average loss: 0.0116 Accuracy: 98.42%
Test set: Average loss: 0.0118, Accuracy: 9787/10000 (97.87%)
Train Epoch: 9 [0/60000 (0%)] Loss: 1.476871
Train Epoch: 9 [12800/60000 (21%)] Loss: 1.483141
Train Epoch: 9 [25600/60000 (43%)] Loss: 1.489299
Train Epoch: 9 [38400/60000 (64%)] Loss: 1.478493
Train Epoch: 9 [51200/60000 (85%)] Loss: 1.472141
Train Epoch: 9 Average loss: 0.0116 Accuracy: 98.89%
Test set: Average loss: 0.0117, Accuracy: 9856/10000 (98.56%)
Train Epoch: 10 [0/60000 (0%)] Loss: 1.468742
Train Epoch: 10 [12800/60000 (21%)] Loss: 1.479285
Train Epoch: 10 [25600/60000 (43%)] Loss: 1.471160
Train Epoch: 10 [38400/60000 (64%)] Loss: 1.468553
Train Epoch: 10 [51200/60000 (85%)] Loss: 1.478799
Train Epoch: 10 Average loss: 0.0115 Accuracy: 99.06%
Test set: Average loss: 0.0117, Accuracy: 9819/10000 (98.19%)
可以看到从第六个epoch开始,acurracy呈现断崖式增长,但是前期增长却很慢,而且训练速度非常慢。