【深度学习】基于pytorch的胶囊网络实现

参考资料：

• 阿里云实践：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呈现断崖式增长，但是前期增长却很慢，而且训练速度非常慢。