# 大綱
在[DECORE: Deep Compression with Reinforcement Learning](https://openaccess.thecvf.com/content/CVPR2022/html/Alwani_DECORE_Deep_Compression_With_Reinforcement_Learning_CVPR_2022_paper.html)這篇論文使用了強化學習來做neural pruning。它在每個channel設置一個agent來學習該channel的"importance"(importance越小代表該channel被pruned之後對該深度神經網路沒有太大的改變)，藉由agent機率性的對channel開開關關來計算該agent的importance，且每個agent只有一個參數需要學習，並且每個agent之間是互相獨立的。透過研究這篇論文也可以順便學習強化學習的理論與實作。
# 理論
* 定義在第$i$個layer的第$j$個channel設置1個agent並賦予他一個權重$w_j$。權重$w_j$就是state，它被定義為$$s_i=\{w_0,w_1,\cdots\}$$那麼該agent有$$p_j=\frac{1}{1+e^{-w_j}}$$的機率讓channel $j$通行，這時候定義$c_j=1$否則為$0$。action定義為$$a_i=\{c_0,c_1,\cdots\}$$代表了在第$i$個layer的channel通行狀況。這邊就很直接了當的定義了強化學習中最重要的state與action。
* $R_{i,C}$定義為壓縮率reward：$$R_{i,C}=\sum_{j}1-a_{i,j}$$當越多的$a_{i,j}=0$(代表越多的channel被關起來)，$R_{i,C}$越大。
* $R_{acc}$定義為準確率reward，當channel的開關被決定之後，model獲得一個預測$\hat{y}$，當預測正確給予一個1的reward，否則給予處罰：$$R_{acc}=\left\{\begin{matrix}
1 & y=\hat{y}\\
-\lambda & y\neq\hat{y}
\end{matrix}\right.$$。$\lambda$是懲罰係數，論文的實驗得到$\lambda$越小有更好的壓縮率。
* $R_i$為整體reward，它被定義為壓縮率reward與準確率reward的相乘。$$R_i=R_{i,C}\times R_{acc}$$代表channel關掉越多還預測正確的話reward越高。所以可以理解為這個強化學習就是希望在壓縮率很高的時候又能預測正確。
* 那麼最重要的強化學習的loss function定義：$$J(\theta)=\sum_{\tau}\prod_t\pi_{\theta}(\tau_t)R(\tau)$$$\pi_\theta$是policy function且每個獨立的episode中step就只會是1，那麼loss function就只剩下：
$$J(\theta)=\sum_{\tau}\pi_{\theta}(\tau)R(\tau)$$這邊$\tau$代表了根據了state($\{s_1, s_2,\cdots\}$)給出action($\{a_1, a_2,\cdots\}$)的1次episode。並且因為agent是互相獨立的，所以
$$\pi_\theta(\tau)=\prod_{i,j}g(p_{i,j},c_{i,j})$$其中
$$g(p,c)=\left\{\begin{matrix}
p & c=1\\
1-p & c=0
\end{matrix}\right.$$
* 這邊強化學習的最佳化問題為
$$\theta^*=\sum_i\arg\max_\theta\mathbb{E}_{\tau_i\sim\pi_\theta(\tau_i)}(R_i)$$根據REINFORCE policy gradient algorithm，
$$\nabla J(\theta)=\frac{1}{N}\sum_b\sum_i\nabla\log\pi_{\theta}(s_{i,b}|a_{i,b})R_{i,b}$$這邊$i$和$b$分別是layer與batch的index，$N$是batch的個數，同一個batch每個input的開關狀況都是不一樣的。
# Pytorch實作
* channel的開關：在model forward的時候根據state $s_i=\{w_0,w_1,\cdots\}$算出機率$p_j$來決定channel的開關，而且同一個batch每個input的開關狀況都是不一樣的。首先需要在每個layer使用函式register_forward_hook來註冊一個forward hook：
```python=
from functools import partial
layer.register_forward_hook(partial(forward_hook, self=self))
```
forward_hook這個函式會在layer計算完之後(ex. 以linear來說就是$Ax=b$計算完之後執行)。這邊self是一個儲存$s_i=\{w_0,w_1,\cdots\}$與開關(gate)的一個class。接下來看forward_hook的程式：
```python=
def forward_hook(model, input, output, self):
self.P = torch.sigmoid(self.s_i)
batch_num = output.shape[0]
channel_num = output.shape[1]
self.gate = torch.zeros((batch_num, channel_num))
for i in range(batch_num):
self.gate[i] = torch.bernoulli(self.P)
for j in range(channel_num):
if self.gate[i][j] == 0:
output[i][j] = 0
```
這邊gate會去紀錄channel開關的情況，然後等$\hat{y}$計算完之後與$R_{acc}$一起計算$R_i$。
# 測試結果
以下面這個model在mnist手寫辨識問題上使用DECORE方法($\lambda$=-10, epochs=20)一樣的程式跑4次。
```python=
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Conv2d(in_channels=1, out_channels=8, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=3, padding=1)
self.conv3 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1)
self.conv4 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1)
self.linear1 = nn.Linear(in_features=64*7*7, out_features=128)
self.linear2 = nn.Linear(in_features=128, out_features=10)
def forward(self, x):
size = x.size(0)
x = F.relu(self.conv1(x))
x = F.max_pool2d(F.relu(self.conv2(x)), 2)
x = F.relu(self.conv3(x))
x = F.max_pool2d(F.relu(self.conv4(x)), 2)
x = x.view(size, -1)
x = F.relu(self.linear1(x))
x = F.softmax(self.linear2(x), dim=1)
return x
```
可以得到以下結果：
雖然DECORE在計算importance的有一定的效果，但是每次程式的執行都有不同的結果，而且變異還蠻大的。在論文裡面似乎沒有提到這一點，會是程式寫錯了嗎？接下來會嘗試測試actor-critic algorithm，預計能夠解決每次程式執行變異比較大的問題。
# 參考資料
https://openaccess.thecvf.com/content/CVPR2022/papers/Alwani_DECORE_Deep_Compression_With_Reinforcement_Learning_CVPR_2022_paper.pdf
Sign in with Wallet
Connect another wallet