mirrors/labmlai--annotated_deep_learning_paper_implementations
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:19:01 +08:00
commit 3b90d1192f
2172 changed files with 594509 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
labml_nn/capsule_networks/__init__.py
+185
View File
@@ -0,0 +1,185 @@
"""
---
title: Capsule Networks
summary: >
PyTorch implementation and tutorial of Capsule Networks.
Capsule network is a neural network architecture that embeds features
as capsules and routes them with a voting mechanism to next layer of capsules.
---
# Capsule Networks
This is a [PyTorch](https://pytorch.org) implementation/tutorial of
[Dynamic Routing Between Capsules](https://arxiv.org/abs/1710.09829).
Capsule network is a neural network architecture that embeds features
as capsules and routes them with a voting mechanism to next layer of capsules.
Unlike in other implementations of models, we've included a sample, because
it is difficult to understand some concepts with just the modules.
[This is the annotated code for a model that uses capsules to classify MNIST dataset](mnist.html)
This file holds the implementations of the core modules of Capsule Networks.
I used [jindongwang/Pytorch-CapsuleNet](https://github.com/jindongwang/Pytorch-CapsuleNet) to clarify some
confusions I had with the paper.
Here's a notebook for training a Capsule Network on MNIST dataset.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/capsule_networks/mnist.ipynb)
"""
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data
class Squash(nn.Module):
"""
## Squash
This is **squashing** function from paper, given by equation $(1)$.
$$\mathbf{v}_j = \frac{{\lVert \mathbf{s}_j \rVert}^2}{1 + {\lVert \mathbf{s}_j \rVert}^2}
\frac{\mathbf{s}_j}{\lVert \mathbf{s}_j \rVert}$$
$\frac{\mathbf{s}_j}{\lVert \mathbf{s}_j \rVert}$
normalizes the length of all the capsules, whilst
$\frac{{\lVert \mathbf{s}_j \rVert}^2}{1 + {\lVert \mathbf{s}_j \rVert}^2}$
shrinks the capsules that have a length smaller than one .
"""
def __init__(self, epsilon=1e-8):
super().__init__()
self.epsilon = epsilon
def forward(self, s: torch.Tensor):
"""
The shape of `s` is `[batch_size, n_capsules, n_features]`
"""
# ${\lVert \mathbf{s}_j \rVert}^2$
s2 = (s ** 2).sum(dim=-1, keepdims=True)
# We add an epsilon when calculating $\lVert \mathbf{s}_j \rVert$ to make sure it doesn't become zero.
# If this becomes zero it starts giving out `nan` values and training fails.
# $$\mathbf{v}_j = \frac{{\lVert \mathbf{s}_j \rVert}^2}{1 + {\lVert \mathbf{s}_j \rVert}^2}
# \frac{\mathbf{s}_j}{\sqrt{{\lVert \mathbf{s}_j \rVert}^2 + \epsilon}}$$
return (s2 / (1 + s2)) * (s / torch.sqrt(s2 + self.epsilon))
class Router(nn.Module):
"""
## Routing Algorithm
This is the routing mechanism described in the paper.
You can use multiple routing layers in your models.
This combines calculating $\mathbf{s}_j$ for this layer and
the routing algorithm described in *Procedure 1*.
"""
def __init__(self, in_caps: int, out_caps: int, in_d: int, out_d: int, iterations: int):
"""
`in_caps` is the number of capsules, and `in_d` is the number of features per capsule from the layer below.
`out_caps` and `out_d` are the same for this layer.
`iterations` is the number of routing iterations, symbolized by $r$ in the paper.
"""
super().__init__()
self.in_caps = in_caps
self.out_caps = out_caps
self.iterations = iterations
self.softmax = nn.Softmax(dim=1)
self.squash = Squash()
# This is the weight matrix $\mathbf{W}_{ij}$. It maps each capsule in the
# lower layer to each capsule in this layer
self.weight = nn.Parameter(torch.randn(in_caps, out_caps, in_d, out_d), requires_grad=True)
def forward(self, u: torch.Tensor):
"""
The shape of `u` is `[batch_size, n_capsules, n_features]`.
These are the capsules from the lower layer.
"""
# $$\hat{\mathbf{u}}_{j|i} = \mathbf{W}_{ij} \mathbf{u}_i$$
# Here $j$ is used to index capsules in this layer, whilst $i$ is
# used to index capsules in the layer below (previous).
u_hat = torch.einsum('ijnm,bin->bijm', self.weight, u)
# Initial logits $b_{ij}$ are the log prior probabilities that capsule $i$
# should be coupled with $j$.
# We initialize these at zero
b = u.new_zeros(u.shape[0], self.in_caps, self.out_caps)
v = None
# Iterate
for i in range(self.iterations):
# routing softmax $$c_{ij} = \frac{\exp({b_{ij}})}{\sum_k\exp({b_{ik}})}$$
c = self.softmax(b)
# $$\mathbf{s}_j = \sum_i{c_{ij} \hat{\mathbf{u}}_{j|i}}$$
s = torch.einsum('bij,bijm->bjm', c, u_hat)
# $$\mathbf{v}_j = squash(\mathbf{s}_j)$$
v = self.squash(s)
# $$a_{ij} = \mathbf{v}_j \cdot \hat{\mathbf{u}}_{j|i}$$
a = torch.einsum('bjm,bijm->bij', v, u_hat)
# $$b_{ij} \gets b_{ij} + \mathbf{v}_j \cdot \hat{\mathbf{u}}_{j|i}$$
b = b + a
return v
class MarginLoss(nn.Module):
"""
## Margin loss for class existence
A separate margin loss is used for each output capsule and the total loss is the sum of them.
The length of each output capsule is the probability that class is present in the input.
Loss for each output capsule or class $k$ is,
$$\mathcal{L}_k = T_k \max(0, m^{+} - \lVert\mathbf{v}_k\rVert)^2 +
\lambda (1 - T_k) \max(0, \lVert\mathbf{v}_k\rVert - m^{-})^2$$
$T_k$ is $1$ if the class $k$ is present and $0$ otherwise.
The first component of the loss is $0$ when the class is not present,
and the second component is $0$ if the class is present.
The $\max(0, x)$ is used to avoid predictions going to extremes.
$m^{+}$ is set to be $0.9$ and $m^{-}$ to be $0.1$ in the paper.
The $\lambda$ down-weighting is used to stop the length of all capsules from
falling during the initial phase of training.
"""
def __init__(self, *, n_labels: int, lambda_: float = 0.5, m_positive: float = 0.9, m_negative: float = 0.1):
super().__init__()
self.m_negative = m_negative
self.m_positive = m_positive
self.lambda_ = lambda_
self.n_labels = n_labels
def forward(self, v: torch.Tensor, labels: torch.Tensor):
"""
`v`, $\mathbf{v}_j$ are the squashed output capsules.
This has shape `[batch_size, n_labels, n_features]`; that is, there is a capsule for each label.
`labels` are the labels, and has shape `[batch_size]`.
"""
# $$\lVert \mathbf{v}_j \rVert$$
v_norm = torch.sqrt((v ** 2).sum(dim=-1))
# $$\mathcal{L}$$
# `labels` is one-hot encoded labels of shape `[batch_size, n_labels]`
labels = torch.eye(self.n_labels, device=labels.device)[labels]
# $$\mathcal{L}_k = T_k \max(0, m^{+} - \lVert\mathbf{v}_k\rVert)^2 +
# \lambda (1 - T_k) \max(0, \lVert\mathbf{v}_k\rVert - m^{-})^2$$
# `loss` has shape `[batch_size, n_labels]`. We have parallelized the computation
# of $\mathcal{L}_k$ for for all $k$.
loss = labels * F.relu(self.m_positive - v_norm) + \
self.lambda_ * (1.0 - labels) * F.relu(v_norm - self.m_negative)
# $$\sum_k \mathcal{L}_k$$
return loss.sum(dim=-1).mean()
labml_nn/capsule_networks/mnist.ipynb
+208
View File
@@ -0,0 +1,208 @@
{
"nbformat": 4,
"nbformat_minor": 0,
"metadata": {
"colab": {
"name": "Capsule Networks",
"provenance": [],
"collapsed_sections": []
},
"kernelspec": {
"name": "python3",
"display_name": "Python 3"
},
"accelerator": "GPU"
},
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "AYV_dMVDxyc2"
},
"source": [
"[![Github](https://img.shields.io/github/stars/labmlai/annotated_deep_learning_paper_implementations?style=social)](https://github.com/labmlai/annotated_deep_learning_paper_implementations)\n",
"[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/capsule_networks/mnist.ipynb) \n",
"\n",
"## Training a Capsule Network to classify MNIST digits\n",
"\n",
"This is an experiment to train a Capsule Network to classify MNIST digits using PyTorch."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "AahG_i2y5tY9"
},
"source": [
"Install the `labml-nn` package"
]
},
{
"cell_type": "code",
"metadata": {
"id": "ZCzmCrAIVg0L",
"colab": {
"base_uri": "https://localhost:8080/"
},
"outputId": "7ab15f72-c99f-4097-ecd2-5740ee9ed61c"
},
"source": [
"!pip install labml-nn"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "SE2VUQ6L5zxI"
},
"source": [
"Imports"
]
},
{
"cell_type": "code",
"metadata": {
"id": "0hJXx_g0wS2C"
},
"source": [
"import torch\n",
"\n",
"from labml import experiment\n",
"from labml_nn.capsule_networks.mnist import Configs"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "Lpggo0wM6qb-"
},
"source": [
"Create an experiment"
]
},
{
"cell_type": "code",
"metadata": {
"id": "bFcr9k-l4cAg"
},
"source": [
"experiment.create(name=\"capsule_networks\")"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "-OnHLi626tJt"
},
"source": [
"Initialize [Capsule Network configurations](https://nn.labml.ai/capsule_networks/mnist.html)"
]
},
{
"cell_type": "code",
"metadata": {
"id": "Piz0c5f44hRo"
},
"source": [
"conf = Configs()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "wwMzCqpD6vkL"
},
"source": [
"Set experiment configurations and assign a configurations dictionary to override configurations"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 17
},
"id": "e6hmQhTw4nks",
"outputId": "ebefa8fa-93d2-4131-db95-e27f15aa3aa0"
},
"source": [
"experiment.configs(conf, {'optimizer.optimizer': 'Adam',\n",
" 'optimizer.learning_rate': 1e-3,\n",
" 'inner_iterations': 5})"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "EvI7MtgJ61w5"
},
"source": [
"Set PyTorch models for loading and saving"
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 102
},
"id": "GDlt7dp-5ALt",
"outputId": "9701092b-c88a-4687-c90e-b193c369e59e"
},
"source": [
"experiment.add_pytorch_models({'model': conf.model})"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "markdown",
"metadata": {
"id": "KJZRf8527GxL"
},
"source": [
"Start the experiment and run the training loop."
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/",
"height": 646
},
"id": "aIAWo7Fw5DR8",
"outputId": "5ddbfce3-91f8-4506-e483-1640cb5a14b3"
},
"source": [
"with experiment.start():\n",
" conf.run()"
],
"outputs": [],
"execution_count": null
},
{
"cell_type": "code",
"metadata": {
"id": "oBXXlP2b7XZO"
},
"source": [
""
],
"outputs": [],
"execution_count": null
}
]
}
labml_nn/capsule_networks/mnist.py
+174
View File
@@ -0,0 +1,174 @@
"""
---
title: Classify MNIST digits with Capsule Networks
summary: Code for training Capsule Networks on MNIST dataset
---
# Classify MNIST digits with Capsule Networks
This is an annotated PyTorch code to classify MNIST digits with PyTorch.
This paper implements the experiment described in paper
[Dynamic Routing Between Capsules](https://arxiv.org/abs/1710.09829).
"""
from typing import Any
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.data
from labml import experiment, tracker
from labml.configs import option
from labml_nn.capsule_networks import Squash, Router, MarginLoss
from labml_nn.helpers.datasets import MNISTConfigs
from labml_nn.helpers.metrics import AccuracyDirect
from labml_nn.helpers.trainer import SimpleTrainValidConfigs, BatchIndex
class MNISTCapsuleNetworkModel(nn.Module):
"""
## Model for classifying MNIST digits
"""
def __init__(self):
super().__init__()
# First convolution layer has $256$, $9 \times 9$ convolution kernels
self.conv1 = nn.Conv2d(in_channels=1, out_channels=256, kernel_size=9, stride=1)
# The second layer (Primary Capsules) s a convolutional capsule layer with $32$ channels
# of convolutional $8D$ capsules ($8$ features per capsule).
# That is, each primary capsule contains 8 convolutional units with a 9 × 9 kernel and a stride of 2.
# In order to implement this we create a convolutional layer with $32 \times 8$ channels and
# reshape and permutate its output to get the capsules of $8$ features each.
self.conv2 = nn.Conv2d(in_channels=256, out_channels=32 * 8, kernel_size=9, stride=2, padding=0)
self.squash = Squash()
# Routing layer gets the $32 \times 6 \times 6$ primary capsules and produces $10$ capsules.
# Each of the primary capsules have $8$ features, while output capsules (Digit Capsules)
# have $16$ features.
# The routing algorithm iterates $3$ times.
self.digit_capsules = Router(32 * 6 * 6, 10, 8, 16, 3)
# This is the decoder mentioned in the paper.
# It takes the outputs of the $10$ digit capsules, each with $16$ features to reproduce the
# image. It goes through linear layers of sizes $512$ and $1024$ with $ReLU$ activations.
self.decoder = nn.Sequential(
nn.Linear(16 * 10, 512),
nn.ReLU(),
nn.Linear(512, 1024),
nn.ReLU(),
nn.Linear(1024, 784),
nn.Sigmoid()
)
def forward(self, data: torch.Tensor):
"""
`data` are the MNIST images, with shape `[batch_size, 1, 28, 28]`
"""
# Pass through the first convolution layer.
# Output of this layer has shape `[batch_size, 256, 20, 20]`
x = F.relu(self.conv1(data))
# Pass through the second convolution layer.
# Output of this has shape `[batch_size, 32 * 8, 6, 6]`.
# *Note that this layer has a stride length of $2$*.
x = self.conv2(x)
# Resize and permutate to get the capsules
caps = x.view(x.shape[0], 8, 32 * 6 * 6).permute(0, 2, 1)
# Squash the capsules
caps = self.squash(caps)
# Take them through the router to get digit capsules.
# This has shape `[batch_size, 10, 16]`.
caps = self.digit_capsules(caps)
# Get masks for reconstructioon
with torch.no_grad():
# The prediction by the capsule network is the capsule with longest length
pred = (caps ** 2).sum(-1).argmax(-1)
# Create a mask to maskout all the other capsules
mask = torch.eye(10, device=data.device)[pred]
# Mask the digit capsules to get only the capsule that made the prediction and
# take it through decoder to get reconstruction
reconstructions = self.decoder((caps * mask[:, :, None]).view(x.shape[0], -1))
# Reshape the reconstruction to match the image dimensions
reconstructions = reconstructions.view(-1, 1, 28, 28)
return caps, reconstructions, pred
class Configs(MNISTConfigs, SimpleTrainValidConfigs):
"""
Configurations with MNIST data and Train & Validation setup
"""
epochs: int = 10
model: nn.Module = 'capsule_network_model'
reconstruction_loss = nn.MSELoss()
margin_loss = MarginLoss(n_labels=10)
accuracy = AccuracyDirect()
def init(self):
# Print losses and accuracy to screen
tracker.set_scalar('loss.*', True)
tracker.set_scalar('accuracy.*', True)
# We need to set the metrics to calculate them for the epoch for training and validation
self.state_modules = [self.accuracy]
def step(self, batch: Any, batch_idx: BatchIndex):
"""
This method gets called by the trainer
"""
# Set the model mode
self.model.train(self.mode.is_train)
# Get the images and labels and move them to the model's device
data, target = batch[0].to(self.device), batch[1].to(self.device)
# Increment step in training mode
if self.mode.is_train:
tracker.add_global_step(len(data))
# Run the model
caps, reconstructions, pred = self.model(data)
# Calculate the total loss
loss = self.margin_loss(caps, target) + 0.0005 * self.reconstruction_loss(reconstructions, data)
tracker.add("loss.", loss)
# Call accuracy metric
self.accuracy(pred, target)
if self.mode.is_train:
loss.backward()
self.optimizer.step()
# Log parameters and gradients
if batch_idx.is_last:
tracker.add('model', self.model)
self.optimizer.zero_grad()
tracker.save()
@option(Configs.model)
def capsule_network_model(c: Configs):
"""Set the model"""
return MNISTCapsuleNetworkModel().to(c.device)
def main():
"""
Run the experiment
"""
experiment.create(name='capsule_network_mnist')
conf = Configs()
experiment.configs(conf, {'optimizer.optimizer': 'Adam',
'optimizer.learning_rate': 1e-3})
experiment.add_pytorch_models({'model': conf.model})
with experiment.start():
conf.run()
if __name__ == '__main__':
main()
labml_nn/capsule_networks/readme.md
+20
View File
@@ -0,0 +1,20 @@
# [Capsule Networks](https://nn.labml.ai/capsule_networks/index.html)
This is a [PyTorch](https://pytorch.org) implementation/tutorial of
[Dynamic Routing Between Capsules](https://arxiv.org/abs/1710.09829).
Capsule network is a neural network architecture that embeds features
as capsules and routes them with a voting mechanism to next layer of capsules.
Unlike in other implementations of models, we've included a sample, because
it is difficult to understand some concepts with just the modules.
[This is the annotated code for a model that uses capsules to classify MNIST dataset](mnist.html)
This file holds the implementations of the core modules of Capsule Networks.
I used [jindongwang/Pytorch-CapsuleNet](https://github.com/jindongwang/Pytorch-CapsuleNet) to clarify some
confusions I had with the paper.
Here's a notebook for training a Capsule Network on MNIST dataset.
[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/labmlai/annotated_deep_learning_paper_implementations/blob/master/labml_nn/capsule_networks/mnist.ipynb)