如何使用Huggingface BERT输出作为输入搭建二分类CNN模型

前置处理优化 #

你当前逐句调用tokenizer的方式仅适合单样本推理，实际训练时可以在调用tokenizer时添加padding='max_length'、truncation=True、max_length=32参数，直接将所有输入补全/截断到统一长度，此时BERT输出的序列维度固定，是最通用的训练方案。
如果确实需要支持完全无截断的可变长度输入，CNN本身的卷积+池化结构天然可以适配，不需要固定序列长度。

适配可变长度输入的CNN实现代码 #

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModel

class BertEinsteinClassifier(nn.Module):
    def __init__(self, embed_dim=768, num_classes=2, filter_sizes=[2,3,4], num_filters=100, dropout=0.5):
        super().__init__()
        # 定义不同尺寸的卷积核，提取不同粒度的文本特征
        self.convs = nn.ModuleList([
            nn.Conv2d(in_channels=1, out_channels=num_filters, kernel_size=(fs, embed_dim)) 
            for fs in filter_sizes
        ])
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(len(filter_sizes) * num_filters, num_classes)

    def forward(self, bert_last_hidden_state):
        # bert_last_hidden_state输入形状：[batch_size, seq_len, embed_dim]，seq_len为任意长度
        # 新增通道维度适配Conv2d输入要求：[batch_size, 1, seq_len, embed_dim]
        x = bert_last_hidden_state.unsqueeze(1)
        
        conv_results = []
        for conv in self.convs:
            # 卷积后过ReLU激活，去掉最后一个冗余维度
            conv_out = F.relu(conv(x)).squeeze(3)
            # *最大池化是适配可变长度的核心：无论序列长度多少，每个卷积核在整个序列维度取最大值，得到固定长度输出*
            pool_out = F.max_pool1d(conv_out, conv_out.size(2)).squeeze(2)
            conv_results.append(pool_out)
        
        # 拼接所有卷积核的特征，得到固定长度的全局特征向量
        concat_feature = torch.cat(conv_results, dim=1)
        concat_feature = self.dropout(concat_feature)
        # 输出二分类logits
        return self.classifier(concat_feature)

调用测试示例 #

# 初始化组件
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
bert_model = AutoModel.from_pretrained("bert-base-uncased")
classifier = BertEinsteinClassifier()

# 测试输入
test_texts = [
    "Hello World", 
    "Two things are infinite: the universe and human stupidity; and I'm not sure about the universe."
]

bert_model.eval()
classifier.eval()
with torch.no_grad():
    for text in test_texts:
        model_input = tokenizer(text, return_tensors="pt")
        bert_out = bert_model(**model_input).last_hidden_state
        logits = classifier(bert_out)
        # 预测结果：1=爱因斯坦所说，0=非爱因斯坦所说
        pred = torch.argmax(logits, dim=1).item()
        print(f"文本：{text}\n预测结果：{pred}\n")

训练注意事项 #

• 批量训练时建议统一输入长度，可大幅提升训练效率，避免逐句计算的性能损耗
• 可以选择冻结BERT参数只训练CNN分类头，或者微调BERT+CNN整体参数，根据你的数据集规模决定

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