如何在Bert序列分类任务中使用大于1的batch size处理大批量数据？

• 分词器和模型每次调用函数都重复加载，资源浪费极其严重，处理几万条数据会严重超时
• 单条样本编码+推理，没有利用批处理并行优势，效率极低
• 没有做统一长度的截断和填充，无法批量输入模型

import tensorflow as tf
from transformers import BertTokenizer, TFBertModel
import pandas as pd

# 全局仅加载一次分词器和模型，避免重复开销
MODEL_NAME = 'bert-base-uncased'
MAX_SEQ_LEN = 128 # 可根据你的文本平均长度调整，最大不超过512
BATCH_SIZE = 64 # 可根据显存大小调整，显存充足可设为128/256，不足可降到32/16

tokenizer = BertTokenizer.from_pretrained(MODEL_NAME)
model = TFBertModel.from_pretrained(MODEL_NAME)

# 加载CSV数据集，替换为你实际的文件路径和文本列名
df = pd.read_csv("你的数据集路径.csv")
text_list = df['text'].tolist() # 如果你已经有了文本列表z，直接替换为z即可

# 批量编码所有文本，统一长度
encoded_input = tokenizer.batch_encode_plus(
    text_list,
    add_special_tokens=True,
    max_length=MAX_SEQ_LEN,
    padding='max_length',
    truncation=True,
    return_tensors='tf'
)
input_ids = encoded_input['input_ids']
attention_mask = encoded_input['attention_mask'] # 加入attention_mask避免padding部分参与模型计算

# 构建批处理数据集，分批推理避免显存溢出
dataset = tf.data.Dataset.from_tensor_slices((input_ids, attention_mask)).batch(BATCH_SIZE)

# 批量提取<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]> embedding
cls_embeddings = []
for batch_input_ids, batch_attention_mask in dataset:
    outputs = model(batch_input_ids, attention_mask=batch_attention_mask)
    last_hidden_states = outputs[0]
    # 取每个样本的第一个token即<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]>，单批次输出shape为 (batch_size, hidden_size)
    batch_cls = last_hidden_states[:, 0, :]
    cls_embeddings.append(batch_cls)

# 拼接所有批次结果，最终输出shape为 (48000, hidden_size)的Tensor
final_cls_tensor = tf.concat(cls_embeddings, axis=0)

如果数据量太大内存无法一次性装下所有embedding，可以在批量生成时同步将结果写入本地.npy文件，不需要全部暂存在内存中。

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