如何将PyTorch版基于全连接的BERT二分类器改造为前置Conv1D、最大池化、展平层的结构

核心实现要点 #

• PyTorch内置的nn.Conv1d要求输入张量格式为(批次大小, 通道数, 序列长度)，原有BERT输出的pooled向量形状为(批次大小, 768)，属于无通道维度的二维张量，无法直接传入卷积层，需要先通过扩维操作调整形状适配输入要求。
• 严格按照需求顺序搭建模块链路：BERT主干输出 -> Dropout层 -> Conv1D层 -> 最大池化层 -> 展平层 -> 全连接分类层 -> Sigmoid输出。
• 全连接层的输入维度必须与展平后的输出维度严格匹配，否则会触发维度不匹配的运行时报错。

改造后完整代码 #

import torch
import torch.nn as nn
from transformers import BertModel

class BertBinaryClassifier(nn.Module):
    def __init__(self, dropout=0.1, conv_out_channels=128, conv_kernel_size=3, pool_kernel_size=2):
        super(BertBinaryClassifier, self).__init__()
        # 保留原有BERT主干结构
        self.bert = BertModel.from_pretrained('bert-base-uncased')
        self.dropout = nn.Dropout(dropout)
        
        # 新增卷积、池化、展平模块
        self.conv1d = nn.Conv1d(
            in_channels=1,
            out_channels=conv_out_channels,
            kernel_size=conv_kernel_size,
            padding=conv_kernel_size // 2  # 保持序列长度一致，简化维度计算
        )
        self.maxpool = nn.MaxPool1d(kernel_size=pool_kernel_size)
        self.flatten = nn.Flatten()
        
        # 自动计算展平后的张量维度，匹配全连接层输入要求
        # 池化后剩余序列长度 = 原BERT输出的768维 / 池化核大小
        pooled_seq_len = 768 // pool_kernel_size
        flatten_dim = conv_out_channels * pooled_seq_len
        self.linear = nn.Linear(flatten_dim, 1)
        self.sigmoid = nn.Sigmoid()
        
    def forward(self, tokens, masks=None):
        # 提取BERT输出的pooled语义向量，形状为(batch_size, 768)
        _, pooled_output = self.bert(tokens, attention_mask=masks, output_all_encoded_layers=False)
        dropout_output = self.dropout(pooled_output)
        
        # 扩维适配Conv1D输入：(batch_size,768) -> (batch_size,1,768)
        conv_input = dropout_output.unsqueeze(1)
        # 依次经过卷积、最大池化、展平操作
        conv_output = self.conv1d(conv_input)
        pool_output = self.maxpool(conv_output)
        flatten_output = self.flatten(pool_output)
        
        # 接入全连接层完成最终分类
        linear_output = self.linear(flatten_output)
        prediction = self.sigmoid(linear_output)
        return prediction

# 原有配置保持不变
BATCH_SIZE = 4
EPOCHS = 5

调优提示 #

• 卷积层的输出通道数、卷积核大小、池化核大小都可以根据实际任务效果灵活调整，示例中的参数值仅作参考。
• 如果出现维度不匹配报错，可以在forward函数中逐次打印各层输出的张量形状，快速定位维度偏差点。
• 若展平后全连接层参数量过大，可在展平层和最终分类全连接层之间新增一层降维线性层，减少计算开销。

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