【HuggingFace Transformers】BertEncoder源码解析

1. BertEncoder类 介绍

BertEncoder 类用于实现 BERT 模型的编码器部分。这个编码器由多个堆叠的 BertLayer 组成，用于处理输入的隐藏状态并生成新的隐藏状态。核心功能为：

• BertLayer 堆叠：编码器的核心是 BertLayer 的堆叠，每一层都处理输入的隐藏状态并生成新的隐藏状态。
• 梯度检查点：这一功能主要用于减少内存消耗，但需要注意它与缓存功能不兼容。
• 多种输出选择：可以选择是否输出所有隐藏状态、注意力权重，以及以何种形式返回结果（字典或元组）。
• 兼容性检查：代码中对梯度检查点与缓存的兼容性进行了检查，并在发现冲突时自动调整。

2. BertEncoder类 源码注释

源码地址：transformers/src/transformers/models/bert/modeling_bert.py

# -*- coding: utf-8 -*-
# @time: 2024/7/12 17:20

import torch

from torch import nn
from transformers import BertLayer
from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions
from transformers.utils import logging
from typing import Optional, Tuple, Union

logger = logging.get_logger(__name__)


class BertEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()

        # 保存配置对象，用于配置编码器的各种参数
        self.config = config

        # 使用 nn.ModuleList 创建一个包含多个 BertLayer 的列表
        # BertLayer 是 BERT 模型的基本组成部分，每个 BertLayer 使用相同的配置
        # config.num_hidden_layers 指定了要创建的 BertLayer 层数，bert-base版是12
        self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])

        # 初始化梯度检查点标志为 False(默认)
        # 梯度检查点（Gradient Checkpointing）是一种节省内存的技术，用于在计算过程中保存中间状态，以减少内存使用
        # 当需要节省内存时，可以通过设置这个标志来启用梯度检查点功能
        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = False,
        output_hidden_states: Optional[bool] = False,
        return_dict: Optional[bool] = True,
    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
        """这个方法执行前向传播操作，计算从输入的隐藏状态到输出隐藏状态的映射过程。"""

        """ 1. 初始化输出变量 """
        # 如果需要输出所有隐藏状态，初始化一个空元组来保存隐藏状态，否则为 None
        all_hidden_states = () if output_hidden_states else None
        # 如果需要输出注意力权重，初始化一个空元组来保存自注意力的权重，否则为 None
        all_self_attentions = () if output_attentions else None
        # 如果需要输出注意力权重且配置中添加了交叉注意力，初始化一个空元组来保存交叉注意力的权重，否则为 None
        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None

        """ 2. 梯度检查点 """
        # 检查是否启用了梯度检查点（gradient_checkpointing），如果启用了且 use_cache 为 True，则输出警告信息并将 use_cache 设置为 False。这是因为梯度检查点功能与缓存不兼容。
        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once("`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...")
                use_cache = False

        # 初始化 `next_decoder_cache`，如果 `use_cache` 为 True，则为一个空元组，否则为 None
        next_decoder_cache = () if use_cache else None

        """ 3. BertLayer 的堆叠计算 """
        # -------------------------------核心部分：多个BertLayer的堆叠-----------------------------------
        # 遍历所有的 BertLayer，对每一层执行前向传播操作
        for i, layer_module in enumerate(self.layer):
            # 如果需要输出所有隐藏状态，则将当前的隐藏状态添加到 all_hidden_states 中
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            # 获取当前层的头掩码，如果 head_mask 为 None，则 layer_head_mask 也为 None
            layer_head_mask = head_mask[i] if head_mask is not None else None
            # 获取当前层的过去的键值对，如果 past_key_values 为 None，则 past_key_value 也为 None
            past_key_value = past_key_values[i] if past_key_values is not None else None

            # 如果启用了梯度检查点并且当前处于训练模式
            if self.gradient_checkpointing and self.training:
                # 使用自定义的梯度检查点函数执行当前层的前向传播
                layer_outputs = self._gradient_checkpointing_func(
                    layer_module.__call__,
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                )
            else:
                # 否则，正常执行当前层的前向传播
                layer_outputs = layer_module(
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                )

            # 在每一层计算完成后，更新 hidden_states 为当前层的输出。
            hidden_states = layer_outputs[0]

            # 如果使用缓存，将当前层的缓存结果添加到 next_decoder_cache 中
            if use_cache:
                next_decoder_cache += (layer_outputs[-1],)

            # 如果需要输出注意力权重，将当前层的自注意力权重添加到 all_self_attentions 中
            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)
                # 如果配置中启用了交叉注意力，将当前层的交叉注意力权重添加到 all_cross_attentions 中
                if self.config.add_cross_attention:
                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)

        # -------------------------------------BertLayer堆叠完成---------------------------------------

        """ 4. 处理输出 """
        # 如果需要输出所有隐藏状态，将最终的 hidden_states 添加到 all_hidden_states 中。
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        # 如果 return_dict 为 False，则返回一个包含多个值的元组，这些值中只包含非 None 的值
        if not return_dict:
            return tuple(
                v
                for v in [
                    hidden_states,  # 隐藏状态
                    next_decoder_cache,  # 下一步解码器缓存
                    all_hidden_states,  # 所有隐藏状态
                    all_self_attentions,  # 所有自注意力
                    all_cross_attentions,  # 所有交叉注意力
                ]
                if v is not None
            )

        # 如果 return_dict 为 True，则返回一个包含详细信息的模型输出对象
        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,  # 最后的隐藏状态
            past_key_values=next_decoder_cache,  # 过去的键值对
            hidden_states=all_hidden_states,  # 所有隐藏状态
            attentions=all_self_attentions,  # 自注意力
            cross_attentions=all_cross_attentions,  # 交叉注意力
        )