可逆神經(jīng)網(wǎng)絡(luò)（Invertible Neural Networks）詳細(xì)解析：讓神經(jīng)網(wǎng)絡(luò)更加輕量化

來(lái)源：PaperWeekly

為什么要用可逆網(wǎng)絡(luò)呢？

1. 因?yàn)榫幋a和解碼使用相同的參數(shù)，所以 model 是輕量級(jí)的?？赡娴慕翟刖W(wǎng)絡(luò) InvDN 只有 DANet 網(wǎng)絡(luò)參數(shù)量的 4.2%，但是 InvDN 的降噪性能更好。
2. 由于可逆網(wǎng)絡(luò)是信息無(wú)損的，所以它能保留輸入數(shù)據(jù)的細(xì)節(jié)信息。
3. 無(wú)論網(wǎng)絡(luò)的深度如何，可逆網(wǎng)絡(luò)都使用恒定的內(nèi)存來(lái)計(jì)算梯度。

其中最主要目的就是為了減少內(nèi)存的消耗，當(dāng)前所有的神經(jīng)網(wǎng)絡(luò)都采用反向傳播的方式來(lái)訓(xùn)練，反向傳播算法需要存儲(chǔ)網(wǎng)絡(luò)的中間結(jié)果來(lái)計(jì)算梯度，而且其對(duì)內(nèi)存的消耗與網(wǎng)絡(luò)單元數(shù)成正比。這也就意味著，網(wǎng)絡(luò)越深越廣，對(duì)內(nèi)存的消耗越大，這將成為很多應(yīng)用的瓶頸。

下面是 Pytorch summary 的結(jié)果，F(xiàn)orward/backward pass size（MB）: 218.59 就是需要保存的中間變量大小，可以看出這部分占據(jù)了很大部分顯存（隨著網(wǎng)絡(luò)深度的增加，中間變量占據(jù)顯存量會(huì)一直增加，resnet152（size=224）的中間變量更是占據(jù)總共內(nèi)存的 606.6÷836.79≈0.725 ）。如果不存儲(chǔ)中間層結(jié)果，那么就可以大幅減少 GPU 的顯存占用，有助于訓(xùn)練更深更廣的網(wǎng)絡(luò)。

import torchfrom torchvision import modelsfrom torchsummary import summary
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')vgg = models.vgg16().to(device)
summary(vgg, (3, 224, 224))

----------------------------------------------------------------        Layer (type)               Output Shape         Param #================================================================            Conv2d-1         [-1, 64, 224, 224]           1,792              ReLU-2         [-1, 64, 224, 224]               0            Conv2d-3         [-1, 64, 224, 224]          36,928              ReLU-4         [-1, 64, 224, 224]               0         MaxPool2d-5         [-1, 64, 112, 112]               0            Conv2d-6        [-1, 128, 112, 112]          73,856              ReLU-7        [-1, 128, 112, 112]               0            Conv2d-8        [-1, 128, 112, 112]         147,584              ReLU-9        [-1, 128, 112, 112]               0        MaxPool2d-10          [-1, 128, 56, 56]               0           Conv2d-11          [-1, 256, 56, 56]         295,168             ReLU-12          [-1, 256, 56, 56]               0           Conv2d-13          [-1, 256, 56, 56]         590,080             ReLU-14          [-1, 256, 56, 56]               0           Conv2d-15          [-1, 256, 56, 56]         590,080             ReLU-16          [-1, 256, 56, 56]               0        MaxPool2d-17          [-1, 256, 28, 28]               0           Conv2d-18          [-1, 512, 28, 28]       1,180,160             ReLU-19          [-1, 512, 28, 28]               0           Conv2d-20          [-1, 512, 28, 28]       2,359,808             ReLU-21          [-1, 512, 28, 28]               0           Conv2d-22          [-1, 512, 28, 28]       2,359,808             ReLU-23          [-1, 512, 28, 28]               0        MaxPool2d-24          [-1, 512, 14, 14]               0           Conv2d-25          [-1, 512, 14, 14]       2,359,808             ReLU-26          [-1, 512, 14, 14]               0           Conv2d-27          [-1, 512, 14, 14]       2,359,808             ReLU-28          [-1, 512, 14, 14]               0           Conv2d-29          [-1, 512, 14, 14]       2,359,808             ReLU-30          [-1, 512, 14, 14]               0        MaxPool2d-31            [-1, 512, 7, 7]               0           Linear-32                 [-1, 4096]     102,764,544             ReLU-33                 [-1, 4096]               0          Dropout-34                 [-1, 4096]               0           Linear-35                 [-1, 4096]      16,781,312             ReLU-36                 [-1, 4096]               0          Dropout-37                 [-1, 4096]               0           Linear-38                 [-1, 1000]       4,097,000================================================================Total params: 138,357,544Trainable params: 138,357,544Non-trainable params: 0----------------------------------------------------------------Input size (MB): 0.57Forward/backward pass size (MB): 218.59Params size (MB): 527.79Estimated Total Size (MB): 746.96----------------------------------------------------------------

可逆網(wǎng)絡(luò)具有的性質(zhì)：

1. 網(wǎng)絡(luò)的輸入、輸出的大小必須一致。

2. 網(wǎng)絡(luò)的雅可比行列式不為 0。

雅可比行列式通常稱為雅可比式（Jacobian），它是以 n 個(gè) n 元函數(shù)的偏導(dǎo)數(shù)為元素的行列式 。事實(shí)上，在函數(shù)都連續(xù)可微（即偏導(dǎo)數(shù)都連續(xù)）的前提之下，它就是函數(shù)組的微分形式下的系數(shù)矩陣（即雅可比矩陣）的行列式。若因變量對(duì)自變量連續(xù)可微，而自變量對(duì)新變量連續(xù)可微,則因變量也對(duì)新變量連續(xù)可微。這可用行列式的乘法法則和偏導(dǎo)數(shù)的連鎖法則直接驗(yàn)證。也類似于導(dǎo)數(shù)的連鎖法則。偏導(dǎo)數(shù)的連鎖法則也有類似的公式；這常用于重積分的計(jì)算中。

為什么神經(jīng)網(wǎng)絡(luò)會(huì)與雅可比行列式有關(guān)系？這里我借用李宏毅老師的 ppt（12-14頁(yè)）。想看視頻的可以到 b 站上看。

簡(jiǎn)單的來(lái)講就是 ，他們的分布之間的關(guān)系就變?yōu)?nbsp;，又因?yàn)橛?nbsp;，所以  這個(gè)網(wǎng)絡(luò)的雅可比行列式不為 0 才行。

順便提一下，flow-based Model 優(yōu)化的損失函數(shù)如下：

其實(shí)這里跟矩陣運(yùn)算很像，矩陣可逆的條件也是矩陣的雅可比行列式不為 0，雅可比矩陣可以理解為矩陣的一階導(dǎo)數(shù)。
假設(shè)可逆網(wǎng)絡(luò)的表達(dá)式為：

它的雅可比矩陣為：

其行列式為 1。

1.3.1 可逆塊結(jié)構(gòu)
可逆神經(jīng)網(wǎng)絡(luò)將每一層分割成兩部分，分別為  和 ，每一個(gè)可逆塊的輸入是 ，輸出是 。其結(jié)構(gòu)如下：
正向計(jì)算圖示：

公式表示：

逆向計(jì)算圖示：

公式表示：

其中 F 和 G 都是相似的殘差函數(shù)，參考上圖殘差網(wǎng)絡(luò)。可逆塊的跨距只能為 1，也就是說(shuō)可逆塊必須一個(gè)接一個(gè)連接，中間不能采用其它網(wǎng)絡(luò)形式銜接，否則的話就會(huì)丟失信息，并且無(wú)法可逆計(jì)算了，這點(diǎn)與殘差塊不一樣。如果一定要采取跟殘差塊相似的結(jié)構(gòu)，也就是中間一部分采用普通網(wǎng)絡(luò)形式銜接，那中間這部分的激活結(jié)果就必須顯式的存起來(lái)。
1.3.2 不用存儲(chǔ)激活結(jié)果的反向傳播
為了更好地計(jì)算反向傳播的步驟，我們修改一下上述正向計(jì)算和逆向計(jì)算的公式：

盡管  和  的值是相同的，但是兩個(gè)變量在圖中卻代表不同的節(jié)點(diǎn)，所以在反向傳播中它們的總體導(dǎo)數(shù)是不一樣的。 的導(dǎo)數(shù)包含通過(guò)  產(chǎn)生的間接影響，而  的導(dǎo)數(shù)卻不受  的任何影響。
在反向傳播計(jì)算流程中，先給出最后一層的激活值  和誤差傳播的總體導(dǎo)數(shù) ，然后要計(jì)算出其輸入值  和對(duì)應(yīng)的導(dǎo)數(shù) ，以及殘差函數(shù) F 和 G 中權(quán)重參數(shù)的總體導(dǎo)數(shù)，求解步驟如下：

1.3.3 計(jì)算開(kāi)銷
一個(gè) N 個(gè)連接的神經(jīng)網(wǎng)絡(luò)，正向計(jì)算的理論加乘開(kāi)銷為 N，反向傳播求導(dǎo)的理論加乘開(kāi)銷為 2N（反向求導(dǎo)包含復(fù)合函數(shù)求導(dǎo)連乘），而可逆網(wǎng)絡(luò)多一步需要反向計(jì)算輸入值的操作，所以理論計(jì)算開(kāi)銷為 4N，比普通網(wǎng)絡(luò)開(kāi)銷約多出 33% 左右。但是在實(shí)際操作中，正向和反向的計(jì)算開(kāi)銷在 GPU 上差不多，可以都理解為 N。那么這樣的話，普通網(wǎng)絡(luò)的整體計(jì)算開(kāi)銷為 2N，可逆網(wǎng)絡(luò)的整體開(kāi)銷為 3N，也就是多出了約 50%。
1.3.4 雅可比行列式的計(jì)算

其編碼公式如下：

其解碼公式如下：

為了計(jì)算雅可比矩陣，我們更直觀的寫(xiě)成下面的編碼公式：

它的雅可比矩陣為：

其實(shí)上面這個(gè)雅可比行列式也是 1，因?yàn)檫@里 ，它們的系數(shù)是一樣的。
有另外一種解釋方式就是把這種對(duì)偶的形式切成兩半：

其行列式為 1。

因?yàn)槭菍?duì)偶的形式，所以這里的行列式也為 1。
因?yàn)?nbsp;，所以其行列式也為 1。

上圖中符號(hào)的含義：

• x1，x2，x3：表示 3 個(gè)輸入層節(jié)點(diǎn)。
• ：表示從 t-1 層到 t 層的權(quán)重參數(shù)，j 表示 t 層的第 j 個(gè)節(jié)點(diǎn)，i 表示 t-1 層的第 i 個(gè)節(jié)點(diǎn)。
• ：表示 t 層的第 i 個(gè)激活后輸出結(jié)果。
• g(x)：表示激活函數(shù)。

正向傳播計(jì)算過(guò)程：

• 隱藏層（網(wǎng)絡(luò)的第二層）

• 輸出層（網(wǎng)絡(luò)的最后一層）

• 輸出層

• 隱藏層

• 梯度消失問(wèn)題；

• 網(wǎng)絡(luò)退化問(wèn)題。