bert pytorch源码_Albert Pytorch 源码解读：共享实现方法

预计阅读时间：8分钟

---

ALBERT使用层级参数共享对BERT进行压缩，这个理论上是非常好理解的，但实现上却有一些细节需要考虑。走读albert pytorch版本的源码之时，觉得写得挺优雅，而且灵活，遂记录在此。

阅读本文前，可能需要先了解一些Albert基础知识以及Pytorch框架。

文本参考的源码版本

albert_pytorch​github.com

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内部的共享方式

分为三种

1. all : ffn和attention都共享
2. ffn :ffn共享
3. attention: attention共享

对于不同的内部共享，在初始化Module时将不共享组件实例化 inner_group_num * num_hidden_groups个 保存在ModuleList之中，在forward时，按照索引定位到指定组件即可。

本文参考的代码并没有实现内部共享，感兴趣的读者可以尝试写写。

---

用碎片的时间，来总结学习

文章系原创，转载请说明出处