預計閱讀時間：8分鐘
ALBERT使用層級引數共享對BERT進行壓縮，這個理論上是非常好理解的，但實現上卻有一些細節需要考慮。走讀albert pytorch版本的原始碼之時，覺得寫得挺優雅，而且靈活，遂記錄在此。
閱讀本文前，可能需要先了解一些Albert基礎知識以及Pytorch框架。
文字參考的原始碼版本
Albert 原始碼整體和bert程式碼差異很小，為了實現靈活的引數共享，作者提出了一個Group的概念。 原始碼中將每一層都分到一個指定Group之中，一個Group包含了多個相鄰的層，
同一個Group裡面的層是引數共享的
，這個group個數由
num_hidden_groups
引數決定，預設為1。即所有的層share同一個Transformer權重。
如num_hidden_groups 為2，
num_hidden_layers
為12，那麼層分為兩組。1~6層是第一組，7~12是第二組。
如num_hidden_groups 為3，
num_hidden_layers
為12 ，那麼層分為三組。1~4為第一組，5~8為第二組，9~12為第三組。
以此類推。。。
Layers Group Split。 相同顏色的layer引數共享
層索引
layer_idx
和組索引
group_idx
由程式碼計算得出：
group_idx = int（layer_idx / num_hidden_layers * num_hidden_groups）
對於group編號較低的組，學習低抽象的知識，group編號較高的組，學習相對較高抽象的知識，這個是make sense的。透過設定
num_hidden_groups
和
num_hidden_layers
可以靈活設定模型深度和共享的程度。
可見group在原始碼中是一個比較重要的型別，其由 AlbertGroup類實現。AlbertTransformer 在初始化之時，預先例項化了
num_hidden_groups
個AlbertGroup，這個AlbertGroup代表新的一層引數，裡面還有一些細節，後文會描述。
self
。
group
=
nn
。
ModuleList
（［
AlbertGroup
（
config
）
for
_
in
range
（
config
。
num_hidden_groups
）］）
當AlbertTransformer計算到某一層時，直接在group列表中找到對應的AlbertGroup去forward，因此在梯度backward的時候，梯度的變化也會傳遞到對應的AlbertGroup，這樣就實現了多級引數共享的引數更新。
AlbertTransformer 完整程式碼：
class
AlbertTransformer
（
nn
。
Module
）：
def
__init__
（
self
，
config
）：
super
（
AlbertTransformer
，
self
）
。
__init__
（）
self
。
output_attentions
=
config
。
output_attentions
self
。
output_hidden_states
=
config
。
output_hidden_states
self
。
num_hidden_layers
=
config
。
num_hidden_layers
self
。
num_hidden_groups
=
config
。
num_hidden_groups
self
。
group
=
nn
。
ModuleList
（［
AlbertGroup
（
config
）
for
_
in
range
（
config
。
num_hidden_groups
）］）
def
forward
（
self
，
hidden_states
，
attention_mask
，
head_mask
）：
all_hidden_states
=
（）
all_attentions
=
（）
for
layer_idx
in
range
（
self
。
num_hidden_layers
）：
if
self
。
output_hidden_states
and
layer_idx
==
0
：
all_hidden_states
=
all_hidden_states
+
（
hidden_states
，）
group_idx
=
int
（
layer_idx
/
self
。
num_hidden_layers
*
self
。
num_hidden_groups
）
layer_module
=
self
。
group
［
group_idx
］
layer_outputs
=
layer_module
（
hidden_states
，
attention_mask
，
head_mask
［
layer_idx
］）
hidden_states
=
layer_outputs
［
0
］［
-
1
］
if
self
。
output_attentions
：
all_attentions
=
all_attentions
+
layer_outputs
［
1
］
if
self
。
output_hidden_states
：
all_hidden_states
=
all_hidden_states
+
layer_outputs
［
0
］
outputs
=
（
hidden_states
，）
if
self
。
output_hidden_states
：
outputs
=
outputs
+
（
all_hidden_states
，）
if
self
。
output_attentions
：
outputs
=
outputs
+
（
all_attentions
，）
其他就是些很常規的程式碼了，和BERT原始碼基本類似。
除了對layers進行分組之外
，
group的內部也分了層級
AlbertGroup 完整程式碼：
class
AlbertGroup
（
nn
。
Module
）：
def
__init__
（
self
，
config
）：
super
（
AlbertGroup
，
self
）
。
__init__
（）
self
。
inner_group_num
=
config
。
inner_group_num
self
。
inner_group
=
nn
。
ModuleList
（［
AlbertLayer
（
config
）
for
_
in
range
（
config
。
inner_group_num
）］）
def
forward
（
self
，
hidden_states
，
attention_mask
，
head_mask
）：
layer_attentions
=
（）
layer_hidden_states
=
（）
for
inner_group_idx
in
range
（
self
。
inner_group_num
）：
layer_module
=
self
。
inner_group
［
inner_group_idx
］
layer_outputs
=
layer_module
（
hidden_states
，
attention_mask
，
head_mask
）
hidden_states
=
layer_outputs
［
0
］
layer_attentions
=
layer_attentions
+
（
layer_outputs
［
1
］，）
layer_hidden_states
=
layer_hidden_states
+
（
hidden_states
，）
return
（
layer_hidden_states
，
layer_attentions
）
AlbertGroup 內部的層級由
inner_group_num
引數確定，預設為1。其內部處理的邏輯也很簡單，即forward多層的AlbertLayer，這個AlbertLayer就代表著一個block。
由此可見，假設一個block的引數量為 m。則實際的encoder的引數量K為：
K = m * inner_group_num * num_hidden_groups
inner_group_num和num_hidden_groups預設均為1。大多數預訓練模型是基於預設引數訓練的，所以這兩個引數一般也不會改動。除非需要嘗試調整共享程度進行重新預訓練。
除了Group分組共享外albert 還可以調整 block內部的共享方式
分為三種
all ： ffn和attention都共享
ffn ：ffn共享
attention： attention共享
對於不同的內部共享，在初始化Module時將
不共享
元件例項化 inner_group_num * num_hidden_groups個 儲存在ModuleList之中，在forward時，按照索引定位到指定元件即可。
本文參考的程式碼並沒有實現內部共享，感興趣的讀者可以嘗試寫寫。
用碎片的時間，來總結學習
文章系原創，轉載請說明出處
