今天 reddit 上一篇帖子引起了熱議,博主 jamesonatfritz 稱他將原本具備 1.7M 參數(shù)的風格遷移網絡減少到只有 11,868 個參數(shù),該網絡仍然能夠輸出風格化的圖像。且量化后的最終網絡體積僅有 17 kB,非常適合移動 app。
jamesonatfritz 想解決神經網絡的過參數(shù)化問題,想要創(chuàng)建體積小但性能優(yōu)的神經網絡。他所試驗的第一個任務便是藝術風格遷移。
GitHub 鏈接:https://github.com/fritzlabs/fritz-style-transfer
現(xiàn)在有很多用來訓練藝術風格遷移模型的現(xiàn)成工具,還有上千種開源實現(xiàn)。其中的多數(shù)工具利用 Johnson 等人在《Perceptual Losses for Real-Time Style Transfer and Super-Resolution》中提出的網絡架構的變體來實現(xiàn)快速、前饋的風格化。因此,多數(shù)遷移模型的大小是 7MB。對于你的應用來說,這個負擔并非不可承受,但也并非無足輕重。
研究表明,神經網絡的體積通常遠遠大于所需,數(shù)百萬的權重中有很多并不重要。因此作者創(chuàng)造了一個體積大大縮小的可靠風格遷移模型:一個只有 11686 個訓練權重的 17KB 神經網絡。
左:原圖;中:來自上述 17KB 模型的風格化圖像;右:來自 7MB 模型的風格化圖像。
快速概覽:
原始模型:
大小:7MB
權重數(shù):1.7M
在 iPhone X 上的速度:18 FPS
小模型:
大小:17KB
權重數(shù):11,868
在 iPhone X 上的速度:29 FPS
如何縮小風格遷移模型
作者主要使用了兩種技術,而且都可以泛化到其他模型:
1. 大刀闊斧地修剪層和權重;
2. 通過量化將 32 位浮點權重轉換為 8 位整型
機器之心Synced剪枝小程序
修剪策略
卷積神經網絡通常包含數(shù)百萬甚至上億個需要在訓練階段進行調整的權重。通常來講,權重越多準確率越高。但這種增加權重提高準確率的做法非常低效。谷歌 MobileNetV2 的 stock 配置具有 347 萬個權重,內存占用達 16MB。InceptionV3 架構大小約為前者的 6 倍,具備 2400 萬個權重,內存占用達 92MB。盡管多了 2000 多萬個權重,但 InceptionV3 在 ImageNet 上的 top-1 分類準確率只比 MobileNetV2 高出 7 個百分點(80% vs 73%)。
因此,我們可以假設神經網絡中的多數(shù)權重沒有那么重要并將其移除。但重點是怎么做呢?我們可以選擇在三個層面進行修剪:單個權重、層、塊。
權重層面:假設某個神經網絡上的多數(shù)(>95%)權重都沒有什么用。如果能找出那些對準確率有影響的權重,就可以將其留下并將其他移除。
層層面:每個層中都包含一些權重。例如,2D 卷積層具有一個權重張量,即卷積核,用戶可以定義其寬度、高度和深度。縮小卷積核可以減小整個網絡的大小。
塊層面:多個層通常可以結合成可重復利用的子圖,即塊。以 ResNet 為例,它的名字來源于重復 10-50 次的「殘差塊」。在塊層面進行修剪可以移除多個層,從而一次性移除多個參數(shù)。
在實踐中,稀疏張量運算沒有很好的實現(xiàn),因此權重層面的修剪沒有多大價值。那么就只剩下層和塊層面的修剪了。
實踐中的修剪
作者使用的層修剪技術是引入 width multiplier 作為超參數(shù)。width multiplier 最初由谷歌在其論文《MobileNets: Efficient Convolutional Neural Networks for Mobile Vision》中提出,非常簡單、高效。
width multiplier 利用一個恒定系數(shù)調整每個卷積層中的卷積核數(shù)量。對于給定的層及 width multiplier alpha,卷積核數(shù)量 F 變?yōu)?alpha * F。
有了這個超參數(shù),我們就可以生成一系列架構相同但權重數(shù)不同的網絡。訓練每種配置,就可以在模型速度、大小及準確率之間做出權衡。
下面是作者模仿 Johnson 等人在《The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks》提出的網絡架構構建快速風格遷移模型的方法,不同之處在于添加了 width multiplier 作為超參數(shù)。
@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 個權重。
你可以構建一些寬度參數(shù)很小的網絡,但是也有相當多的重復塊。作者決定修剪掉一些,但實際操作后卻發(fā)現(xiàn)不能移除太多。即使參數(shù)量保持不變,較深的網絡能夠產生更好的結果。作者最終刪除了五個殘差塊中的兩個,并將每層的默認濾波器數(shù)量減少至 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)
帶有寬度參數(shù)的較小風格遷移網絡。
通過反復嘗試,作者發(fā)現(xiàn)仍然可以用上述架構實現(xiàn)良好的風格遷移,一直到寬度參數(shù)為 0.3,在每一層上留下 9 個濾波器。最終結果是一個只有 11,868 個權重的神經網絡。任何權重低于 10000 的網絡都不能持續(xù)訓練,并且會產生糟糕的風格化圖像。
值得一提的是,剪枝技術是在網絡訓練之前應用的。在訓練期間或訓練后反復修剪,你可以在很多任務上實現(xiàn)更高的性能。
量化
最后一段壓縮是在網絡訓練完成后進行的。神經網絡權重通常存儲為 64 位或 32 位浮點數(shù)。量化過程將每一個浮點權重映射到具有較低位寬的整數(shù)。從 32 位浮點權重變?yōu)?8 位整型,使得存儲大小減少了 4 倍。作者利用 Alexis Creuzot 在博客中提出的方法(https://heartbeat.fritz.ai/reducing-coreml2-model-size-by-4x-with-quantization-in-ios12-b1c854651c4),在不怎么影響風格的情況下使浮點數(shù)降低到了 8 位量化。
現(xiàn)在所有主要的移動框架都支持量化,如 TensorFlow Mobile、TensorFlow Lite、Core ML 和 Caffe2Go。
最終結果
該微型網絡架構有 11,868 個參數(shù),相比之下,Johnson 最初的模型具有 170 萬個參數(shù),大小為 1.7MB。當轉化為 Core ML 并量化時,最終大小僅為 17KB——為原始大小的 1/400。以下是在梵高的《Starry Night》上的訓練結果。
此微型風格遷移結果的實時視頻可在 Heartbeat App 上查看:
http://bit.ly/heartbeat-ios
作者驚訝地發(fā)現(xiàn),盡管尺寸相差 400 倍,但在 iPhone X 上,這款微型模型的運行速度僅快了 50%。原因可能是計算與這一通用架構相關,也可能是將圖像遷移到 GPU 進行處理時造成的。
如果你對結果表示懷疑,可以自己下載并運行此微型模型。甚至訓練自己的模型!
下載地址:https://github.com/fritzlabs/fritz-style-transfer/blob/master/example/starry_night_640x480_small_a03_q8.mlmodel
總而言之,作者用兩種簡單的技術將風格遷移神經網絡的規(guī)模減小了 99.75%。使用簡單的 width multiplier 超參數(shù)修剪層,訓練后的權重從 32 位浮點數(shù)量化為 8 位整數(shù)。未來,作者期待看到將這些方法泛化到其它神經網絡的效果。風格遷移相對簡單,因為「準確率」肉眼可見。對于圖像識別這樣更加可以量化的任務而言,如此極端的修剪可能帶來更明顯的性能下降。
Reddit 討論
這篇帖子下有一些 reddit 網友對該項目提出了質疑:
gwern:
看你的博客,剪枝部分似乎沒有移除任何層,只是更改了層的寬度/濾波器,然后對所有參數(shù)進行量化。如果所有層都在(因為你沒有做任何類似于訓練較寬的淺層網絡的工作來模仿原始深度教師網絡),那么它們仍將從后續(xù)計算的每一層中引入大量延遲,即使每一層都很小。(由于你可以在手機 GPU 上安裝更多模型,每個模型使用較少的 FLOPS,因此整體吞吐量會變得更好。但是每個模型的迭代仍然需要一段時間,在特定大小之后,每一層基本上是即時的。)
jamesonatfritz 回復:
你說得對,濾波器剪枝部分確實沒有移除層,但是我去掉了兩個殘差塊,從而消除了一些層。整體 FLOPs 的降低情況不如全部權重數(shù)量的減少情況,這一點你說得對。不幸的是,Apple 沒法讓你較好地控制模型運行的位置。你無法強制該模型使用 GPU。一些啟發(fā)式方法導致較小的模型僅在 CPU 上運行,這是可能的。
gwern:
「去掉了兩個殘差塊,從而消除了一些層。」
你借此獲得了一些加速,但是使用更扁平的模型或許會實現(xiàn)更多加速。
jamesonatfritz 回復:
確實如此。我試過的最扁平模型只有一個卷積層、一個殘差模塊和一個上采樣模塊,但我發(fā)現(xiàn)這些變體無法收斂。
gwern:
這似乎有些過了:只有一個層有些過于難了。我想的是三四個層這樣,在預訓練風格遷移模型的確切像素輸出上進行訓練。或許值得一試。
Ikuyas:
這個方法和直接使用小模型有什么區(qū)別嗎?我確定使用 11,868 個參數(shù)進行訓練結果會更好。另外,1.7M 參數(shù)太大了,每個參數(shù)的貢獻估計會很小。但是,真實情況是只有幾百個參數(shù)是真正重要的,其他參數(shù)只是到處吸收一點微小的噪聲。
從標準回歸的角度來看,這似乎是完美的預期結果。
作者回復:
我應該在文章里寫清楚的,事實上你所說的正是我所做的。剪枝發(fā)生在訓練之前。反直覺的一件事是,實際上使用較少的參數(shù)從頭開始訓練模型無法確保能得到一樣的結果。盡管一小部分權重比較重要,但你很難提前知道哪些權重是重要的。詳情參見論文:《The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks》。
Ikuyas:
神經網絡中的參數(shù)缺乏有意義的解釋,這是第一堂機器學習課程中就學過的。這并不反直覺,而是預料之中。剪枝后的參數(shù)甚至并不被認為是吸收噪聲的神經元。對于標準回歸模型來說,噪聲有時似乎像是正態(tài)分布的實現(xiàn)。而神經網絡擬合并不假設任何此類事情。因此使用較少的參數(shù),你可以用完全不同的模型擬合數(shù)據(jù)。剪枝技術并沒有什么用。