An implementation of a variant of a Capsule Network with attention. We use a form of attention to improve signal to noise levels and weight sharing to reduce parameter count. We force the capsule pose to contain a geometric pose and the attention mechanism uses that information to guide the 'routing by agreement' algorithm. The attention mechanism also gives priority to nearby objects as a means to improve the signal to noise ratios for pose estimation. The attention is and algorithm and is not learned, but the part-to-whole relationships are learned. The capsule shares the part-to-whole transformation weights for all children instead of having a separate transformation for every child capsule

Attention algorithm

• uses geometric pose information to guide the 'routing by agreement' algorithm.
• not learned
• is not differentiable.

Weight sharing

• a Capsule is assigned a fixed number of parts
• part-to-whole transformations weights are shared based on the parent-part-child relationship
• capsules of the same type in the same layer share the part-to-whole weights

Differences with the paper:

• Capsule outputs are forced to include geometric pose information that we use for attention routing.
• Capsule outputs are weighted by the distance between the child and parent capsule
• Part-to-whole transformation weights are associated with the parent capsule.
• We train without the reconstruction regularizaton.

TODO

• test with multiple instances per capsule type
• implement rotation and scaling geometric pose
• test with larger images
• more flexible specification of instance and part counts

Step 1. Install Keras>=2.0.9 with TensorFlow>=1.4 backend.

pip install tensorflow-gpu
pip install keras

Step 2. Clone this repository to local.

git clone https://github.com/eastbayml/CAN.git can
cd can

Step 3. Train a CAN Net on synthetic data

First generate some data

python gen_images.py

Train with default settings using sample images include in the data/images.npz:

python train.py

Or use this script to generate new data

python gen_images.py

This code borrows heavily from the excellent implementation of the Capsule network by XifengGuo. E-mail [email protected]

And of course the paper defining the Capsule concept

Sara Sabour, Nicholas Frosst, Geoffrey E Hinton. Dynamic Routing Between Capsules. NIPS 2017