CapsNet入門系列番外：基于TensorFlow實現(xiàn)膠囊網(wǎng)絡(luò)

編者按：全棧開發(fā)者Debarko De簡明扼要地介紹了膠囊網(wǎng)絡(luò)的概念，同時給出了基于numpy和TensorFlow的膠囊網(wǎng)絡(luò)實現(xiàn)。

什么是膠囊網(wǎng)絡(luò)？什么是膠囊？膠囊網(wǎng)絡(luò)比卷積神經(jīng)網(wǎng)絡(luò)（CNN）更好嗎？本文將討論這些關(guān)于Hinton提出的CapsNet（膠囊網(wǎng)絡(luò)）的話題。

注意，本文討論的不是制藥學(xué)中的膠囊，而是神經(jīng)網(wǎng)絡(luò)和機器學(xué)習(xí)中的膠囊。

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

基本上，CNN是堆疊一大堆神經(jīng)元構(gòu)成的系統(tǒng)。CNN很擅長處理圖像分類問題。讓神經(jīng)網(wǎng)絡(luò)映射一張圖像的所有像素，從算力上來講，太昂貴了。而卷積在保留數(shù)據(jù)本質(zhì)的前提下大大簡化了計算?；旧希矸e是一大堆矩陣乘法，再將乘積相加。

圖像傳入網(wǎng)絡(luò)后，一組核或過濾器掃描圖像并進(jìn)行卷積操作，從而創(chuàng)建特征映射。這些特征接著傳給之后的激活層和池化層。取決于網(wǎng)絡(luò)的層數(shù)，這一組合可能反復(fù)堆疊。激活網(wǎng)絡(luò)給網(wǎng)絡(luò)帶來了一些非線性（比如ReLU）。池化（比如最大池化）有助于減少訓(xùn)練時間。池化的想法是為每個子區(qū)域創(chuàng)建“概要”。同時池化也提供了一些目標(biāo)檢測的位置和平移不變性。網(wǎng)絡(luò)的最后是一個分類器，比如softmax分類器，分類器返回類別。訓(xùn)練基于對應(yīng)標(biāo)注數(shù)據(jù)的錯誤進(jìn)行反向傳播。在這一步驟中，非線性有助于解決梯度衰減問題。

CNN有什么問題？

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

畸形變換

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

比例不變性

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

什么是膠囊網(wǎng)絡(luò)？

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

膠囊網(wǎng)絡(luò)架構(gòu)

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

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

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

當(dāng)前的膠囊網(wǎng)絡(luò)實現(xiàn)還有改進(jìn)的空間。不過別忘了Hinton的論文一開始就提到了：

這篇論文的目標(biāo)不是探索整個空間，而是簡單地展示一個相當(dāng)直接的實現(xiàn)表現(xiàn)良好，同時動態(tài)路由有所裨益。

好了，我們已經(jīng)談了夠多理論了。讓我們找點樂子，構(gòu)建一個膠囊網(wǎng)絡(luò)。我將引領(lǐng)你閱讀一些為MNIST數(shù)據(jù)配置一個膠囊網(wǎng)絡(luò)的代碼。我會在代碼里加上注釋，這樣你可以逐行理解這些代碼是如何工作的。本文將包括兩個重要的代碼片段。其余代碼見GitHub倉庫：

# 只依賴numpy和tensorflow

import numpy as np

import tensorflow as tf

from config import cfg

# 定義卷積膠囊類，該類由多個神經(jīng)網(wǎng)絡(luò)層組成

#

classCapsConv(object):

''' 膠囊層

參數(shù)：

input：一個4維張量。

num_units：整數(shù)，膠囊的輸出向量的長度。

with_routing：布爾值，該膠囊路由經(jīng)過低層膠囊。

num_outputs：該層中的膠囊數(shù)目。

返回：

一個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

ifnot 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:

# 數(shù)字膠囊（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中的單個膠囊的路由算法。

參數(shù)：

input： 張量 [batch_size, num_caps_l=1152, length(u_i)=8, 1]

num_caps_l為l層的膠囊數(shù)

返回：

張量 [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)

# 重復(fù)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加權(quán)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函數(shù)，得到：[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):

'''壓縮函數(shù)

參數(shù)：

vector：一個4維張量 [batch_size, num_caps, vec_len, 1],

返回：

一個和vector形狀相同的4維張量，

但第3維和第4維經(jīng)過壓縮

'''

vec_abs = tf.sqrt(tf.reduce_sum(tf.square(vector))) # 一個標(biāo)量

scalar_factor = tf.square(vec_abs) / (1 + tf.square(vec_abs))

vec_squashed = scalar_factor * tf.divide(vector, vec_abs) # 對應(yīng)元素相乘

return(vec_squashed)

上面是一整個膠囊層。堆疊膠囊層以構(gòu)成膠囊網(wǎng)絡(luò)。

import tensorflow as tf

from config import cfg

from utils import get_batch_data

from capsLayer importCapsConv

classCapsNet(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'類重寫為函數(shù)，

# capsLay函數(shù)應(yīng)該封裝為兩個函數(shù)，

# 一個類似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]

# 數(shù)字膠囊層，[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)

# 前文示意圖中的編碼器結(jié)構(gòu)

# 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]，對應(yīng)元素相乘

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

# 總結(jié)

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()

完整代碼（含訓(xùn)練和驗證模型）見此（https://github.com/debarko/CapsNet-Tensorflow）。代碼以Apache 2.0許可發(fā)布。我參考了naturomics的代碼（https://github.com/naturomics）。

總結(jié)

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

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