PyTorch中完整前向传播f_x1能否拆分为两步执行？

你的目标是把BERT的完整前向传播拆成两步：先拿到第k层隐藏态h_x1，再基于h_x1继续计算后面的层，让最终输出f_x2和完整前向的f_x1近似相等。但当前代码存在逻辑错误，导致达不到预期，下面具体分析并给出修正方案：

核心问题所在 #

当你把h_x1作为inputs_embeds传入BERT时，BERT会从头执行全部12层Transformer，而不是从第k层之后继续计算。举个例子，k=3时，完整前向是「embedding→层1→层2→层3→…→层12」，但你当前的第二步是把层3的结果当输入，重新跑一遍「层1→层2→…→层12」，得到的结果和原本的层12输出完全无关，自然f_x1和f_x2不会相等。

正确实现方式 #

要实现从第k层继续计算，需要把BERT的Transformer层拆成前k层和后12-k层，分别处理完整前向和拆分后的第二步计算。修改后的模型结构和前向逻辑如下：

def __init__(self, bert_model_name='bert-base-uncased', num_classes=2):
    super().__init__()
    # 加载预训练BERT
    self.bert = BertModel.from_pretrained(bert_model_name)
    self.k = 3  # 可根据需求调整，也可作为参数传入
    # 拆分Transformer编码器层
    self.bert_first_k_layers = nn.ModuleList(self.bert.encoder.layer[:self.k])
    self.bert_last_12mk_layers = nn.ModuleList(self.bert.encoder.layer[self.k:])
    # 保留BERT的嵌入层和池化层
    self.bert_embeddings = self.bert.embeddings
    self.bert_pooler = self.bert.pooler
    
    # 自定义后续全连接层
    self.dropout = nn.Dropout(0.1)
    self.linear = nn.Linear(768, 768)
    self.out = nn.Linear(768, num_classes)

def forward(self, x_input_ids=None, x_seg_ids=None, x_atten_masks=None, inputs_embeds=None, get_hidden=False):
    if inputs_embeds is not None:
        # 拆分第二步：从第k层开始继续计算到第12层
        hidden_states = inputs_embeds
        # 遍历后12-k层Transformer
        for layer in self.bert_last_12mk_layers:
            hidden_states = layer(hidden_states, attention_mask=x_atten_masks)[0]
        # 取<s> token的表示（和完整前向逻辑一致）
        query = hidden_states[:, 0]
        # 也可使用BERT原生池化层：query = self.bert_pooler(hidden_states)
    else:
        # 完整前向：从嵌入层开始跑完整12层
        embedding_output = self.bert_embeddings(input_ids=x_input_ids, token_type_ids=x_seg_ids)
        hidden_states = embedding_output
        # 先跑前k层，拿到第k层隐藏态h_x1
        for layer in self.bert_first_k_layers:
            hidden_states = layer(hidden_states, attention_mask=x_atten_masks)[0]
        h_x1 = hidden_states.clone()
        # 再跑后12-k层，得到最终隐藏态
        for layer in self.bert_last_12mk_layers:
            hidden_states = layer(hidden_states, attention_mask=x_atten_masks)[0]
        query = hidden_states[:, 0]
    
    # 后续全连接层计算
    query = self.dropout(query)
    linear = self.relu(self.linear(query))
    out = self.out(linear)
    
    if get_hidden:
        return out, h_x1
    else:
        return out

验证与训练注意事项 #

• 修正后，拆分两步的计算路径和完整前向完全一致，在model.eval()模式下（关闭dropout等随机操作），f_x1和f_x2应该完全相等；训练模式下因dropout随机性会有微小差异，属于正常现象。
• 训练逻辑可以保留原有的MSELoss设置，只要模型结构正确，损失会快速收敛到接近0的水平。

内容的提问来源于stack exchange，提问作者kkgarg