基于HuggingFace Transformers的大规模DNA序列嵌入提取优化求助

核心方案可行性确认 #

你提到的流式完成tokenize并逐批提取嵌入、不保存中间token文件的方案完全可行，且能从根源上避免中间分片占用磁盘与内存，更适配6000万+行的超大规模数据集，同时保证最终输出包含Chromosome、label、embeddings三列。

具体实现代码 #

方案一：Pandas分块流式处理（直接从原始数据端到端生成嵌入） #

该方案无需加载全量数据到内存，按固定块大小分批处理，处理完一批立即释放内存并追加写入结果：

import pandas as pd
import torch
import gc
from transformers import AutoTokenizer, AutoModel

# 初始化模型与tokenizer
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()

# 配置参数
chunk_size = 2000  # 可根据内存/显存调整，显存不足则调小
output_columns = ['Chromosome', 'label', 'embeddings']
output_path = "final_embeddings.parquet"

# 初始化输出文件（写入表头）
pd.DataFrame(columns=output_columns).to_parquet(output_path, index=False)

# 分块遍历原始数据（替换为你的数据源路径，比如原element_final对应的存储文件）
for chunk in pd.read_csv("your_large_dna_dataset.csv", chunksize=chunk_size):
    # 保留必要列并过滤空值
    chunk = chunk[['Chromosome', 'sequence', 'label']].dropna()
    
    # 1. 对当前批次序列做tokenize
    tokenized = tokenizer.batch_encode_plus(
        chunk["sequence"].tolist(),
        return_tensors="pt",
        truncation=False,
        padding=False,
        max_length=80
    )
    input_ids = tokenized["input_ids"].to(device)
    
    # 2. 提取嵌入向量
    with torch.no_grad():
        outputs = model(input_ids=input_ids, output_hidden_states=True)
    embeddings = outputs.hidden_states[-1]
    averaged_embeddings = torch.mean(embeddings, dim=1).to(torch.float32).cpu().numpy()
    
    # 3. 整理结果并追加写入输出文件
    chunk['embeddings'] = list(averaged_embeddings)
    result_chunk = chunk[output_columns]
    result_chunk.to_parquet(
        output_path,
        mode='a',
        index=False,
        engine='pyarrow'
    )
    
    # 强制清理内存，避免泄漏
    del tokenized, input_ids, outputs, embeddings, averaged_embeddings, chunk, result_chunk
    torch.cuda.empty_cache()
    gc.collect()

方案二：用IterableDataset流式处理（适配Hugging Face数据集格式） #

如果原始数据已转为Hugging Face Dataset，可转为流式数据集处理：

import pandas as pd
import torch
import gc
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModel

# 初始化模型与tokenizer
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()

# 配置参数
batch_size = 2000
output_columns = ['Chromosome', 'label', 'embeddings']
output_path = "final_embeddings.parquet"

# 加载为流式数据集（避免全量加载）
dataset = load_dataset("csv", data_files="your_large_dna_dataset.csv", streaming=True)['train']
dataset = dataset.select_columns(['Chromosome', 'sequence', 'label'])

# 初始化输出文件
pd.DataFrame(columns=output_columns).to_parquet(output_path, index=False)

# 逐批处理流式数据
for batch in dataset.iter(batch_size=batch_size):
    # Tokenize当前批次
    tokenized = tokenizer.batch_encode_plus(
        batch["sequence"],
        return_tensors="pt",
        truncation=False,
        padding=False,
        max_length=80
    )
    input_ids = tokenized["input_ids"].to(device)
    
    # 提取嵌入
    with torch.no_grad():
        outputs = model(input_ids=input_ids, output_hidden_states=True)
    embeddings = outputs.hidden_states[-1]
    averaged_embeddings = torch.mean(embeddings, dim=1).to(torch.float32).cpu().numpy()
    
    # 整理结果并追加写入
    result_df = pd.DataFrame({
        'Chromosome': batch['Chromosome'],
        'label': batch['label'],
        'embeddings': list(averaged_embeddings)
    })
    result_df.to_parquet(
        output_path,
        mode='a',
        index=False,
        engine='pyarrow'
    )
    
    # 清理内存
    del tokenized, input_ids, outputs, embeddings, averaged_embeddings, batch, result_df
    torch.cuda.empty_cache()
    gc.collect()

关键优化点 #

• 批次大小适配：根据GPU显存和内存情况调整chunk_size/batch_size，显存不足时可降至1000甚至更低
• 内存强制清理：每批处理后手动删除临时变量+清理CUDA缓存，避免内存累积泄漏
• 高效存储格式：用Parquet格式存储嵌入，比CSV占用空间更小，且支持高效追加写入与后续读取
• 冗余数据剔除：全程仅保留必要列，处理完的临时数据立即删除，减少内存占用

原分片方案的兼容修复（可选） #

如果仍想保留分片模式，可将原200分片拆分为更多小分片（比如1000个），同时调小单批处理的batch_size，避免加载分片时内存溢出：

# Step1 修改分片数量
tokenized_dataset.save_to_disk(f"tokenized_elements/tokenized_{ELEMENT}", num_shards=1000)

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