文章目錄

• • 前言
  • 1. 交替方向乘子法
  • 2. 論文中的表述
  • 3. 對論文中的公式進行推導
  • 4. 代碼流程
  • 5. 主要函數實現
  • 6. dense vs. prune(finetune)
  • 結束語

前言

??本篇博客記錄一下自己根據對論文 GRIM: A General, Real-Time Deep Learning Inference Framework for Mobile Devices based on Fine-Grained Structured Weight Sparsity 中提到的ADMM算法的理解，給出了ADMM算法的推導過程，并在文章的末尾提供了實現的代碼。

1. 交替方向乘子法

??交替方向乘子法(Alternating Direction Method of Multipliers, ADMM)作為一種求解優化問題的計算框架，適用于求解凸優化問題。ADMM算法的思想根源可以追溯到20世紀50年代，在20世紀八九十年代中期存在大量的文章分析這種方法的性質，但是當時ADMM主要用于解決偏微分方程問題。1970年由 R. Glowinski 和 D. Gabay 等提出的一種適用于可分離凸優化的簡單有效方法，并在統計機器學習、數據挖掘和計算機視覺等領域中得到了廣泛應用。ADMM算法主要解決帶有等式約束的關于兩個變量的目標函數的最小化問題，可以看作在增廣拉朗格朗日算法基礎上發展的算法，混合了對偶上升算法(Dual Ascent)的可分解性和乘子法(Method of Multipliers)的算法優越的收斂性。相對于乘子法，ADMM算法最大的優勢在于其能夠充分利用目標函數的可分解性，對目標函數中的多變量進行交替優化。在解決大規模問題上，利用ADMM算法可以將原問題的目標函數等價地分解成若干個可求解的子問題，然后并行求解每一個子問題，最后協調子問題的解得到原問題的全局解。¹

??優化問題
$$\begin{gathered} minimizef(x)+g(z) \\ subjecttoAx+Bz=c \end{gathered}$$??其中，$x\in R^n,z\in R^m,A\in R^{p\times n},B\in R^{p\times m},c\in R^p$，構造拉格朗日函數為
$$L_p(x,z,\lambda )=f(x)+g(z)+{\lambda }^T(Ax+Bz?c)$$??其增廣拉格朗日函數(augmented Lagrangian function)為
$$L_p(x,z,\lambda )=f(x)+g(z)+{\lambda }^T(Ax+Bz?c)+\frac{\rho }{2}||Ax+Bz?c|{|}^2$$??對偶上升法迭代更新
$$\begin{gathered} (x^{k+1},z^{k+1})=\underset{x,z}{argmin}{L}_p(x,z,{\lambda }^k) \\ {\lambda }^{k+1}={\lambda }^k+\rho (A{x}^{k+1}+B{z}^{k+1}?c) \end{gathered}$$??交替方向乘子法則是在$(x,z)$一起迭代的基礎上將$x,z$分別固定單獨交替迭代，即
$$\begin{gathered} x^{k+1}=\underset{x}{argmin}{L}_p(x,z^k,{\lambda }^k) \\ {z}^{k+1}=\underset{z}{argmin}{L}_p(x^{k+1},z,{\lambda }^k) \\ {\lambda }^{k+1}={\lambda }^k+\rho (A{x}^{k+1}+B{z}^{k+1}?c) \end{gathered}$$??交替方向乘子的另一種等價形式，將殘差定義為$r^k=A{x}^k+B{z}^k?c$，同時定義$u^k=\frac{1}{\rho }{\lambda }^k$作為縮放的對偶變量(dual variable)，有
$$({\lambda }^k{)}^T{r}^k+\frac{\rho }{2}||r^k|{|}^2=\frac{\rho }{2}||r^k+u^k|{|}^2?\frac{\rho }{2}||u^k|{|}^2$$??改寫 ADMM 的迭代過程
$$\begin{gathered} x^{k+1}=\underset{x}{argmin}\{f(x)+\frac{\rho }{2}||Ax+B{z}^k?c+u^k|{|}^2\} \\ {z}^{k+1}=\underset{z}{argmin}\{g(z)+\frac{\rho }{2}||A{x}^{k+1}+Bz?c+u^k|{|}^2\} \\ {u}^{k+1}=u^k+A{x}^{k+1}+B{z}^{k+1}?c \end{gathered}$$

2. 論文中的表述

3. 對論文中的公式進行推導

??為便于推導公式，將論文中的進行簡化，參數W和b簡記為W，此時的優化問題變為
$$\begin{gathered} minimizef(W_i)+\sum _{i=1}^{N}g(Z_i) \\ subjectto{W}_i=Z_i,i=1,2,...,N \end{gathered}$$??構造拉格朗日函數為
$$L_p(w,z,\lambda )=f(w)+\sum g(z)+{\lambda }^T(w?z)$$??其增廣拉格朗日函數為
$$L_p(w,z,\lambda )=f(w)+\sum g(z)+{\lambda }^T(w?z)+\sum \frac{\rho }{2}||w?z|{|}^2$$??交替方向乘子法：在(x, z)一起迭代的基礎上將 x, z 分別固定，單獨交替迭代，即
$$\begin{gathered} w^{k+1}=\underset{w}{argmin}{L}_p(w,z^k,{\lambda }^k) \\ {z}^{k+1}=\underset{z}{argmin}{L}_p(w^{k+1},z,{\lambda }^k) \\ {\lambda }^{k+1}={\lambda }^k+\sum \rho (w?z) \end{gathered}$$??定義一個對偶變量
$$u^k=\frac{1}{\rho }{\lambda }^k$$??改寫ADMM的迭代過程
$$\begin{gathered} w^{k+1}=\underset{w}{argmin}\{f(w)+\sum \frac{\rho }{2}||w?z^k+u^k|{|}^2\} \\ {z}^{k+1}=\underset{z}{argmin}\{\sum g(z)+\sum \frac{\rho }{2}||w^{k+1}?z+u^k|{|}^2\} \\ {u}^{k+1}=u^k+w^{k+1}?z^{k+1} \end{gathered}$$

4. 代碼流程

# 初始化參數Z和U Z, U = initialize_Z_and_U(model)# 訓練model，并更新X，Z，U，損失函數為admm loss for epoch in range(epochs):for data, target in train_loader:optimizer.zero_grad()output = model(data)loss = admm_loss(model, Z, U, output, target)loss.backward()optimizer.step()W = update_W(model)Z = update_Z(W, U, percent)U = update_U(U, W, Z)# 對weight進行剪枝，返回 mask mask = apply_prune(model, percent)# 對剪枝后的model進行finetune finetune(model, mask, train_loader, test_loader, optimizer)

5. 主要函數實現

def admm_loss(args, device, model, Z, U, output, target):idx = 0loss = F.nll_loss(output, target)for name, param in model.named_parameters():if name.split('.')[-1] == "weight":u = U[idx].to(device)z = Z[idx].to(device)# 這里就是推導出來的admm的表達式loss += args.rho / 2 * (param - z + u).norm()return lossdef update_W(model):W = ()for name, param in model.named_parameters():if name.split('.')[-1] == "weight":W += (param.detach().cpu().clone(),)return Wdef update_Z(W, U, args):new_Z = ()idx = 0for w, u in zip(W, U):z = w + upcen = np.percentile(abs(z), 100*args.percent[idx])under_threshold = abs(z) < pcen# percent剪枝率，小于percent分位數的置為0z.data[under_threshold] = 0new_Z += (z,)idx += 1return new_Zdef update_U(U, W, Z):new_U = ()for u, w, z in zip(U, W, Z):new_u = u + w - znew_U += (new_u,)return new_Udef prune_weight(weight, device, percent):# to work with admm, we calculate percentile based on all elements instead of nonzero elements.weight_numpy = weight.detach().cpu().numpy()pcen = np.percentile(abs(weight_numpy), 100*percent)under_threshold = abs(weight_numpy) < pcen# 非結構化剪枝weight_numpy[under_threshold] = 0mask = torch.Tensor(abs(weight_numpy) >= pcen).to(device)return mask

6. dense vs. prune(finetune)

結束語

??對論文中算法的推導僅限于自己的理解，可能還存在一些問題，歡迎來評論區交流哦^_^

參考教程

---

《分布式機器學習：交替方向乘子法在機器學習中的應用》---- 雷大江著 ??

總結

以上是生活随笔為你收集整理的ADMM算法在神经网络模型剪枝方面的应用的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。