license: cc-by-nc-sa-4.0
widget:
- text: ACCTGATTCTGAGTC
tags:
- DNA
- 生物学
- 基因组学
- 分割
SegmentNT
SegmentNT是一个基于Nucleotide Transformer(NT)DNA基础模型的分割模型,能够以单核苷酸分辨率预测序列中多种基因组元素的位置。该模型在长达30kb的人类基因组输入序列上训练,覆盖14种不同类型的基因组元素,包括基因(蛋白质编码基因、长链非编码RNA、5'UTR、3'UTR、外显子、内含子、剪接受体和供体位点)和调控元件(polyA信号、组织不变和组织特异性启动子与增强子,以及CTCF结合位点)。
开发团队: InstaDeep
模型来源
使用方法
在下一个版本发布之前,需要通过以下命令从源代码安装transformers
库以使用这些模型:
pip install --upgrade git+https://github.com/huggingface/transformers.git
这里提供了一个小代码片段,用于从虚拟DNA序列中获取logits和嵌入。
⚠️ 默认情况下,最大序列长度设置为训练时的30,000个核苷酸,或5001个token(包括CLS token)。然而,SegmentNT-multi-species已证明可以推广到50,000 bp的序列。如果您需要在30kbp到50kbp之间的序列上进行推理,请确保更改Rotary Embedding层的rescaling_factor
,即num_dna_tokens_inference / max_num_tokens_nt
,其中num_dna_tokens_inference
是推理时的token数量(例如,40008个碱基对的序列为6669个token),而max_num_tokens_nt
是骨干nucleotide-transformer训练时的最大token数,即2048
。
./inference_segment_nt.ipynb
可以在Google Colab中运行,点击图标即可。该笔记本展示了如何处理需要更改重缩放因子的序列长度以及不需要更改的序列长度。您可以运行该笔记本并重现SegmentNT论文中的图1和图3。
from transformers import AutoTokenizer, AutoModel
import torch
tokenizer = AutoTokenizer.from_pretrained("InstaDeepAI/segment_nt", trust_remote_code=True)
model = AutoModel.from_pretrained("InstaDeepAI/segment_nt", trust_remote_code=True)
max_length = 12 + 1
assert (max_length - 1) % 4 == 0, (
"DNA token的数量(不包括预加的CLS token)需要能被下采样块的数量(即4)的2次幂整除。")
sequences = ["ATTCCGATTCCGATTCCG", "ATTTCTCTCTCTCTCTGAGATCGATCGATCGAT"]
tokens = tokenizer.batch_encode_plus(sequences, return_tensors="pt", padding="max_length", max_length = max_length)["input_ids"]
attention_mask = tokens != tokenizer.pad_token_id
outs = model(
tokens,
attention_mask=attention_mask,
output_hidden_states=True
)
logits = outs.logits.detach()
probabilities = torch.nn.functional.softmax(logits, dim=-1)
print(f"概率形状: {probabilities.shape}")
idx_intron = model.config.features.index("intron")
probabilities_intron = probabilities[:,:,idx_intron]
print(f"内含子概率形状: {probabilities_intron.shape}")
训练数据
SegmentNT模型在除20号和21号染色体(作为测试集保留)以及22号染色体(用作验证集)外的所有人类染色体上进行了训练。训练过程中,序列在基因组中随机采样并附带注释。然而,我们通过使用长度为30,000的滑动窗口在20号和21号染色体上固定验证集和测试集中的序列。验证集用于监控训练和早期停止。
训练过程
预处理
DNA序列使用Nucleotide Transformer分词器进行分词,该分词器将序列分词为6-mer token,如相关代码库的分词部分所述。该分词器的词汇量为4105。模型的输入形式如下:
<CLS> <ACGTGT> <ACGTGC> <ACGGAC> <GACTAG> <TCAGCA>
训练
模型在配备8个GPU的DGXH100节点上训练,共处理了23B个token,耗时3天。模型分别在3kb、10kb、20kb和30kb的序列上进行训练,每次训练的有效批次大小为256个序列。
架构
该模型由nucleotide-transformer-v2-500m-multi-species编码器组成,我们移除了其语言模型头部,并用一个由2个下采样卷积块和2个上采样卷积块组成的1维U-Net分割头部[4]替换。每个块由2个卷积层组成,分别具有1,024和2,048个核。这个额外的分割头部包含5300万个参数,使总参数数达到562M。
BibTeX条目和引用信息
@article{de2024segmentnt,
title={SegmentNT: 基于DNA基础模型的单核苷酸分辨率基因组注释},
author={de Almeida, Bernardo P and Dalla-Torre, Hugo and Richard, Guillaume and Blum, Christopher and Hexemer, Lorenz and Gelard, Maxence and Pandey, Priyanka and Laurent, Stefan and Laterre, Alexandre and Lang, Maren and others},
journal={bioRxiv},
pages={2024--03},
year={2024},
publisher={Cold Spring Harbor Laboratory}
}