Capsule Layer
PyTorch Capsule Layer, include conv2d and linear layers.
Requirements
- Anaconda
- PyTorch
conda install pytorch torchvision -c pytorch
Installation
pip install git+https://github.com/leftthomas/CapsuleLayer.git@master
To update:
pip install --upgrade git+https://github.com/leftthomas/CapsuleLayer.git@master
Examples
CapsuleConv2d
import torch
from capsule_layer import capsule_cov2d
x = torch.rand(4, 8, 28, 50)
w = torch.randn(2, 8, 4, 3, 5)
if torch.cuda.is_available():
x, w = x.to('cuda'), w.to('cuda')
# routing_type options: ['dynamic', 'k_means']
y, prob = capsule_cov2d(x, w, stride=1, padding=1, routing_type='k_means')or with modules interface:
import torch
from capsule_layer import CapsuleConv2d
x = torch.rand(4, 8, 28, 50)
module = CapsuleConv2d(in_channels=8, out_channels=16, kernel_size=(3, 5), in_length=4, out_length=8, stride=1, padding=1, routing_type='k_means')
if torch.cuda.is_available():
x, module = x.to('cuda'), module.to('cuda')
y, prob = module(x)CapsuleLinear
import torch
from capsule_layer import capsule_linear
x = torch.rand(64, 128, 8)
w = torch.randn(10, 16, 8)
if torch.cuda.is_available():
x, w = x.to('cuda'), w.to('cuda')
# routing_type options: ['dynamic', 'k_means']
y, prob = capsule_linear(x, w, share_weight=True, routing_type='dynamic')or with modules interface:
import torch
from capsule_layer import CapsuleLinear
x = torch.rand(64, 128, 8)
module = CapsuleLinear(out_capsules=10, in_length=8, out_length=16, in_capsules=None, routing_type='dynamic', num_iterations=3)
if torch.cuda.is_available():
x, module = x.to('cuda'), module.to('cuda')
y, prob = module(x)Routing Algorithm
- dynamic routing
import torch
import capsule_layer.functional as F
x = torch.rand(64, 10, 128, 8)
if torch.cuda.is_available():
x = x.to('cuda')
y, prob = F.dynamic_routing(x, num_iterations=10)- k-means routing
import torch
import capsule_layer.functional as F
x = torch.rand(64, 5, 64, 8)
if torch.cuda.is_available():
x = x.to('cuda')
# similarity options: ['dot', 'cosine', 'tonimoto', 'pearson']
y, prob = F.k_means_routing(x, num_iterations=100, similarity='tonimoto')Similarity Algorithm
- tonimoto similarity
import torch
import capsule_layer.functional as F
x1 = torch.rand(64, 16)
x2 = torch.rand(1, 16)
if torch.cuda.is_available():
x1, x2 = x1.to('cuda'), x2.to('cuda')
y = F.tonimoto_similarity(x1, x2, dim=-1)- pearson similarity
import torch
import capsule_layer.functional as F
x1 = torch.rand(32, 8, 16)
x2 = torch.rand(32, 8, 1)
if torch.cuda.is_available():
x1, x2 = x1.to('cuda'), x2.to('cuda')
y = F.pearson_similarity(x1, x2, dim=1)Dynamic Scheduler
- routing iterations
from capsule_layer import CapsuleLinear
from capsule_layer.optim import MultiStepRI
model = CapsuleLinear(3, 4, 7, num_iterations=2)
scheduler = MultiStepRI(model, milestones=[5, 20], addition=3, verbose=True)
# scheduler = MultiStepRI(model, milestones=[5, 20], addition=[3, 3], verbose=True)
for epoch in range(30):
model.train()
...
model.eval()
...
scheduler.step()Contribution
Any contributions to Capsule Layer are welcome!
Copyright and License
Capsule Layer is provided under the MIT License.