使用BERT/BART处理超长句生成向量的报错与内存问题求助

错误原因 #

报错ValueError: too many values to unpack (expected 2)核心原因是输入张量维度错误：

• tokenizer.encode_plus返回的input_ids是[1, seq_len]的二维张量，直接存入列表后，循环中len(input_ids)取的是第一维度长度（固定为1），导致窗口计算完全失效。
• 后续对窗口切片后的张量执行unsqueeze(0)，会把维度变成[1,1, window_len]，而模型期望的输入格式是[batch_size, seq_len]，因此input_shape变成三维，解包batch_size和seq_len时触发错误。

修改方案及代码 #

关键修改点： #

1. 存储input_ids和attention_mask时，去掉多余的batch维度（从二维转一维）。
2. 修正窗口数量计算逻辑，避免生成无效窗口。
3. 窗口输入直接使用一维张量，无需额外unsqueeze(0)（模型会自动处理单样本的batch维度）。

from transformers import AutoModel, AutoTokenizer
import torch
import numpy as np

tokenizer = AutoTokenizer.from_pretrained("sagorsarker/bangla-bert-base")
model = AutoModel.from_pretrained("sagorsarker/bangla-bert-base")
sentences = data['Text_stemmed_without_no']

input_ids_list = []
attention_mask_list = []
for sentence in sentences:
    encoded_dict = tokenizer.encode_plus(
                        sentence,
                        add_special_tokens=True,
                        return_attention_mask=True,
                        return_tensors='pt'
                   )
    # 关键：去掉多余的batch维度，转成一维张量
    input_ids_list.append(encoded_dict['input_ids'].squeeze(0))
    attention_mask_list.append(encoded_dict['attention_mask'].squeeze(0))

max_length = 512
stride = 256
bert_embeddings_list = []

for input_ids, attention_mask in zip(input_ids_list, attention_mask_list):
    seq_len = input_ids.size(0)
    # 修正窗口数量计算：确保覆盖整个序列且不生成无效窗口
    num_windows = max(1, (seq_len - max_length + stride - 1) // stride + 1)
    embeddings = []
    
    for i in range(num_windows):
        start_index = i * stride
        end_index = min(start_index + max_length, seq_len)
        window_input_ids = input_ids[start_index:end_index]
        window_attention_mask = attention_mask[start_index:end_index]
        
        # 直接传入一维张量，模型会自动添加batch维度
        with torch.no_grad():
            last_hidden_states = model(
                input_ids=window_input_ids.unsqueeze(0), 
                attention_mask=window_attention_mask.unsqueeze(0)
            )
        # 取<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]> token的嵌入作为窗口表征
        window_emb = last_hidden_states[0][:, 0, :].numpy()
        embeddings.append(window_emb)
    
    # 可选：若需单句统一向量，推荐取所有窗口embedding的均值：np.mean(embeddings, axis=0)
    merged_emb = np.concatenate(embeddings, axis=0)
    bert_embeddings_list.append(merged_emb)

bert_embeddings = np.array(bert_embeddings_list)

错误原因 #

一次性将200个长句子的编码张量拼接成大张量，加上bart-large模型本身参数规模极大（约1.4亿参数），导致显存/内存占用超出Kaggle会话的限制阈值。

修改方案及代码 #

关键修改点： #

1. 逐个处理句子，计算完单个句子的embedding后立即转存为numpy，释放PyTorch张量占用的内存。
2. 避免一次性存储所有编码后的张量，减少内存累积。

from transformers import BartTokenizer, BartModel
import torch
import numpy as np

# 加载模型和分词器
tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')
model = BartModel.from_pretrained('facebook/bart-large')
sentences = data['Text_stemmed_without_no']

bart_embeddings_list = []

# 逐个处理句子，避免内存累积
for sentence in sentences:
    encoded_dict = tokenizer.encode_plus(
                        sentence,
                        add_special_tokens=True,
                        max_length=600,
                        pad_to_max_length=True,
                        return_attention_mask=True,
                        return_tensors='pt'
                   )
    input_ids = encoded_dict['input_ids']
    attention_mask = encoded_dict['attention_mask']
    
    # 计算embedding并释放张量内存
    with torch.no_grad():
        last_hidden_states = model(input_ids=input_ids, attention_mask=attention_mask)
    # 取<[BOS_never_used_51bce0c785ca2f68081bfa7d91973934]> token的嵌入作为句子表征
    sent_emb = last_hidden_states[0][:, 0, :].numpy()
    bart_embeddings_list.append(sent_emb)
    
    # 手动清理当前句子的张量，释放内存（显式清理更稳妥）
    del input_ids, attention_mask, last_hidden_states
    torch.cuda.empty_cache()  # 使用GPU时，添加此句释放显存

# 转换为最终的numpy数组
bart_embeddings = np.concatenate(bart_embeddings_list, axis=0)

额外优化建议： #

• 若仍出现内存问题，可缩小max_length（比如设为512），或改用bart-base替代bart-large（参数规模仅为前者的1/3）。
• torch.no_grad()必须开启，否则会存储梯度信息占用额外内存。

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