CapsNet-Keras

这是论文中CapsNet的Keras实现（分支tf2.2支持TensorFlow 2）：
Sara Sabour, Nicholas Frosst, Geoffrey E Hinton. 动态路由在胶囊之间的应用。NIPS 2017
目前的平均测试误差 = 0.34%，最佳测试误差 = 0.30%。

与论文的不同之处：

• 我们使用了学习率衰减，衰减因子为0.9，每1个epoch衰减一次，而论文中并未给出详细参数（或者他们根本没有使用？）。
• 我们只报告了训练50个epoch后的测试误差。
  在论文中，根据图A.1，我推测他们可能训练了1250个epoch？听起来有些不可思议，也许是我理解错了。
• 我们使用MSE（均方误差）作为重建损失，并且损失系数为lam_recon=0.0005*784=0.392。
  这应该与论文中使用SSE（平方和误差）并设置lam_recon=0.0005是等价的。

警告

请使用Keras==2.0.7配合TensorFlow==1.2后端，否则K.batch_dot函数可能无法正常工作。

然而，如果你使用TensorFlow>=2.0，则请检出分支tf2.2。

使用方法

步骤1. 将本仓库克隆到本地。

git clone https://github.com/XifengGuo/CapsNet-Keras.git capsnet-keras
cd capsnet-keras
git checkout tf2.2 # 仅当使用TensorFlow>=2.0时

步骤2. 安装Keras==2.0.7，并配置TensorFlow==1.2后端。

pip install tensorflow-gpu==1.2
pip install keras==2.0.7

或安装TensorFlow>=2.0

pip install tensorflow==2.2

步骤3. 在MNIST上训练一个CapsNet

使用默认设置进行训练：

python capsulenet.py

更多详细用法可运行帮助命令：

python capsulenet.py -h

步骤4. 测试一个预训练的CapsNet模型

假设你已经使用上述命令训练了一个模型，那么该模型将被保存到result/trained_model.h5。现在只需运行以下命令即可获得测试结果。

$ python capsulenet.py -t -w result/trained_model.h5

它会输出测试准确率，并展示重构后的图像。测试数据与验证数据相同。如果要对新数据进行测试，只需根据需要修改代码即可。

你也可以直接从以下链接下载我训练好的模型：
https://pan.baidu.com/s/1sldqQo1
或
https://drive.google.com/open?id=1A7pRxH7iWzYZekzr-O0nrwqdUUpUpkik

步骤5. 多GPU训练

这需要Keras>=2.0.9。更新Keras后：

python capsulenet-multi-gpu.py --gpus 2

它会自动在多张GPU上训练50个epoch，然后输出在测试集上的性能。但在训练过程中不会报告验证准确率。

结果

测试误差

CapsNet在MNIST上的分类测试误差。平均值和标准差结果基于3次实验。可以通过运行以下命令重现这些结果。

python capsulenet.py --routings 1 --lam_recon 0.0    #CapsNet-v1   
python capsulenet.py --routings 1 --lam_recon 0.392  #CapsNet-v2
python capsulenet.py --routings 3 --lam_recon 0.0    #CapsNet-v3 
python capsulenet.py --routings 3 --lam_recon 0.392  #CapsNet-v4

方法	路由次数	重建	MNIST (%)	论文
基线	--	--	--	0.39
CapsNet-v1	1	否	0.39 (0.024)	0.34 (0.032)
CapsNet-v2	1	是	0.36 (0.009)	0.29 (0.011)
CapsNet-v3	3	否	0.40 (0.016)	0.35 (0.036)
CapsNet-v4	3	是	0.34 (0.016)	0.25 (0.005)

损失和准确率：

训练速度

单块GTX 1070 GPU约为100秒/epoch。
单块GTX 1080Ti GPU约为80秒/epoch。
使用capsulenet-multi-gpu.py在两块GTX 1080Ti GPU上约为55秒/epoch。

重构结果

通过运行以下命令得到的CapsNet-v4重构结果：

python capsulenet.py -t -w result/trained_model.h5

顶部5行是MNIST的真实图像，底部则是对应的重构图像。

操控潜在编码

python capsulenet.py -t --digit 5 -w result/trained_model.h5 

对于每个数字，第i行对应胶囊的第i维，从左到右的列则表示向胶囊某一维的值依次添加[-0.25, -0.2, -0.15, -0.1, -0.05, 0, 0.05, 0.1, 0.15, 0.2, 0.25]。

我们可以看到，每一维都捕捉到了数字的一些特征。不同数字胶囊的同一维可能代表不同的特征。这是因为不同数字是由不同的特征向量（数字胶囊）重构出来的，而这些向量在重构过程中是相互独立的。

其他实现

• PyTorch：

  • XifengGuo/CapsNet-Pytorch
  • timomernick/pytorch-capsule
  • gram-ai/capsule-networks
  • nishnik/CapsNet-PyTorch
  • leftthomas/CapsNet

• TensorFlow：

  • naturomics/CapsNet-Tensorflow
    我参考了该仓库中的部分函数。
  • InnerPeace-Wu/CapsNet-tensorflow
  • chrislybaer/capsules-tensorflow

• MXNet：

  • AaronLeong/CapsNet_Mxnet

• Chainer：

  • soskek/dynamic_routing_between_capsules

• Matlab：

  • yechengxi/LightCapsNet

CapsNet-Keras 快速上手指南

CapsNet-Keras 是基于 Keras 实现的胶囊网络（Capsule Network），复现了 NIPS 2017 论文《Dynamic Routing Between Capsules》中的核心算法。该工具在 MNIST 数据集上表现优异，平均测试错误率低至 0.34%。

环境准备

本工具支持两种运行环境，请根据你的 TensorFlow 版本选择其一：

• 方案 A（推荐，现代环境）：
  • TensorFlow >= 2.0 (例如 2.2)
  • 需切换至项目的 tf2.2 分支
• 方案 B（原始环境）：
  • TensorFlow == 1.2
  • Keras == 2.0.7
  • 注意：若使用此方案，必须严格匹配版本，否则 K.batch_dot 函数可能无法正常工作。

硬件要求：

• 支持 CUDA 的 NVIDIA GPU（可选，但强烈推荐以加速训练）
• 单卡训练速度参考：GTX 1070 约 100s/epoch，GTX 1080Ti 约 80s/epoch。

安装步骤

1. 克隆项目代码

首先将仓库克隆到本地。如果你使用的是 TensorFlow 2.x，请务必切换到 tf2.2 分支。

git clone https://github.com/XifengGuo/CapsNet-Keras.git capsnet-keras
cd capsnet-keras
git checkout tf2.2 # 仅当使用 TensorFlow>=2.0 时执行此命令

2. 安装依赖

根据你的环境方案执行对应的安装命令。国内用户如遇下载缓慢，可配置 pip 使用清华或阿里镜像源（例如添加 -i https://pypi.tuna.tsinghua.edu.cn/simple）。

方案 A：安装 TensorFlow 2.x 环境

pip install tensorflow==2.2

方案 B：安装 TensorFlow 1.2 旧环境

pip install tensorflow-gpu==1.2
pip install keras==2.0.7

基本使用

以下示例展示如何在 MNIST 数据集上进行训练和测试。

1. 训练模型

使用默认参数开始训练胶囊网络。训练将持续 50 个 epoch，完成后模型将保存至 result/trained_model.h5。

python capsulenet.py

如需查看更详细的参数选项（如调整路由次数、重建损失系数等），可运行：

python capsulenet.py -h

2. 测试预训练模型

训练完成后，或在下载了作者提供的预训练模型后，运行以下命令进行测试。该命令会输出测试准确率并生成重建图像对比。

python capsulenet.py -t -w result/trained_model.h5

注：如果没有本地模型，可从 Google Drive 或百度网盘下载作者预训练的模型文件，放置于对应路径后运行上述命令。

3. 多 GPU 训练（可选）

如果你拥有多张 GPU 且已升级 Keras 至 2.0.9 以上版本，可以使用多卡加速训练：

python capsulenet-multi-gpu.py --gpus 2

注意：多卡模式下训练过程中不会实时报告验证集准确率，仅在训练结束后输出测试结果。