僅17 KB、一萬個權重的微型風格遷移網路！
jamesonatfritz 想解決神經網路的過引數化問題，想要建立體積小但效能優的神經網路。他所試驗的第一個任務便是藝術風格遷移。
-
GitHub 連結：https://github.com/fritzlabs/fritz-style-transfer
現在有很多用來訓練藝術風格遷移模型的現成工具，還有上千種開源實現。其中的多數工具利用 Johnson 等人在《Perceptual Losses for Real-Time Style Transfer and Super-Resolution》中提出的網路架構的變體來實現快速、前饋的風格化。因此，多數遷移模型的大小是 7MB。對於你的應用來說，這個負擔並非不可承受，但也並非無足輕重。
研究表明，神經網路的體積通常遠遠大於所需，數百萬的權重中有很多並不重要。因此作者創造了一個體積大大縮小的可靠風格遷移模型：一個只有 11686 個訓練權重的 17KB神經網路。
左：原圖；中：來自上述 17KB 模型的風格化影象；右：來自 7MB 模型的風格化影象。
快速概覽：
原始模型：
-
大小：7MB
-
權重數：1.7M
-
在 iPhone X 上的速度：18 FPS
小模型：
-
大小：17KB
-
權重數：11,868
-
在 iPhone X 上的速度：29 FPS
如何縮小風格遷移模型
作者主要使用了兩種技術，而且都可以泛化到其他模型：
1. 大刀闊斧地修剪層和權重；
2. 通過量化將 32 位浮點權重轉換為 8 位整型
修剪策略
卷積神經網路通常包含數百萬甚至上億個需要在訓練階段進行調整的權重。通常來講，權重越多準確率越高。但這種增加權重提高準確率的做法非常低效。谷歌MobileNetV2 的 stock 配置具有 347 萬個權重，記憶體佔用達 16MB。InceptionV3 架構大小約為前者的 6 倍，具備 2400 萬個權重，記憶體佔用達 92MB。儘管多了 2000 多萬個權重，但 InceptionV3 在ImageNet上的 top-1 分類準確率只比MobileNetV2 高出 7 個百分點（80% vs 73%）。
因此，我們可以假設神經網路中的多數權重沒有那麼重要並將其移除。但重點是怎麼做呢？我們可以選擇在三個層面進行修剪：單個權重、層、塊。
權重層面：假設某個神經網路上的多數（>95%）權重都沒有什麼用。如果能找出那些對準確率有影響的權重，就可以將其留下並將其他移除。
層層面：每個層中都包含一些權重。例如，2D 卷積層具有一個權重張量，即卷積核，使用者可以定義其寬度、高度和深度。縮小卷積核可以減小整個網路的大小。
塊層面：多個層通常可以結合成可重複利用的子圖，即塊。以 ResNet 為例，它的名字來源於重複 10-50 次的「殘差塊」。在塊層面進行修剪可以移除多個層，從而一次性移除多個引數。
在實踐中，稀疏張量運算沒有很好的實現，因此權重層面的修剪沒有多大價值。那麼就只剩下層和塊層面的修剪了。
實踐中的修剪
作者使用的層修剪技術是引入 width multiplier 作為超引數。width multiplier 最初由谷歌在其論文《MobileNets: Efficient ConvolutionalNeural Networks for Mobile Vision》中提出，非常簡單、高效。
width multiplier 利用一個恆定係數調整每個卷積層中的卷積核數量。對於給定的層及 width multiplier alpha，卷積核數量 F 變為 alpha * F。
有了這個超引數，我們就可以生成一系列架構相同但權重數不同的網路。訓練每種配置，就可以在模型速度、大小及準確率之間做出權衡。
下面是作者模仿 Johnson 等人在《The Lottery Ticket Hypothesis: Finding Sparse, TrainableNeural Networks》提出的網路架構構建快速風格遷移模型的方法，不同之處在於添加了 width multiplier 作為超引數。
@classmethod def build( cls, image_size, alpha=1.0, input_tensor=None, checkpoint_file=None): """Build a Transfer Network Model using keras' functional API. Args: image_size - the size of the input and output image (H, W) alpha - a width <mark data-id="2e982b73-88e2-41e8-a430-f7ae5a9af4bf" data-type="technologies">parameter</mark> to scale the number of channels by Returns: model: a keras model object """ x = keras.layers.Input( shape=(image_size[0], image_size[1], 3), tensor=input_tensor) out = cls._convolution(x, int(alpha * 32), 9, strides=1) out = cls._convolution(out, int(alpha * 64), 3, strides=2) out = cls._convolution(out, int(alpha * 128), 3, strides=2) out = cls._residual_block(out, int(alpha * 128)) out = cls._residual_block(out, int(alpha * 128)) out = cls._residual_block(out, int(alpha * 128)) out = cls._residual_block(out, int(alpha * 128)) out = cls._residual_block(out, int(alpha * 128)) out = cls._upsample(out, int(alpha * 64), 3) out = cls._upsample(out, int(alpha * 32), 3) out = cls._convolution(out, 3, 9, relu=False, padding='same') # Restrict outputs of pixel values to -1 and 1. out = keras.layers.Activation('tanh')(out) # Deprocess the image into valid image data. Note we'll need to define # a custom layer for this in Core ML as well. out = layers.DeprocessStylizedImage()(out) model = keras.models.Model(inputs=x, outputs=out)
注意，模型構建器類的其餘部分沒有顯示。
當 alpha=1.0 時，得到的網路包含 170 萬個權重。當 alpha=0.5 時，得到的網路僅有 424,102 個權重。
你可以構建一些寬度引數很小的網路，但是也有相當多的重複塊。作者決定修剪掉一些，但實際操作後卻發現不能移除太多。即使引數量保持不變，較深的網路能夠產生更好的結果。作者最終刪除了五個殘差塊中的兩個，並將每層的預設濾波器數量減少至 32 個。得到的微型網路如下所示：
@classmethod
def build(
cls,
image_size,
alpha=1.0,
input_tensor=None,
checkpoint_file=None):
"""Build a Small Transfer Network Model using keras' functional API.
This architecture removes some blocks of layers and reduces the size
of convolutions to save on computation.
Args:
image_size - the size of the input and output image (H, W)
alpha - a width <mark data-id="2e982b73-88e2-41e8-a430-f7ae5a9af4bf" data-type="technologies">parameter</mark> to scale the number of channels by
Returns:
model: a keras model object
"""
x = keras.layers.Input(
shape=(image_size[0], image_size[1], 3), tensor=input_tensor)
out = cls._convolution(x, int(alpha * 32), 9, strides=1)
out = cls._convolution(out, int(alpha * 32), 3, strides=2)
out = cls._convolution(out, int(alpha * 32), 3, strides=2)
out = cls._residual_block(out, int(alpha * 32))
out = cls._residual_block(out, int(alpha * 32))
out = cls._residual_block(out, int(alpha * 32))
out = cls._upsample(out, int(alpha * 32), 3)
out = cls._upsample(out, int(alpha * 32), 3)
out = cls._convolution(out, 3, 9, relu=False, padding='same')
# Restrict outputs of pixel values to -1 and 1.
out = keras.layers.Activation('tanh')(out)
# Deprocess the image into valid image data. Note we'll need to define
# a custom layer for this in Core ML as well.
out = layers.DeprocessStylizedImage()(out)
model = keras.models.Model(inputs=x, outputs=out)
帶有寬度引數的較小風格遷移網路。
通過反覆嘗試，作者發現仍然可以用上述架構實現良好的風格遷移，一直到寬度引數為 0.3，在每一層上留下 9 個濾波器。最終結果是一個只有 11,868 個權重的神經網路。任何權重低於 10000 的網路都不能持續訓練，並且會產生糟糕的風格化影象。
值得一提的是，剪枝技術是在網路訓練之前應用的。在訓練期間或訓練後反覆修剪，你可以在很多工上實現更高的效能。
量化
最後一段壓縮是在網路訓練完成後進行的。神經網路權重通常儲存為 64 位或 32 位浮點數。量化過程將每一個浮點權重對映到具有較低位寬的整數。從 32 位浮點權重變為 8 位整型，使得儲存大小減少了 4 倍。作者利用 Alexis Creuzot 在部落格中提出的方法（https://heartbeat.fritz.ai/reducing-coreml2-model-size-by-4x-with-quantization-in-ios12-b1c854651c4），在不怎麼影響風格的情況下使浮點數降低到了 8 位量化。
現在所有主要的移動框架都支援量化，如TensorFlow Mobile、TensorFlow Lite、Core ML 和 Caffe2Go。
最終結果
該微型網路架構有 11,868 個引數，相比之下，Johnson 最初的模型具有 170 萬個引數，大小為 1.7MB。當轉化為 Core ML 並量化時，最終大小僅為 17KB——為原始大小的 1/400。以下是在梵高的《Starry Night》上的訓練結果。
此微型風格遷移結果的實時視訊可在 Heartbeat App 上檢視：
http://bit.ly/heartbeat-ios
作者驚訝地發現，儘管尺寸相差 400 倍，但在 iPhone X 上，這款微型模型的執行速度僅快了 50%。原因可能是計算與這一通用架構相關，也可能是將影象遷移到 GPU 進行處理時造成的。
如果你對結果表示懷疑，可以自己下載並執行此微型模型。甚至訓練自己的模型！
-
下載地址：https://github.com/fritzlabs/fritz-style-transfer/blob/master/example/starry_night_640x480_small_a03_q8.mlmodel
總而言之，作者用兩種簡單的技術將風格遷移神經網路的規模減小了 99.75%。使用簡單的 width multiplier 超引數修剪層，訓練後的權重從 32 位浮點數量化為 8 位整數。未來，作者期待看到將這些方法泛化到其它神經網路的效果。風格遷移相對簡單，因為「準確率」肉眼可見。對於影象識別這樣更加可以量化的任務而言，如此極端的修剪可能帶來更明顯的效能下降。
Reddit 討論
這篇帖子下有一些 reddit 網友對該專案提出了質疑：
gwern：
看你的部落格，剪枝部分似乎沒有移除任何層，只是更改了層的寬度／濾波器，然後對所有引數進行量化。如果所有層都在（因為你沒有做任何類似於訓練較寬的淺層網路的工作來模仿原始深度教師網路），那麼它們仍將從後續計算的每一層中引入大量延遲，即使每一層都很小。（由於你可以在手機 GPU 上安裝更多模型，每個模型使用較少的 FLOPS，因此整體吞吐量會變得更好。但是每個模型的迭代仍然需要一段時間，在特定大小之後，每一層基本上是即時的。）
jamesonatfritz 回覆：
你說得對，濾波器剪枝部分確實沒有移除層，但是我去掉了兩個殘差塊，從而消除了一些層。整體 FLOPs 的降低情況不如全部權重數量的減少情況，這一點你說得對。不幸的是，Apple 沒法讓你較好地控制模型執行的位置。你無法強制該模型使用 GPU。一些啟發式方法導致較小的模型僅在 CPU 上執行，這是可能的。
gwern：
「去掉了兩個殘差塊，從而消除了一些層。」
你藉此獲得了一些加速，但是使用更扁平的模型或許會實現更多加速。
jamesonatfritz 回覆：
確實如此。我試過的最扁平模型只有一個卷積層、一個殘差模組和一個上取樣模組，但我發現這些變體無法收斂。
gwern：
這似乎有些過了：只有一個層有些過於難了。我想的是三四個層這樣，在預訓練風格遷移模型的確切畫素輸出上進行訓練。或許值得一試。
Ikuyas：
這個方法和直接使用小模型有什麼區別嗎？我確定使用 11,868 個引數進行訓練結果會更好。另外，1.7M引數太大了，每個引數的貢獻估計會很小。但是，真實情況是隻有幾百個引數是真正重要的，其他引數只是到處吸收一點微小的噪聲。
從標準迴歸的角度來看，這似乎是完美的預期結果。
作者回復：
我應該在文章裡寫清楚的，事實上你所說的正是我所做的。剪枝發生在訓練之前。反直覺的一件事是，實際上使用較少的引數從頭開始訓練模型無法確保能得到一樣的結果。儘管一小部分權重比較重要，但你很難提前知道哪些權重是重要的。詳情參見論文：《The Lottery Ticket Hypothesis: Finding Sparse, TrainableNeural Networks》。
Ikuyas：
神經網路中的引數缺乏有意義的解釋，這是第一堂機器學習課程中就學過的。這並不反直覺，而是預料之中。剪枝後的引數甚至並不被認為是吸收噪聲的神經元。對於標準迴歸模型來說，噪聲有時似乎像是正態分佈的實現。而神經網路擬合併不假設任何此類事情。因此使用較少的引數，你可以用完全不同的模型擬合數據。剪枝技術並沒有什麼用。