从源码解析 Bert 的 BertPooler 模块

前文链接：

1. 从源码层面，深入理解 Bert 框架
2. 从源码解析 Bert 的 Embedding 模块
3. ​从源码解析 Bert 的 BertEncoder 模块

以上文章解析了Bert的BertEmbedding、BertEncoder模块，接下来分析bert的第三个重要模块 BertPooler；
BertPooler模块本质上就是一个全连接层网络，或者说一个分类器；BertPooler模块接在bert的堆叠encoder层后，主要用于解决序列级的NLP任务

BertPooler模块源码

class BertPooler(nn.Module):
    def __init__(self, config):
        super(BertPooler, self).__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states):
        # We "pool" the model by simply taking the hidden state corresponding
        # to the first token.
        first_token_tensor = hidden_states[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output

1. __ init __(self, config)函数

__ init __(self, config)函数根据配置文件里的信息确定了全连接层网络的隐含层节点个数，并且定义了网络的激活函数为 Tanh() 函数

2. forward(self, hidden_states)函数

forward()函数的目的就是将BertEncoder模块得到数据放入__init__()函数定义好好的分类器中，并将分类器的结果返回

这里着重解释一下这行代码

first_token_tensor = hidden_states[:, 0]

这句代码的作用是：使BertPooler模块的全连接层只接收BertEncoder模块中第一个token（[CLS]）对应的输出，因为BertPooler模块用于解决序列级任务，而序列级任务只会用到[CLS]位置对应的输出