前幾天，Sara Sabour 開源了一份 Capsule 代碼，該代碼是論文 Dynamic Routing between Capsules 中所採用的實現。其實早在去年剛公布此論文，機器之心就曾詳解解讀過核心思想與基本代碼，我們採用的代碼也是各研究者嘗試復現論文結果的模型。而最近 Sara 開放的代碼是標準的官方實現，因此我們希望能解讀部分核心代碼，並探討其與 naturomics 等人實現過程的差異。
Sara 實現地址：https://github.com/Sarasra/models/tree/master/research/capsules
我們主要根據 Sara 的代碼解釋了 CapsNet 的架構與實現方法，包括 Squash 非線性函數、動態路由更新方法、PrimaryCaps 層與 DigitCaps 的實現過程，還有最後的 Margin Loss 度量函數。我們希望入門讀者先了解 Capsule 的概念與 CapsNet 的基本架構，以下我們提供了 Capsule 論文解讀與基本概念。
先讀懂 CapsNet 架構然後用 TensorFlow 實現，這應該是最詳細的教程了
終於，Geoffrey Hinton 那篇備受關注的 Capsule 論文公開了
淺析 Geoffrey Hinton 最近提出的 Capsule 計劃
本文章所解釋的代碼均來自於 capsule_model.py 和 layers.py 兩個文件，它們也是整個實現的核心部分。下面，我們將從基本的 Capsule 概念與 Squash 非線性函數開始解析 Sara 所完成的實現。
在論文中，Geoffrey Hinton 介紹 Capsule 為：「Capsule 是一組神經元，其輸入輸出向量表示特定實體類型的實例化參數（即特定物體、概念實體等出現的概率與某些屬性）。我們使用輸入輸出向量的長度表徵實體存在的概率，向量的方向表示實例化參數（即實體的某些圖形屬性）。同一層級的 capsule 通過變換矩陣對更高級別的 capsule 的實例化參數進行預測。當多個預測一致時（本論文使用動態路由使預測一致），更高級別的 capsule 將變得活躍。」
Capsule 中神經元的激活情況表示了圖像中存在的特定實體的各種性質。這些性質可以包含很多種不同的參數，例如姿勢（位置、大小、方向）、變形、速度、反射率、色彩、紋理等等。而輸入輸出向量的長度表示了某個實體出現的概率，所以它的值必須在 0 到 1 之間。
為了實現這種壓縮，並完成 Capsule 層級的激活功能，Hinton 等人使用了一個被稱為「squashing」的非線性函數。該非線性函數確保短向量的長度能夠縮短到幾乎等於零，而長向量的長度壓縮到接近但不超過 1 的情況。以下是該非線性函數的表達式：
其中 v_j 為 Capsule j 的輸出向量，s_j 為上一層所有 Capsule 輸出到當前層 Capsule j 的向量加權和，簡單說 s_j 是 Capsule j 的輸入向量。
在 Sara 提供的實現中，她們使用以下方法定義非線性激活函數。其中輸入張量對於全連接 Capsule 層來說維度為 [batch, num_channels, num_atoms]，對於卷積 Capsule 層來說，輸入的維度為 [batch, num_channels, num_atoms, height, width]。該函數將輸出一組激活張量 v_j，其維度等於輸入張量的維度。
def _squash(input_tensor):
"""Applies norm nonlinearity (squash) to a capsule layer.
Args:
input_tensor: Input tensor. Shape is [batch, num_channels, num_atoms] for a
fully connected capsule layer or
[batch, num_channels, num_atoms, height, width] for a convolutional
capsule layer.
Returns:
A tensor with same shape as input (rank 3) for output of this layer.
"""
with tf.name_scope('norm_non_linearity'):
norm = tf.norm(input_tensor, axis=2, keep_dims=True)
norm_squared = norm * norm
return (input_tensor / norm) * (norm_squared / (1 + norm_squared))
因為按照 Hinton 的思想，找到最好的處理路徑就等價於正確處理了圖像，所以在 Capsule 中加入 Routing 機制可以找到一組係數 c_ij，它們能令預測向量 u_j|i hat 最符合輸出向量 v_j，即最符合輸出的輸入向量，這樣我們就找到了最好的路徑。
圖1：該圖展示了 Capsule 的層級結構與動態 Routing 的過程
按照論文 Dynamic Routing between Capsules 所述，c_ij 為耦合係數，該係數由動態 Routing 過程迭代地更新與確定。Capsule i 和後一層級所有 Capsule 間的耦合係數和為 1。此外，該耦合係數由「routing softmax」決定，且 softmax 函數中的 logits b_ij 初始化為 0，耦合係數 c_ij 的 softmax 計算方式為：
b_ij 依賴於兩個 Capsule 的位置與類型，但不依賴於當前的輸入圖像。我們可以通過測量後面層級中每一個 Capsule j 的當前輸出 v_j 和 前面層級 Capsule i 的預測向量間的一致性，然後藉助該測量的一致性迭代地更新耦合係數。本論文簡單地通過內積度量這種一致性，即
，這一部分也就涉及到使用 Routing 更新耦合係數。
有意思的是，Sara 的實現會添加一個 leaky_routing 函數，按照該函數的定義，它會添加額外的維度以路由分對數（logits）。如果需要執行路由的張量維度與上層任意 Capsule 單元不匹配，那麼該函數將允許激活的 Capsule 單元在額外的維度中進行路由。如下參數 logits 為需要路由的張量，其中它的維度在全連接層的情況下為 [input_capsule_num, output_capsule_num]，在卷積層的情況下回多增加兩個維度。output_dim 為分對數的第二個維度，即輸出的 Capsule 單元數。
def _leaky_routing(logits, output_dim):
"""Adds extra dimmension to routing logits.
This enables active capsules to be routed to the extra dim if they are not a
good fit for any of the capsules in layer above.
Args:
logits: The original logits. shape is
[input_capsule_num, output_capsule_num] if fully connected. Otherwise, it
has two more dimmensions.
output_dim: The number of units in the second dimmension of logits.
Returns:
Routing probabilities for each pair of capsules. Same shape as logits.
"""
# leak is a zero matrix with same shape as logits except dim(2) = 1 because
# of the reduce_sum.
leak = tf.zeros_like(logits, optimize=True)
leak = tf.reduce_sum(leak, axis=2, keep_dims=True)
leaky_logits = tf.concat([leak, logits], axis=2)
leaky_routing = tf.nn.softmax(leaky_logits, dim=2)
return tf.split(leaky_routing, [1, output_dim], 2)[1]
如上所示，tf.zeros_like 將構建一個與 logits 維度相同的張量，其中每個元素都為 0。在求和處理後，leaky_logits 將在第三個維度拼接 leak 和 logits 張量。在對 leaky_logits 第 3 個維度進行 Softmax 後就相當於計算了以下偽代碼中的耦合係數 c_ij，我們需要使用它執行進一步的路由。此外，_leaky_routing 應該是沒有應用 Squash 非線性激活。因此該函數真實的意義與過程可能需要進一步探討，這裡的理解不是很完全。
以上_leaky_routing 函數會在完整執行路由和非線性壓縮的_update_routing 函數中調用，所以在我們查看完整的路由函數前可以先複習以下原論文中所描述的動態路由偽代碼。
Routing 過程就是圖1右邊表述的更新過程，我們會計算 v_j 與 u_j|i hat 的乘積並將它與原來的 b_ij 相加而更新 b_ij，然後利用 softmax(b_ij) 更新 c_ij 而進一步修正了後一層的 Capsule 輸入 s_j。當輸出新的 v_j 後又可以迭代地更新 c_ij，這樣我們不需要反向傳播而直接通過計算輸入與輸出的一致性更新參數。
對於所有在 l 層的 Capsule i 和在 l+1 層的 Capsule j，先初始化 b_ij 等於零。然後迭代 r 次，每次先根據 b_i 計算 c_i，然後在利用 c_ij 與 u_j|i hat 計算 s_j 與 v_j。利用計算出來的 v_j 更新 b_ij 以進入下一個迭代循環更新 c_ij。該 Routing 算法十分容易收斂，基本上通過 3 次迭代就能有不錯的效果。
在以下定義的路由更新過程中，_update_routing 函數會對經精煉的輸入張量求和並執行 Squash 非線性變換。它的輸出激活值可作為 PrimaryCaps 層和 DigitCaps 層的最終輸出，後面我們將詳細討論這兩個層的實現。按照 _update_routing 函數的說明，它會基於當前層的激活值與前一層投票結果（即線性組合結果）之間的相似性，迭代地更新 logits 路由結果，即對輸入張量進行精煉。
def _update_routing(votes, biases, logit_shape, num_dims, input_dim, output_dim,
num_routing, leaky):
votes_t_shape = [3, 0, 1, 2]
for i in range(num_dims - 4):
votes_t_shape += [i + 4]
r_t_shape = [1, 2, 3, 0]
for i in range(num_dims - 4):
r_t_shape += [i + 4]
votes_trans = tf.transpose(votes, votes_t_shape)
def _body(i, logits, activations):
"""Routing while loop."""
# route: [batch, input_dim, output_dim, ...]
if leaky:
route = _leaky_routing(logits, output_dim)
else:
route = tf.nn.softmax(logits, dim=2)
preactivate_unrolled = route * votes_trans
preact_trans = tf.transpose(preactivate_unrolled, r_t_shape)
preactivate = tf.reduce_sum(preact_trans, axis=1) + biases
activation = _squash(preactivate)
activations = activations.write(i, activation)
# distances: [batch, input_dim, output_dim]
act_3d = tf.expand_dims(activation, 1)
tile_shape = np.ones(num_dims, dtype=np.int32).tolist()
tile_shape[1] = input_dim
act_replicated = tf.tile(act_3d, tile_shape)
distances = tf.reduce_sum(votes * act_replicated, axis=3)
logits += distances
return (i + 1, logits, activations)
activations = tf.TensorArray(
dtype=tf.float32, size=num_routing, clear_after_read=False)
logits = tf.fill(logit_shape, 0.0)
i = tf.constant(0, dtype=tf.int32)
_, logits, activations = tf.while_loop(
lambda i, logits, activations: i < num_routing,
_body,
loop_vars=[i, logits, activations],
swap_memory=True)
return activations.read(num_routing - 1)
如上所示，votes 為前一層經轉換的輸出張量。num_dims 為輸入 votes 的維度數量，對於全連接 Capsule 層來說，它的維度為 4，對於卷積層來說，它的維度為 6。input_dim 為 輸入層的 Capsule 單元數，output_dim 為輸出層的 Capsule 單元數。num_routing 為路由的迭代次數，而 leaky 則代表著是否使用前面定義的滲漏路由 _leaky_routing。
以上 _update_routing 函數最終會輸出一個激活張量，即上面原論文偽代碼中的 v_j。在初步討論了路由算法後，我們可以查看它到底用在了哪些地方，即哪些運算需要執行路由算法。從 Sara 的代碼上看，CapsNet 應該使用了兩次 Routing，即在第二層的卷積層和第三層的 Capsule 全連接層後各調用了一次。這一點與論文的描述和其他研究者所實現的復現有所不同。
下面，我們將依據原論文與 Sara 開源的實現討論 CapsNet 主體架構和 Margin loss 度量。這一部分是該論文與實現的核心，因此我們將重點關注這一部分而忽略後面構建的重構網絡與重構損失。
以下是 CapsNet 的整體架構：
第一個卷積層使用了 256 個 9×9 卷積核，步幅為 1，且使用了 ReLU 激活函數。該卷積操作應該沒有使用 Padding，輸出的張量才能是 20×20×256。此外，CapsNet 的卷積核感受野使用的是 9×9，相比於其它 3×3 或 5×5 的要大一些，這個能是因為較大的感受野在 CNN 層次較少的情況下能感受的信息越多。這兩層間的權值數量應該為 9×9×256+256=20992。
隨後，第二個卷積層 PrimaryCaps 開始作為 Capsule 層的輸入而構建相應的張量結構，我們可以從上圖看出第二層卷積操作後生成的張量維度為 6×6×8×32。如果我們先考慮 32 個（32 channel）9×9 的卷積核在步幅為 2 的情況下做卷積，那麼實際上得到的是傳統的 6×6×32 的張量，即等價於 6×6×1×32。因為傳統卷積操作每次計算的輸出都是一個標量，而 PrimaryCaps 的輸出需要是一個長度為 8 的向量，因此傳統卷積下的三維輸出張量 6×6×1×32 就需要變化為四維輸出張量 6×6×8×32。
第三層 DigitCaps 在第二層輸出的向量基礎上進行傳播與 Routing 更新。第二層共輸出 6×6×32=1152 個向量，每一個向量的維度為 8，即第 i 層共有 1152 個 Capsule 單元。而第三層 j 有 10 個標準的 Capsule 單元，每個 Capsule 的輸出向量有 16 個元素。前一層的 Capsule 單元數是 1152 個，那麼 w_ij 將有 1152×10 個，且每一個 w_ij 的維度為 8×16。當 u_i 與對應的 w_ij 相乘得到預測向量後，我們會有 1152×10 個耦合係數 c_ij，對應加權求和後會得到 10 個 16×1 的輸入向量。將該輸入向量輸入到「squashing」非線性函數中求得最終的輸出向量 v_j，其中 v_j 的長度就表示識別為某個類別的概率。
如下所示定義了構建 Capsule 層級的主體函數：
def _build_capsule(self, input_tensor, num_classes):
capsule1 = layers.conv_slim_capsule(
input_tensor,
input_dim=1,
output_dim=self._hparams.num_prime_capsules,
layer_name='conv_capsule1',
num_routing=1,
input_atoms=256,
output_atoms=8,
stride=2,
kernel_size=9,
padding=self._hparams.padding,
leaky=self._hparams.leaky,)
capsule1_atom_last = tf.transpose(capsule1, [0, 1, 3, 4, 2])
capsule1_3d = tf.reshape(capsule1_atom_last,
[tf.shape(input_tensor)[0], -1, 8])
_, _, _, height, width = capsule1.get_shape()
input_dim = self._hparams.num_prime_capsules * height.value * width.value
return layers.capsule(
input_tensor=capsule1_3d,
input_dim=input_dim,
output_dim=num_classes,
layer_name='capsule2',
input_atoms=8,
output_atoms=16,
num_routing=self._hparams.routing,
leaky=self._hparams.leaky,)
該函數在輸入 5 維張量 [batch, 1, 256, height, width] 和目標類別的數量 num_classes 後會輸出一個 3 維張量，它將表示 10 個類別的 Capsule 嵌入向量。也就是說，該函數將構建完整的 CapsNet 架構，並輸出 DigitCaps 層最後得到的 10 個 16 維向量。
該主體函數主要調用了一個 slim 卷積 Capsule 層和一個 Capsule 層。slim 卷積 Capsule 層主要將輸入張量轉換為 Capsule 格式，即上圖的 PrimaryCaps 層。而為了連接卷積 Capsule 層和頂部的全連接 Capsule 層，卷積 Capsule 層的網格位置將與不同類型的 Capsule 維度相合併，並且 Capsule 將為嵌入向量學習不同的變換。下面我們會詳細討論構建以上網絡主體的兩個函數。
在主體模型的代碼中（capsule_model.py 第 54 行），第二個卷積層需要通過調用 Sara 等人定義的 conv_slim_capsule 函數實現，以下的代碼構建了原論文中的 PrimaryCaps 層，其中 input_tensor 為原圖片經過一次卷積後的特徵圖，並增加一個維度以作為一個 Capsule 單元包含神經元的個數（構成向量）。
capsule1 = layers.conv_slim_capsule(
input_tensor,
input_dim=1,
output_dim=self._hparams.num_prime_capsules,
layer_name='conv_capsule1',
num_routing=1,
input_atoms=256,
output_atoms=8,
stride=2,
kernel_size=9,
padding=self._hparams.padding,
leaky=self._hparams.leaky,)
在上面這些參數中，input_tensor 為五維張量即標準卷積的四維張量再加上一維 Capsule 單元數（capsule_model.py 第 194 行）。input_dim 為上一個 Capsule 層的單元數或維度，output_dim 為多個並行卷積操作後所得到的 Capsule 單元數或維度。input_atoms 為前一層 Capsule 單元的元素數，即一個 Capsule 單元包含的神經元數量，這裡 256 代表第一個卷積層所產生的 256 張特徵圖，而 output_atoms 表示當前層的 Capsule 單元元素數，這裡的 8 可以是代表 8 張 6×6×1×32 的特徵圖。其實小編認為我們可以如 Sara 等人的實現將 PrimaryCaps 層看成 32 個 Capsule 單元，每個單元包含 8 個標量神經元，或者將其看成 8 個 Capsule 單元，每個單元包含 32 個標量神經元，這兩種表示方法應該是等價的。剩下其它的參數就和標準卷積層所定義的參數意義一樣，所以讀者可以閱讀原始碼詳細地了解。
在 layers.py 文件中（268 行），Sara 等人定義了 conv_slim_capsule 函數以完成 PrimaryCaps 層的構建。該函數使用 slim 對給定五維的輸入張量執行二維卷積，輸入張量的維度一般為 [batch, input_dim, input_atoms, input_height, input_width]。然後該函數將使用動態路由算法精煉前面卷積運算的結果，並對每一個 Capsule 單元應用非線性 Squash 函數。
conv_slim_capsule 函數所輸出的激活值張量維度為 [batch, output_dim, output_atoms, out_height, out_width]。如果 Padding 選擇的是『SAME』，那麼輸出特徵圖的高和寬就與輸入張量的寬和高。我們注意到執行卷積操作具體的函數為前面定義的_depthwise_conv3d，該函數將返回經過 2 維卷積的 6 維張量。
_depthwise_conv3d 函數在給定一個 5 維輸入張量的情況下會執行 2 維卷積運算，輸入張量的維度與 conv_slim_capsule 函數的輸入相同。_depthwise_conv3d 函數會將輸入 5 維張量中 Batch 和 input_dim 的乘積作為 1 維而壓縮為 4 維張量，即壓縮輸入張量的第一維與第二維為一個維度。之所以需要壓縮為 4 維，是因為我們需要將其作為 tf.nn.conv2d 的輸入。在執行卷積後，我們需要重新將這 4 維張量分解為 6 維張量。即將第一維分解為 Batch 和 input_dim，將第二維分解為 output_dim 和 output_atom。
def _depthwise_conv3d(input_tensor,
kernel,
input_dim,
output_dim,
input_atoms=8,
output_atoms=8,
stride=2,
padding='SAME'):
with tf.name_scope('conv'):
input_shape = tf.shape(input_tensor)
_, _, _, in_height, in_width = input_tensor.get_shape()
# Reshape input_tensor to 4D by merging first two dimmensions.
# tf.nn.conv2d only accepts 4D tensors.
input_tensor_reshaped = tf.reshape(input_tensor, [
input_shape[0] * input_dim, input_atoms, input_shape[3], input_shape[4]
])
input_tensor_reshaped.set_shape((None, input_atoms, in_height.value,
in_width.value))
conv = tf.nn.conv2d(
input_tensor_reshaped,
kernel,
[1, 1, stride, stride],
padding=padding,
data_format='NCHW')
conv_shape = tf.shape(conv)
_, _, conv_height, conv_width = conv.get_shape()
# Reshape back to 6D by splitting first dimmension to batch and input_dim
# and splitting second dimmension to output_dim and output_atoms.
conv_reshaped = tf.reshape(conv, [
input_shape[0], input_dim, output_dim, output_atoms, conv_shape[2],
conv_shape[3]
])
conv_reshaped.set_shape((None, input_dim, output_dim, output_atoms,
conv_height.value, conv_width.value))
return conv_reshaped, conv_shape, input_shape
如上所示為 _depthwise_conv3d 函數，其中參數 input_dim、output_dim 和 input_atoms 等參數的意義與 conv_slim_capsule 函數一致。該函數會返回 6 維張量 [batch, input_dim, output_dim, output_atoms, out_height, out_width]、卷積後的維度和輸入張量的維度，並在 conv_slim_capsule 函數中做進一步處理，下面我們將回頭繼續討論構建 PrimaryCaps 層的函數。
如下定義了 conv_slim_capsule 函數，其層級的每一個 Capsule 單元都為卷積單元，它們在位置網格和不同的下層 Capsule 單元間共享卷積核權重。因此，該函數可訓練的變量為卷積核權重 [kernel_size, kernel_size, input_atoms, output_dim * output_atoms] 和偏置項 [output_dim, output_atoms]。二維卷積的輸出為一個單 Capsule 單元，其通道為 Capsule 單元（atoms）的數量。因此，conv_slim_capsule 函數可以構建在二維卷積層的頂部，其中該二維卷積的 num_routing=1、input_dim=1 和 input_atoms=conv_channels。
通過觀察定義卷積權重的 kernel 變量，我們可以了解該運算本質上就是執行上一層特徵圖數為 input_atoms、本層卷積核數為 output_dim * output_atoms 的卷積操作。在完成卷積運算後，Sara 的實現接著調用了一次前面定義的路由算法，這似乎與 naturomics 等人復現的代碼有一些不同，他們在實現卷積後會將卷積結果直接投入 Squash 非線性函數。當然，原論文似乎也沒有體現這一點，我們都以為只有在 DigitCaps 層才會進行動態路由過程。
在後面的 capsule 函數中（layers.py 第 138 行），Sara 確實又調用了一次動態路由算法，我們會在後面討論該函數。因為路由算法計算了 Squash 值，因此返回的激活值可作為 PrimaryCaps 層的輸出。
def conv_slim_capsule(input_tensor,
input_dim,
output_dim,
layer_name,
input_atoms=8,
output_atoms=8,
stride=2,
kernel_size=5,
padding='SAME',
**routing_args):
with tf.variable_scope(layer_name):
kernel = variables.weight_variable(shape=[
kernel_size, kernel_size, input_atoms, output_dim * output_atoms
])
biases = variables.bias_variable([output_dim, output_atoms, 1, 1])
votes, votes_shape, input_shape = _depthwise_conv3d(
input_tensor, kernel, input_dim, output_dim, input_atoms, output_atoms,
stride, padding)
with tf.name_scope('routing'):
logit_shape = tf.stack([
input_shape[0], input_dim, output_dim, votes_shape[2], votes_shape[3]
])
biases_replicated = tf.tile(biases,
[1, 1, votes_shape[2], votes_shape[3]])
activations = _update_routing(
votes=votes,
biases=biases_replicated,
logit_shape=logit_shape,
num_dims=6,
input_dim=input_dim,
output_dim=output_dim,
**routing_args)
return activations
在返回 PrimaryCaps 層的 5 維張量後，主體函數會將其轉換維 3 維張量並饋送到 capsule 函數，從而構建一個全連接 Capsule 層。在給定輸入張量 [batch, input_dim, input_atoms] 後，Capsule 層將執行以下操作：
對於每一個輸入 Capsule 單元，將它與權重變量相乘以得到線性組合的結果（函數中表示為 votes 變量），這一步將得到原論文所述的 u_j|i hat，即 u_j|i hat = W_ij * u_i。線性組合後的結果維度為 [batch, input_dim, output_dim, output_atoms]。
通過迭代地執行路由過程更新與精煉前面線性組合的結果，即原論文中的 s_j = ∑ c_ij * u_j|i hat，其中 c_ij = softmax(b_ij)。
最後使用 Squash 函數將每個 Capsule 單元的輸出壓縮到 L2 範數小於 1 的情況。
此外，當前層的每一個 Capsule 單元對前一層的 Capsule 單元都保留一個權重張量。因此在訓練中，capsule 函數的權重 [input_dim * num_in_atoms, output_dim * num_out_atoms] 和偏置項 [output_dim * num_out_atoms] 都是需要更新的參數。
如下展示了 capsule 函數的定義，其中 input_dim 同樣為前一層的 Capsule 單元數，input_atoms 同樣為前一層每個 Capsule 單元的元素數，其它參數的意義與上面幾個函數都差不多。該函數將輸出張量 [batch, output_dim, output_atoms]。
def capsule(input_tensor,
input_dim,
output_dim,
layer_name,
input_atoms=8,
output_atoms=8,
**routing_args):
with tf.variable_scope(layer_name):
# weights variable will hold the state of the weights for the layer
weights = variables.weight_variable(
[input_dim, input_atoms, output_dim * output_atoms])
biases = variables.bias_variable([output_dim, output_atoms])
with tf.name_scope('Wx_plus_b'):
# Depthwise matmul: [b, d, c] ** [d, c, o_c] = [b, d, o_c]
# To do this: tile input, do element-wise multiplication and reduce
# sum over input_atoms dimmension.
input_tiled = tf.tile(
tf.expand_dims(input_tensor, -1),
[1, 1, 1, output_dim * output_atoms])
votes = tf.reduce_sum(input_tiled * weights, axis=2)
votes_reshaped = tf.reshape(votes,
[-1, input_dim, output_dim, output_atoms])
with tf.name_scope('routing'):
input_shape = tf.shape(input_tensor)
logit_shape = tf.stack([input_shape[0], input_dim, output_dim])
activations = _update_routing(
votes=votes_reshaped,
biases=biases,
logit_shape=logit_shape,
num_dims=4,
input_dim=input_dim,
output_dim=output_dim,
**routing_args)
return activations
如上所示，tf.expand_dims 會將輸入張量在最後擴充一維，而 tf.tile 會將擴充後的 4 維張量在最後一維複製 output_dim * output_atoms 次。在執行逐元素的乘積後，沿著第三維 input_atoms 對乘積結果求和。求和後的張量可作為最後的線性組合結果而投入路由算法中進行迭代更新，返回的張量即 DigitCaps 層最終輸出的 10 個 16 維的向量，每個向量編碼並表徵著 10 類手寫數字。
在調用完 capsule 函數後，整個_build_capsule 函數所構建的 CapsNet 架構就完成了。原論文使用了 Margin loss 來衡量這 10 個輸出向量預測類別的準確度，而後面也可以使用全連接網絡將這 10 個向量重構為不同手寫數字的圖像，並使用歐幾裡得距離度量重構損失。
在論文解讀中，我們已經了解 DigitCaps 層輸出向量的長度即某個類別的概率，那麼我們該如何構建損失函數，並根據該損失函數迭代地更新整個網絡？前面我們耦合係數 c_ij 是通過一致性 Routing 進行更新的，他並不需要根據損失函數更新，但整個網絡其它的卷積參數和 Capsule 內的 W_ij 都需要根據損失函數進行更新。一般我們就可以對損失函數直接使用標準的反向傳播更新這些參數，而在原論文中，作者採用了 SVM 中常用的 Margin loss，該損失函數的表達式為：
其中 c 是分類類別，T_c 為分類的指示函數（c 存在為 1，c 不存在為 0），m+ 為上邊界，m- 為下邊界。此外，v_c 的模即向量的 L2 距離。
因為實例化向量的長度來表示 Capsule 要表徵的實體是否存在，所以若且唯若圖片裡出現屬於類別 k 的手寫數字時，我們希望類別 k 的最頂層 Capsule 的輸出向量長度很大（在本論文 CapsNet 中為 DigitCaps 層的輸出）。為了允許一張圖裡有多個數字，我們對每一個表徵數字 k 的 Capsule 分別給出單獨的 Margin loss。
以下_margin_loss 定義了 CapsNet 的損失函數，它會懲罰每個輸入分對數偏離邊緣的程度。如函數說明所示，該函數將衡量每一個錯誤分對數對於邊緣的距離。對於負的分對數來說，邊緣為 0.1，對於正的分對數，邊緣為 0.9。若我們同時對這兩個邊緣先減去 0.5，那麼當前的邊緣將會都變為 0.4。
def _margin_loss(labels, raw_logits, margin=0.4, downweight=0.5):
"""Penalizes deviations from margin for each logit.
Each wrong logit costs its distance to margin. For negative logits margin is
0.1 and for positives it is 0.9. First subtract 0.5 from all logits. Now
margin is 0.4 from each side.
Args:
labels: tensor, one hot encoding of ground truth.
raw_logits: tensor, model predictions in range [0, 1]
margin: scalar, the margin after subtracting 0.5 from raw_logits.
downweight: scalar, the factor for negative cost.
Returns:
A tensor with cost for each data point of shape [batch_size].
"""
logits = raw_logits - 0.5
positive_cost = labels * tf.cast(tf.less(logits, margin),
tf.float32) * tf.pow(logits - margin, 2)
negative_cost = (1 - labels) * tf.cast(
tf.greater(logits, -margin), tf.float32) * tf.pow(logits + margin, 2)
return0.5 * positive_cost + downweight * 0.5 * negative_cost
如上所示，margin 為 raw_logits 減去 0.5 的邊緣，而 downweight 負成本的因素。
以上是 CapsNet 的主體代碼，也是整個 Capsule 的核心。Sara 開源的實現還有很多重要的代碼與函數，我們也將繼續探討與思考其中具體的過程，尤其是理解_leaky_routing 的作用。我們也希望有讀者和我們一起解析與分析 Sara 的實現，並探討其最終實現的基線結果。
以下是 Sara 實現的說明，其使用的是 Python 2.7，不過 naturomics 等研究者已經在 GitHub 上修改為了 Python 3。讀者可進一步測試與實現它們：
要求：
TensorFlow
NumPy
GPU
運行以下命令以測試配置是否正確：
python layers_test.py
若我們下載了 Sara 等人提供的 MNIST 數據集和預訓練模型，並把它們放入$DATA_DIR/和$CKPT_DIR/ 目錄下。那麼我們可以運行以下命令而快速獲得 CapsNet 在 MNIST 上的測試結果。
數據集：https://storage.googleapis.com/capsule_toronto/mnist_data.tar.gz
模型：https://storage.googleapis.com/capsule_toronto/mnist_checkpoints.tar.gz
python experiment.py --data_dir=$DATA_DIR/mnist_data/ --train=false \
--summary_dir=/tmp/ --checkpoint=$CKPT_DIR/mnist_checkpoint/model.ckpt-1
下載並抽取二進位版的 cifar10 數據集到 $DATA_DIR/，cifar 10 預訓練模型到 $CKPT_DIR/。
數據集：https://www.cs.toronto.edu/~kriz/cifar.html
模型：https://storage.googleapis.com/capsule_toronto/cifar_checkpoints.tar.gz
python experiment.py --data_dir=$DATA_DIR --train=false --dataset=cifar10 \
--hparams_override=num_prime_capsules=64,padding=SAME,leaky=true,remake=false \
--summary_dir=/tmp/ --checkpoint=$CKPT_DIR/cifar/cifar{}/model.ckpt-600000 \
--num_trials=7
Sample CIFAR10 訓練命令：
python experiment.py --data_dir=$DATA_DIR --dataset=cifar10 --max_steps=600000\
--hparams_override=num_prime_capsules=64,padding=SAME,leaky=true,remake=false \
--summary_dir=/tmp/
Sample MNIST 全部訓練命令：
訓練-驗證傳播模式： --validate=true
在更多 GPU 上訓練：—num_gpus=NUM_GPUS
python experiment.py --data_dir=$DATA_DIR/mnist_data/ --max_steps=300000\
--summary_dir=/tmp/attempt0/
Sample MNIST 基線訓練命令：
python experiment.py --data_dir=$DATA_DIR/mnist_data/ --max_steps=300000\
--summary_dir=/tmp/attempt1/ --model=baseline
該實現的詳細使用方法請查閱 README 文件，我們以上只簡要介紹了一部分。我們也嘗試著實現 Sara 開源的模型，最開始是 xrange 等 Python 2.7 的函數會報錯，在改為 Python 3 後運行仍然會有問題。所以我們更希望有讀者能完成該模型的測試，並向大家展示 Dynamic Routing Between Capsules 論文所實現的結果。