大DNA数据集滑动窗口序列获取崩溃问题的优化问询

问题背景 #

原始数据集仅包含DNA序列坐标，示例如下：

chr   start    stop   label  
chr1   9000    9100    1  
chr1   8803    8903     1  
chr1   8903    9000     0  

目标是为每个坐标生成±5000nt的滑动窗口上下文，预期输出：

chr    start         stop        label  
chr1   9000-5000    9000+5000      1  
chr1   9001-5000    9001+5000      1  
...
chr1   9100-5000    9100+5000      1  

现有6步流程（坐标扩展→序列获取→分词→嵌入→XGBoost）存在原始700K行扩展至10亿行，内存溢出崩溃的问题，核心原因是全量存储中间结果、未做流式/批处理优化。

优化方案 #

1. 用生成器替代全量存储，实现流式坐标扩展 #

原expand_coordinates函数将所有扩展坐标存入列表再转DataFrame，直接耗尽内存。改为生成器逐行产出窗口坐标，不占用全量内存：

def generate_window_coords(element_locs, context=5000):
    for _, row in element_locs.iterrows():
        chr_name = row['Chromosome']
        chr_start = int(row['Start'])
        chr_end = int(row['End'])
        chr_max_len = max_sizes[chr_name]
        
        for i in range(chr_start, chr_end + 1):
            window_start = max((i - 1) - context, 0)
            window_end = min(i + context, chr_max_len)
            yield (chr_name, window_start, window_end, row['label'])

def fetch_sequences_stream(genome, coords_generator):
    for chr_name, start, end, label in coords_generator:
        seq = genome.fetch(chr_name, start, end)
        yield {'Chromosome': chr_name, 'sequence': seq, 'label': label}

使用方式：

genome = pysam.Fastafile(ref_genome)
coords_gen = generate_window_coords(original_df, context=5000)
seq_stream = fetch_sequences_stream(genome, coords_gen)

2. 对接Hugging Face流式数据集，避免全量加载 #

将生成器转为IterableDataset，后续分词、嵌入全程流式处理：

from datasets import IterableDataset

stream_dataset = IterableDataset.from_generator(
    lambda: seq_stream,
    features=Dataset.from_pandas(original_df[['Chromosome', 'sequence', 'label']]).features
)

3. 流式批处理：分词+嵌入+保存一体化 #

跳过中间token存储步骤，每处理一批序列直接生成嵌入并保存，减少内存占用：

import torch
from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-human-ref")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AutoModel.from_pretrained("InstaDeepAI/nucleotide-transformer-500m-human-ref").to(device)

def process_batch(batch):
    # 分词
    tokenized = tokenizer.batch_encode_plus(
        batch["sequence"], 
        return_tensors="pt", 
        truncation=False, 
        padding=False, 
        max_length=10001
    )
    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].mean(dim=1).to(torch.float32)
    
    return {
        'Chromosome': batch['Chromosome'],
        'label': batch['label'],
        'embeddings': embeddings.cpu().numpy()
    }

# 流式处理并分块保存
output_dataset = stream_dataset.map(
    process_batch,
    batched=True,
    batch_size=1000,  # 根据GPU内存调整
    remove_columns=["sequence"]
)
output_dataset.save_to_disk(f"embedded_elements/embeddings_{ELEMENT_LABEL}", num_shards=100)

4. 可选：窗口重叠优化，减少重复序列读取 #

针对连续区间的原始坐标，先提取完整长序列再滑动生成子窗口，避免重复从genome中fetch：

def generate_window_seqs_from_full(element_locs, context=5000):
    genome = pysam.Fastafile(ref_genome)
    for _, row in element_locs.iterrows():
        chr_name = row['Chromosome']
        # 提取原始区间±context的完整长序列
        full_start = max(row['Start'] - context, 0)
        full_end = min(row['End'] + context, max_sizes[chr_name])
        full_seq = genome.fetch(chr_name, full_start, full_end)
        
        # 在长序列上滑动生成窗口
        offset = context - (row['Start'] - full_start)
        for i in range(row['Start'], row['End'] + 1):
            window_start_idx = offset + (i - row['Start']) - context
            window_end_idx = window_start_idx + 2*context + 1
            window_seq = full_seq[window_start_idx:window_end_idx]
            yield {'Chromosome': chr_name, 'sequence': window_seq, 'label': row['label']}

stream_dataset = IterableDataset.from_generator(
    lambda: generate_window_seqs_from_full(original_df, context=5000),
    features=Dataset.from_pandas(original_df[['Chromosome', 'sequence', 'label']]).features
)

5. XGBoost对接优化：流式加载嵌入 #

用DMatrix流式迭代加载分块保存的嵌入，避免全量读入内存：

import xgboost as xgb
from datasets import load_from_disk

embedded_dataset = load_from_disk(f"embedded_elements/embeddings_{ELEMENT_LABEL}")

def dmatrix_generator():
    for batch in embedded_dataset.iter(batch_size=10000):
        yield xgb.DMatrix(batch['embeddings'], label=batch['label'])

model = xgb.XGBClassifier()
model.fit(dmatrix_generator())

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