#

カプセルネットワーク

カプセルネットワークは、フィーチャをカプセルとして埋め込み、投票メカニズムを使用して次のカプセル層にルーティングするニューラルネットワークアーキテクチャです。

他のモデルの実装とは異なり、モジュールだけでは一部の概念を理解するのが難しいため、サンプルを用意しています。

これは、カプセルを使用して MNIST データセットを分類するモデルの注釈付きコードです。

このファイルには、Capsule Networks のコアモジュールの実装が格納されています。

Jindongwang/Pytorch-Capsulenetを使って、論文に関する混乱を解消しました。

これは、MNISTデータセットでカプセルネットワークをトレーニングするためのノートブックです。

32import torch.nn as nn
33import torch.nn.functional as F
34import torch.utils.data
35
36from labml_helpers.module import Module

#

スカッシュ

これは、方程式で与えられる紙からの押しつぶし関数です。 $(1)$

$v_{j} = \frac{∥ s _{j} ∥ ^{2}}{1 + ∥ s _{j} ∥ ^{2}} \frac{s _{j}}{∥ s _{j} ∥}$

$\frac{s _{j}}{∥ s _{j} ∥}$ すべてのカプセルの長さを正規化し、長さが 1 $\frac{∥ s _{j} ∥ ^{2}}{1 + ∥ s _{j} ∥ ^{2}}$ より短いカプセルを縮小します。

39class Squash(Module):

#

54    def __init__(self, epsilon=1e-8):
55        super().__init__()
56        self.epsilon = epsilon

#

s の形は [batch_size, n_capsules, n_features]

58    def forward(self, s: torch.Tensor):

#

$∥ s_{j} ∥^{2}$

64        s2 = (s ** 2).sum(dim=-1, keepdims=True)

#

ゼロにならないように、 $∥ s_{j} ∥$ 計算時にイプシロンを追加します。これがゼロになると、nan 値が与えられ始め、トレーニングは失敗します。 $v_{j} = \frac{∥ s _{j} ∥ ^{2}}{1 + ∥ s _{j} ∥ ^{2}} \frac{s _{j}}{∥ s _{j} ∥ ^{2} + ϵ}$

70        return (s2 / (1 + s2)) * (s / torch.sqrt(s2 + self.epsilon))

#

ルーティングアルゴリズム

これは、このホワイトペーパーで説明されているルーティングメカニズムです。モデルでは複数のルーティングレイヤーを使用できます。

これは、 $s_{j}$ このレイヤーの計算と手順1で説明したルーティングアルゴリズムを組み合わせたものです。

73class Router(Module):

#

in_caps はカプセルの数で、in_d は下のレイヤーのカプセルあたりのフィーチャ数です。out_caps out_d このレイヤーでも同じです。

iterations はルーティングの反復回数で、 $r$ 論文では以下のように表示されています。

84    def __init__(self, in_caps: int, out_caps: int, in_d: int, out_d: int, iterations: int):

#

91        super().__init__()
92        self.in_caps = in_caps
93        self.out_caps = out_caps
94        self.iterations = iterations
95        self.softmax = nn.Softmax(dim=1)
96        self.squash = Squash()

#

$W_{i j}$ これはウェイトマトリックスです。下位レイヤーの各カプセルをこのレイヤーの各カプセルにマッピングします

。

100        self.weight = nn.Parameter(torch.randn(in_caps, out_caps, in_d, out_d), requires_grad=True)

#

u の形は[batch_size, n_capsules, n_features] .これらは下層のカプセルです

。

102    def forward(self, u: torch.Tensor):

#

$\hat{u}_{j ∣ i} = W_{i j} u_{i}$ ここでは $j$ 、このレイヤーのカプセルのインデックスを作成し、下のレイヤー（前のレイヤー）のカプセルのインデックスに使用します。 $i$

111        u_hat = torch.einsum('ijnm,bin->bijm', self.weight, u)

#

$b_{i j}$ 初期ロジットは、 $i$ カプセルと組み合わせるべき対数事前確率です。 $j$ これらはゼロで初期化します

。

116        b = u.new_zeros(u.shape[0], self.in_caps, self.out_caps)
117
118        v = None

#

繰り返し

121        for i in range(self.iterations):

#

ルーティングソフトマックス $c_{i j} = \frac{exp ( b _{i j} )}{\sum _{k} exp ( b _{i k} )}$

123            c = self.softmax(b)

#

$s_{j} = i \sum c_{i j} \hat{u}_{j ∣ i}$

125            s = torch.einsum('bij,bijm->bjm', c, u_hat)

#

$v_{j} = s q u a s h (s_{j})$

127            v = self.squash(s)

#

$a_{i j} = v_{j} \cdot \hat{u}_{j ∣ i}$

129            a = torch.einsum('bjm,bijm->bij', v, u_hat)

#

$b_{i j} \leftarrow b_{i j} + v_{j} \cdot \hat{u}_{j ∣ i}$

131            b = b + a
132
133        return v

#

クラス存在によるマージンロス

出力カプセルごとに個別のマージンロスが使用され、合計損失はそれらの合計になります。各出力カプセルの長さは、入力にクラスが存在する確率です。

$k$ 各出力カプセルまたはクラスの損失は、 $L_{k} = T_{k} m a x (0, m^{+} - ∥ v_{k} ∥)^{2} + λ (1 - T_{k}) m a x (0, ∥ v_{k} ∥ - m^{-})^{2}$

$T_{k}$ $1$ クラスが存在するかどうか、 $k$ $0$ そうでない場合です。 $0$ 損失の最初の要素はクラスが存在しない場合で、 $0$ 2番目の要素はクラスが存在する場合です。 $max (0, x)$ 予測が極端になるのを防ぐために使用されます。 $m^{+}$ $0 .9$ $m^{-}$ $0.1$ 新聞に掲載される予定で、掲載される予定です。

$λ$ ダウンウエイトは、トレーニングの初期段階ですべてのカプセルの長さが落ちるのを防ぐために使用されます。

136class MarginLoss(Module):

#

156    def __init__(self, *, n_labels: int, lambda_: float = 0.5, m_positive: float = 0.9, m_negative: float = 0.1):
157        super().__init__()
158
159        self.m_negative = m_negative
160        self.m_positive = m_positive
161        self.lambda_ = lambda_
162        self.n_labels = n_labels

#

v 、 $v_{j}$ は押しつぶされた出力カプセルです。これには形があります[batch_size, n_labels, n_features] 。つまり、ラベルごとにカプセルがあります。

labels はラベルで、形をしています[batch_size] 。

164    def forward(self, v: torch.Tensor, labels: torch.Tensor):

#

$∥ v_{j} ∥$

172        v_norm = torch.sqrt((v ** 2).sum(dim=-1))

#

$L$ labels ワンホットエンコードされた形状のラベルです [batch_size, n_labels]

176        labels = torch.eye(self.n_labels, device=labels.device)[labels]

#

$L_{k} = T_{k} m a x (0, m^{+} - ∥ v_{k} ∥)^{2} + λ (1 - T_{k}) m a x (0, ∥ v_{k} ∥ - m^{-})^{2}$ loss 形があります[batch_size, n_labels] 。 $L_{k}$ for の計算を並列化しました

。

k

182        loss = labels * F.relu(self.m_positive - v_norm) + \
183               self.lambda_ * (1.0 - labels) * F.relu(v_norm - self.m_negative)

#

$k \sum L_{k}$

186        return loss.sum(dim=-1).mean()