僅17 KB、一萬個權重的微型風格遷移網絡!

選自fritz.ai,機器之心編譯,機器之心編輯部。

今天 reddit 上一篇帖子引起了熱議,博主 jamesonatfritz 稱他將原本具備 1.7M 參數的風格遷移網絡減少到只有 11,868 個參數,該網絡仍然能夠輸出風格化的圖像。且量化後的最終網絡體積僅有 17 kB,非常適合移動 app。

jamesonatfritz 想解決神經網絡的過參數化問題,想要創建體積小但性能優的神經網絡。他所試驗的第一個任務便是藝術風格遷移。

  • GitHub 鏈接:https://github.com/fritzlabs/fritz-style-transfer

現在有很多用來訓練藝術風格遷移模型的現成工具,還有上千種開源實現。其中的多數工具利用 Johnson 等人在《Perceptual Losses for Real-Time Style Transfer and Super-Resolution》中提出的網絡架構的變體來實現快速、前饋的風格化。因此,多數遷移模型的大小是 7MB。對於你的應用來說,這個負擔並非不可承受,但也並非無足輕重。

研究表明,神經網絡的體積通常遠遠大於所需,數百萬的權重中有很多並不重要。因此作者創造了一個體積大大縮小的可靠風格遷移模型:一個只有 11686 個訓練權重的 17KB 神經網絡。

僅17 KB、一萬個權重的微型風格遷移網絡!

左:原圖;中:來自上述 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 Convolutional Neural Networks for Mobile Vision》中提出,非常簡單、高效。

width multiplier 利用一個恆定係數調整每個卷積層中的卷積核數量。對於給定的層及 width multiplier alpha,卷積核數量 F 變為 alpha * F。

有了這個超參數,我們就可以生成一系列架構相同但權重數不同的網絡。訓練每種配置,就可以在模型速度、大小及準確率之間做出權衡。

下面是作者模仿 Johnson 等人在《The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural 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 parameter 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 parameter 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

僅17 KB、一萬個權重的微型風格遷移網絡!

作者驚訝地發現,儘管尺寸相差 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, Trainable Neural Networks》。

Ikuyas:

神經網絡中的參數缺乏有意義的解釋,這是第一堂機器學習課程中就學過的。這並不反直覺,而是預料之中。剪枝後的參數甚至並不被認為是吸收噪聲的神經元。對於標準迴歸模型來說,噪聲有時似乎像是正態分佈的實現。而神經網絡擬合併不假設任何此類事情。因此使用較少的參數,你可以用完全不同的模型擬合數據。剪枝技術並沒有什麼用。


分享到:


相關文章: