CapsNet入門系列番外：基於TensorFlow實現膠囊網路

【編者按】全棧開發者Debarko De簡明扼要地介紹了膠囊網路的概念，同時給出了基於numpy和TensorFlow的膠囊網路實現。

什麼是膠囊網路？什麼是膠囊？膠囊網路比卷積神經網路（CNN）更好嗎？本文將討論這些關於Hinton提出的CapsNet（膠囊網路）的話題。

注意，本文討論的不是製藥學中的膠囊，而是神經網路和機器學習中的膠囊。

閱讀本文前，你需要對CNN有基本的了解，否則建議你先看下我之前寫的Deep Learning for Noobs。下面我將稍稍溫習下與本文相關的CNN的知識，這樣你能更容易理解下文CNN與CapsNet的對比。閑話就不多說了，讓我們開始吧。

基本上，CNN是堆疊一大堆神經元構成的系統。CNN很擅長處理圖像分類問題。讓神經網路映射一張圖像的所有像素，從算力上來講，太昂貴了。而卷積在保留數據本質的前提下大大簡化了計算。基本上，卷積是一大堆矩陣乘法，再將乘積相加。

圖像傳入網路後，一組核或過濾器掃描圖像並進行卷積操作，從而創建特徵映射。這些特徵接著傳給之後的激活層和池化層。取決於網路的層數，這一組合可能反覆堆疊。激活網路給網路帶來了一些非線性（比如ReLU）。池化（比如最大池化）有助於減少訓練時間。池化的想法是為每個子區域創建「概要」。同時池化也提供了一些目標檢測的位置和平移不變性。網路的最後是一個分類器，比如softmax分類器，分類器返回類別。訓練基於對應標註數據的錯誤進行反向傳播。在這一步驟中，非線性有助於解決梯度衰減問題。

CNN有什麼問題？

在分類非常接近數據集的圖像時，CNN表現極為出色。但CNN在顛倒、傾斜或其他朝向不同的圖像上表現很差。訓練時添加同一圖像的不同變體可以解決這一問題。在CNN中，每層對圖像的理解粒度更粗。舉個例子，假設你試圖分類船和馬。最內層（第一層）理解細小的曲線和邊緣。第二層可能理解直線或小形狀，例如船的桅杆和整個尾巴的曲線。更高層開始理解更複雜的形狀，例如整條尾巴或船體。最後層嘗試總覽全圖（例如整條船或整匹馬）。我們在每層之後使用池化，以便在合理的時間內完成計算，但本質上池化同時丟失了位置信息。

池化有助於建立位置不變性。否則CNN將只能擬合非常接近訓練集的圖像或數據。這樣的不變性同時導致具備船的部件但順序錯誤的圖像被誤認為船。所以系統會把上圖右側的圖像誤認為船，而人類則能很清楚地觀察到兩者的區別。另外，池化也有助於建立比例不變性。

池化本來是用來引入位置、朝向、比例不變性的，然而這一方法非常粗糙。事實上池化加入了各種位置不變性，以致將部件順序錯誤的圖像也誤認為船了。我們需要的不是不變性，而是等價性。不變性使CNN可以容忍視角中的小變動，而等價性使CNN理解朝向和比例變動，並相應地適應圖像，從而不損失圖像的空間位置信息。CNN會減少自身尺寸以檢測較小的船。這導向了最近發展出的膠囊網路。

什麼是膠囊網路？

Sara Sabour、Nicholas Frost、Geoffrey Hinton在2017年10月發表了論文Dynamic Routing Between Capsules。當深度學習的祖父之一Geoffrey Hinton發表了一篇論文，這論文註定會是一項重大突破。整個深度學習社區都為此瘋狂。這篇論文討論了膠囊、膠囊網路以及在MNIST上的試驗。MNIST是已標註的手寫數字圖像數據集。相比當前最先進的CNN，膠囊網路在重疊數字上的表現明顯提升。論文的作者提出人腦有一個稱為「膠囊」的模塊，這些膠囊特別擅長處理不同的視覺刺激，以及編碼位姿（位置、尺寸、朝向）、變形、速度、反射率、色調、紋理等信息。大腦肯定具備「路由」低層視覺信息至最擅長處理該信息的卷囊的機制。

膠囊是一組嵌套的神經網路層。在通常的神經網路中，你不斷添加更多層。在膠囊網路中，你會在單個網路層中加入更多的層。換句話說，在一個神經網路層中嵌套另一個。膠囊中的神經元的狀態刻畫了圖像中的一個實體的上述屬性。膠囊輸出一個表示實體存在性的向量。向量的朝向表示實體的屬性。向量發送至神經網路中所有可能的親本膠囊。膠囊可以為每個可能的親本計算出一個預測向量，預測向量是通過將自身權重乘以權重矩陣得出的。預測向量乘積標量最大的親本膠囊的聯繫將增強，而剩下的親本膠囊聯繫將減弱。這一基於合意的路由方法比諸如最大池化之類的現有機制更優越。最大池化路由基於低層網路檢測出的最強烈的特徵。動態路由之外，膠囊網路給膠囊加上了squash函數。squash屬於非線性函數。與CNN給每個網路層添加squash函數不同，膠囊網路給每組嵌套的網路層添加squash函數，從而將squash函數應用到每個膠囊的輸出向量。

論文引入了一個全新的squash函數（見上圖）。ReLU及類似的非線性函數在單個神經元上表現良好，不過論文發現在膠囊上squash函數表現最好。squash函數壓縮膠囊的輸出向量的長度：當向量較小時，結果為0；當向量較大時，結果為1。動態路由增加了一些額外的運算開銷，但毫無疑問帶來了優勢。

當然我們也要注意，這篇論文剛發不久，膠囊的概念還沒有經過全面的測試。它在MNIST數據集上表現良好，但在其他更多種類、更大的數據集上的表現還有待證明。在論文發布的幾天之內，就有人提出一些意見。

當前的膠囊網路實現還有改進的空間。不過別忘了Hinton的論文一開始就提到了：

這篇論文的目標不是探索整個空間，而是簡單地展示一個相當直接的實現表現良好，同時動態路由有所裨益。

好了，我們已經談了夠多理論了。讓我們找點樂子，構建一個膠囊網路。我將引領你閱讀一些為MNIST數據配置一個膠囊網路的代碼。我會在代碼里加上注釋，這樣你可以逐行理解這些代碼是如何工作的。本文將包括兩個重要的代碼片段。其餘代碼見GitHub倉庫：

# 只依賴numpy和tensorflowimport numpy as npimport tensorflow as tffrom config import cfg# 定義卷積膠囊類，該類由多個神經網路層組成#class CapsConv(object): 膠囊層 參數： input：一個4維張量。 num_units：整數，膠囊的輸出向量的長度。 with_routing：布爾值，該膠囊路由經過低層膠囊。 num_outputs：該層中的膠囊數目。 返回： 一個4維張量。 def __init__(self, num_units, with_routing=True): self.num_units = num_units self.with_routing = with_routing def __call__(self, input, num_outputs, kernel_size=None, stride=None): self.num_outputs = num_outputs self.kernel_size = kernel_size self.stride = stride if not self.with_routing: # 主膠囊（PrimaryCaps）層 # 輸入： [batch_size, 20, 20, 256] assert input.get_shape() == [cfg.batch_size, 20, 20, 256] capsules = [] for i in range(self.num_units): # 每個膠囊i： [batch_size, 6, 6, 32] with tf.variable_scope(ConvUnit_ + str(i)): caps_i = tf.contrib.layers.conv2d(input, self.num_outputs, self.kernel_size, self.stride, padding="VALID") caps_i = tf.reshape(caps_i, shape=(cfg.batch_size, -1, 1, 1)) capsules.append(caps_i) assert capsules[0].get_shape() == [cfg.batch_size, 1152, 1, 1] # [batch_size, 1152, 8, 1] capsules = tf.concat(capsules, axis=2) capsules = squash(capsules) assert capsules.get_shape() == [cfg.batch_size, 1152, 8, 1] else: # 數字膠囊（DigitCaps）層 # reshape輸入至：[batch_size, 1152, 8, 1] self.input = tf.reshape(input, shape=(cfg.batch_size, 1152, 8, 1)) # b_IJ: [1, num_caps_l, num_caps_l_plus_1, 1] b_IJ = tf.zeros(shape=[1, 1152, 10, 1], dtype=np.float32) capsules = [] for j in range(self.num_outputs): with tf.variable_scope(caps_ + str(j)): caps_j, b_IJ = capsule(input, b_IJ, j) capsules.append(caps_j) # 返回一個張量：[atch_size, 10, 16, 1] capsules = tf.concat(capsules, axis=1) assert capsules.get_shape() == [cfg.batch_size, 10, 16, 1] return(capsules)def capsule(input, b_IJ, idx_j): 層l+1中的單個膠囊的路由演算法。 參數： input： 張量 [batch_size, num_caps_l=1152, length(u_i)=8, 1] num_caps_l為l層的膠囊數 返回： 張量 [batch_size, 1, length(v_j)=16, 1] 表示 l+1層的膠囊j的輸出向量`v_j` 注意： u_i表示l層膠囊i的輸出向量， v_j則表示l+1層膠囊j的輸出向量 with tf.variable_scope(routing): w_initializer = np.random.normal(size=[1, 1152, 8, 16], scale=0.01) W_Ij = tf.Variable(w_initializer, dtype=tf.float32) # 重複batch_size次W_Ij：[batch_size, 1152, 8, 16] W_Ij = tf.tile(W_Ij, [cfg.batch_size, 1, 1, 1]) # 計算 u_hat # [8, 16].T x [8, 1] => [16, 1] => [batch_size, 1152, 16, 1] u_hat = tf.matmul(W_Ij, input, transpose_a=True) assert u_hat.get_shape() == [cfg.batch_size, 1152, 16, 1] shape = b_IJ.get_shape().as_list() size_splits = [idx_j, 1, shape[2] - idx_j - 1] for r_iter in range(cfg.iter_routing): # 第4行： # [1, 1152, 10, 1] c_IJ = tf.nn.softmax(b_IJ, dim=2) assert c_IJ.get_shape() == [1, 1152, 10, 1] # 第5行： # 在第三維使用c_I加權u_hat # 接著在第二維累加，得到[batch_size, 1, 16, 1] b_Il, b_Ij, b_Ir = tf.split(b_IJ, size_splits, axis=2) c_Il, c_Ij, b_Ir = tf.split(c_IJ, size_splits, axis=2) assert c_Ij.get_shape() == [1, 1152, 1, 1] s_j = tf.multiply(c_Ij, u_hat) s_j = tf.reduce_sum(tf.multiply(c_Ij, u_hat), axis=1, keep_dims=True) assert s_j.get_shape() == [cfg.batch_size, 1, 16, 1] # 第六行: # 使用上文提及的squash函數，得到：[batch_size, 1, 16, 1] v_j = squash(s_j) assert s_j.get_shape() == [cfg.batch_size, 1, 16, 1] # 第7行: # 平鋪v_j，由[batch_size ,1, 16, 1] 至[batch_size, 1152, 16, 1] # [16, 1].T x [16, 1] => [1, 1] # 接著在batch_size維度遞歸運算均值，得到 [1, 1152, 1, 1] v_j_tiled = tf.tile(v_j, [1, 1152, 1, 1]) u_produce_v = tf.matmul(u_hat, v_j_tiled, transpose_a=True) assert u_produce_v.get_shape() == [cfg.batch_size, 1152, 1, 1] b_Ij += tf.reduce_sum(u_produce_v, axis=0, keep_dims=True) b_IJ = tf.concat([b_Il, b_Ij, b_Ir], axis=2) return(v_j, b_IJ)def squash(vector): 壓縮函數 參數： vector：一個4維張量 [batch_size, num_caps, vec_len, 1], 返回： 一個和vector形狀相同的4維張量， 但第3維和第4維經過壓縮 vec_abs = tf.sqrt(tf.reduce_sum(tf.square(vector))) # 一個標量 scalar_factor = tf.square(vec_abs) / (1 + tf.square(vec_abs)) vec_squashed = scalar_factor * tf.divide(vector, vec_abs) # 對應元素相乘 return(vec_squashed)

上面是一整個膠囊層。堆疊膠囊層以構成膠囊網路。

import tensorflow as tffrom config import cfgfrom utils import get_batch_datafrom capsLayer import CapsConvclass CapsNet(object): def __init__(self, is_training=True): self.graph = tf.Graph() with self.graph.as_default(): if is_training: self.X, self.Y = get_batch_data() self.build_arch() self.loss() # t_vars = tf.trainable_variables() self.optimizer = tf.train.AdamOptimizer() self.global_step = tf.Variable(0, name=global_step, trainable=False) self.train_op = self.optimizer.minimize(self.total_loss, global_step=self.global_step) # var_list=t_vars) else: self.X = tf.placeholder(tf.float32, shape=(cfg.batch_size, 28, 28, 1)) self.build_arch() tf.logging.info(Seting up the main structure) def build_arch(self): with tf.variable_scope(Conv1_layer): # Conv1（第一卷積層）, [batch_size, 20, 20, 256] conv1 = tf.contrib.layers.conv2d(self.X, num_outputs=256, kernel_size=9, stride=1, padding=VALID) assert conv1.get_shape() == [cfg.batch_size, 20, 20, 256] # TODO: 將CapsConv類重寫為函數， # capsLay函數應該封裝為兩個函數， # 一個類似conv2d，另一個為TensorFlow的fully_connected（全連接）。 # 主膠囊，[batch_size, 1152, 8, 1] with tf.variable_scope(PrimaryCaps_layer): primaryCaps = CapsConv(num_units=8, with_routing=False) caps1 = primaryCaps(conv1, num_outputs=32, kernel_size=9, stride=2) assert caps1.get_shape() == [cfg.batch_size, 1152, 8, 1] # 數字膠囊層，[batch_size, 10, 16, 1] with tf.variable_scope(DigitCaps_layer): digitCaps = CapsConv(num_units=16, with_routing=True) self.caps2 = digitCaps(caps1, num_outputs=10) # 前文示意圖中的編碼器結構 # 1. 掩碼： with tf.variable_scope(Masking): # a). 計算 ||v_c||，接著計算softmax(||v_c||) # [batch_size, 10, 16, 1] => [batch_size, 10, 1, 1] self.v_length = tf.sqrt(tf.reduce_sum(tf.square(self.caps2), axis=2, keep_dims=True)) self.softmax_v = tf.nn.softmax(self.v_length, dim=1) assert self.softmax_v.get_shape() == [cfg.batch_size, 10, 1, 1] # b). 選取10個膠囊的最大softmax值的索引 # [batch_size, 10, 1, 1] => [batch_size] (index) argmax_idx = tf.argmax(self.softmax_v, axis=1, output_type=tf.int32) assert argmax_idx.get_shape() == [cfg.batch_size, 1, 1] # c). 索引 # 由於我們是三維生物， # 理解argmax_idx的索引過程並不容易 masked_v = [] argmax_idx = tf.reshape(argmax_idx, shape=(cfg.batch_size, )) for batch_size in range(cfg.batch_size): v = self.caps2[batch_size][argmax_idx[batch_size], :] masked_v.append(tf.reshape(v, shape=(1, 1, 16, 1))) self.masked_v = tf.concat(masked_v, axis=0) assert self.masked_v.get_shape() == [cfg.batch_size, 1, 16, 1] # 2. 使用3個全連接層重建MNIST圖像 # [batch_size, 1, 16, 1] => [batch_size, 16] => [batch_size, 512] with tf.variable_scope(Decoder): vector_j = tf.reshape(self.masked_v, shape=(cfg.batch_size, -1)) fc1 = tf.contrib.layers.fully_connected(vector_j, num_outputs=512) assert fc1.get_shape() == [cfg.batch_size, 512] fc2 = tf.contrib.layers.fully_connected(fc1, num_outputs=1024) assert fc2.get_shape() == [cfg.batch_size, 1024] self.decoded = tf.contrib.layers.fully_connected(fc2, num_outputs=784, activation_fn=tf.sigmoid) def loss(self): # 1. 邊際損失 # [batch_size, 10, 1, 1] # max_l = max(0, m_plus-||v_c||)^2 max_l = tf.square(tf.maximum(0., cfg.m_plus - self.v_length)) # max_r = max(0, ||v_c||-m_minus)^2 max_r = tf.square(tf.maximum(0., self.v_length - cfg.m_minus)) assert max_l.get_shape() == [cfg.batch_size, 10, 1, 1] # reshape: [batch_size, 10, 1, 1] => [batch_size, 10] max_l = tf.reshape(max_l, shape=(cfg.batch_size, -1)) max_r = tf.reshape(max_r, shape=(cfg.batch_size, -1)) # 計算 T_c: [batch_size, 10] # T_c = Y，我的理解沒錯吧？試試看。 T_c = self.Y # [batch_size, 10]，對應元素相乘 L_c = T_c * max_l + cfg.lambda_val * (1 - T_c) * max_r self.margin_loss = tf.reduce_mean(tf.reduce_sum(L_c, axis=1)) # 2. 重建損失 orgin = tf.reshape(self.X, shape=(cfg.batch_size, -1)) squared = tf.square(self.decoded - orgin) self.reconstruction_err = tf.reduce_mean(squared) # 3. 總損失 self.total_loss = self.margin_loss + 0.0005 * self.reconstruction_err # 總結 tf.summary.scalar(margin_loss, self.margin_loss) tf.summary.scalar(reconstruction_loss, self.reconstruction_err) tf.summary.scalar(total_loss, self.total_loss) recon_img = tf.reshape(self.decoded, shape=(cfg.batch_size, 28, 28, 1)) tf.summary.image(reconstruction_img, recon_img)self.merged_sum = tf.summary.merge_all()

完整代碼（含訓練和驗證模型）見此。代碼以Apache 2.0許可發布。我參考了naturomics的代碼。

總結

我們介紹了膠囊網路的概念以及如何實現膠囊網路。我們嘗試理解膠囊是高層的嵌套神經網路層。我們也查看了膠囊網路是如何交付朝向和其他不變性的——對圖像中的每個實體而言，保持空間配置的等價性。我確信存在一些本文沒有回答的問題，其中最主要的大概是膠囊及其最佳實現。不過本文是解釋這一主題的初步嘗試。如果你有任何疑問，請評論。我會盡我所能回答。

Siraj Raval及其演講對本文影響很大。請在Twitter上分享本文。在twitter關注我以便在未來獲取更新信息。

原文 What is a CapsNet or Capsule Network?

感謝原作者Debarko De授權論智編譯，未經授權禁止轉載。

推薦閱讀：