在PyTorch/HuggingFace中映射Embedding到标签遇索引错误求助

问题背景 #

使用预训练的InstaDeepAI/nucleotide-transformer-500m-human-ref模型，为带0/1标签的2253万条DNA序列生成Embedding。因内存限制采用批处理，但出现IndexError，无法将生成的Embedding与对应标签正确映射。当前仅测试了前4个样本，得到的Embedding形状为torch.Size([4, 86, 1280])，与全量数据集的标签维度不匹配，需为每条序列生成对应Embedding用于下游任务，且已完成训练/测试/验证集拆分。

当前代码 #

from datasets import Dataset
from transformers import AutoTokenizer, AutoModel
import torch

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-human-ref")
model = AutoModel.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-human-ref")

# Load the dataset
ds1 = Dataset.from_file('training.arrow') #this is already tokenized

# Convert tokenized sequences to tensor
inputs = torch.tensor(ds1['input_ids']).to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))

# Reduce batch size
batch_size = 4

# Pass tokenized sequences through the model with reduced batch size
with torch.no_grad():
    outputs = model(input_ids=inputs[:batch_size], output_hidden_states=True)

# Extract embeddings
hidden_states = outputs.hidden_states
embeddings1 = hidden_states[-1]

尺寸信息 #

embeddings1.shape
torch.Size([4, 86, 1280])

ds1
Dataset({
    features: ['labels', 'input_ids', 'attention_mask'],
    num_rows: 22535512
})

解决方案 #

1. 无需按标签拆分生成Embedding #

按标签拆分再合并的做法会增加流程复杂度，还可能打乱已划分好的训练/测试/验证集，完全没有必要。正确做法是直接在原数据集上批量处理所有样本，同时保留每个样本的标签关联。

2. 全量批处理实现（保留标签映射） #

利用Hugging Face Dataset的map函数自动分批次处理，既避免内存溢出，又保证Embedding与标签一一对应。同时需对模型输出的序列级Embedding做池化，得到单条序列的向量（适配下游任务需求）：

from datasets import Dataset
from transformers import AutoTokenizer, AutoModel
import torch

# 配置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 加载模型与分词器
tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-human-ref")
model = AutoModel.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-human-ref").to(device)
model.eval() # 切换到评估模式

# 加载数据集
ds1 = Dataset.from_file('training.arrow')
batch_size = 4 # 根据内存调整

# 定义批处理函数：生成单条序列的Embedding
def generate_sequence_embedding(batch):
    # 将输入移到设备上
    input_ids = torch.tensor(batch['input_ids']).to(device)
    attention_mask = torch.tensor(batch['attention_mask']).to(device)
    
    with torch.no_grad(): # 关闭梯度计算，节省内存
        outputs = model(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True)
    
    # 取最后一层隐藏状态，用attention mask做平均池化（避免padding干扰）
    last_hidden = outputs.hidden_states[-1]
    mask_expanded = attention_mask.unsqueeze(-1).expand(last_hidden.size())
    pooled_embedding = (last_hidden * mask_expanded).sum(1) / mask_expanded.sum(1)
    
    # 返回numpy格式的Embedding，方便存入Dataset
    return {'embeddings': pooled_embedding.cpu().numpy()}

# 分批次处理整个数据集，自动保留原标签
ds_with_embeddings = ds1.map(
    generate_sequence_embedding,
    batched=True,
    batch_size=batch_size,
    num_proc=4 # 可根据CPU核心数调整，加速处理
)

# 查看处理后的数据集：每个样本包含labels、input_ids、attention_mask、embeddings
print(ds_with_embeddings)

# 保存处理后的数据集，方便下游任务直接使用
ds_with_embeddings.save_to_disk('training_with_embeddings')

3. 关键说明 #

• 池化操作：模型输出的hidden_states[-1]是[batch_size, seq_len, hidden_dim]的序列级Embedding，下游任务需要的是单条序列的向量，因此用带attention mask的平均池化过滤padding部分的影响。
• 内存优化：使用torch.no_grad()关闭梯度计算，避免不必要的内存占用；map函数自动分批次加载数据，不会一次性把全量数据读入内存。
• 标签映射保证：map函数会严格按照原数据集的样本顺序处理，生成的embeddings特征与原labels一一对应，不会出现索引错位问题。

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