基于pytorch的采用胶囊网络分类mnist数据集用CPU的代码

以下是基于PyTorch的胶囊网络分类MNIST数据集的CPU代码：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.datasets as datasets
import torchvision.transforms as transforms

class CapsuleNet(nn.Module):
    def __init__(self):
        super(CapsuleNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=256, kernel_size=9)
        self.primary_capsules = nn.Conv2d(in_channels=256, out_channels=32*8, kernel_size=9, stride=2)
        self.digit_capsules = nn.ModuleList([nn.Linear(in_features=8*32, out_features=16) for _ in range(10)])
        self.decoder = nn.Sequential(
            nn.Linear(in_features=16, out_features=512),
            nn.ReLU(inplace=True),
            nn.Linear(in_features=512, out_features=1024),
            nn.ReLU(inplace=True),
            nn.Linear(in_features=1024, out_features=784),
            nn.Sigmoid()
        )
        
    def forward(self, x):
        x = self.conv1(x)
        x = nn.functional.relu(x)
        x = self.primary_capsules(x)
        x = x.view(x.size(0), 32*8, -1)
        x = self.squash(x)
        x = self.routing(x)
        classes = (x ** 2).sum(dim=-1) ** 0.5
        classes = nn.functional.softmax(classes, dim=-1)
        return classes
    
    def squash(self, x):
        norm = x.norm(dim=-1, keepdim=True)
        scale = (norm ** 2) / (1 + norm ** 2)
        return scale * x / norm
    
    def routing(self, x):
        batch_size = x.size(0)
        b = torch.zeros(batch_size, 10, 32, 1)
        if torch.cuda.is_available():
            b = b.cuda()
        for i in range(3):
            c = nn.functional.softmax(b, dim=1)
            s = (c * x).sum(dim=-1, keepdim=True)
            v = self.squash(s)
            b += (x * v).sum(dim=-1, keepdim=True)
        return v
    
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])

train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True)

capsule_net = CapsuleNet()

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(capsule_net.parameters(), lr=0.001)

for epoch in range(10):
    capsule_net.train()
    train_loss = 0
    train_correct = 0
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = capsule_net(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        train_loss += loss.item()
        train_correct += torch.sum(torch.argmax(output, dim=1) == target)
    print('Epoch: {} Loss: {:.6f} Accuracy: {:.6f}'.format(epoch+1, train_loss/len(train_loader), train_correct/len(train_dataset)))

注意：由于胶囊网络的计算量较大，建议在GPU上运行以获得更好的性能。如果只能使用CPU，则需要耐心等待训练过程的完成。