GENA-LM
GENA-LM 是一系列用于长DNA序列的开源基础模型。
GENA-LM 模型是在人类DNA序列上训练的transformer掩码语言模型。
我们的GENA-LM模型的主要特点:
- 使用BPE分词而非k-mers(DNABERT、Nucleotide Transformer)
- 最大输入序列长度范围从4.5k到36k碱基对,相比之下DNABERT为512碱基对,Nucleotide Transformer为1000碱基对
- 在最新的T2T人类基因组组装上进行预训练,而非GRCh38/hg38
预训练模型
模型 | 架构 | 最大序列长度,标记(碱基对) | 参数数 | 分词器数据 | 训练数据 |
---|---|---|---|---|---|
bert-base | BERT-12L | 512(4500) | 110M | T2T 拆分 v1 | T2T 拆分 v1 |
bert-base-t2t | BERT-12L | 512(4500) | 110M | T2T+1000G SNPs+多物种 | T2T+1000G SNPs |
bert-base-lastln-t2t | BERT-12L | 512(4500) | 110M | T2T+1000G SNPs+多物种 | T2T+1000G SNPs |
bert-base-t2t-multi | BERT-12L | 512(4500) | 110M | T2T+1000G SNPs+多物种 | T2T+1000G SNPs+多物种 |
bert-large-t2t | BERT-24L | 512(4500) | 336M | T2T+1000G SNPs+多物种 | T2T+1000G SNPs |
bigbird-base-sparse | BERT-12L, DeepSpeed稀疏操作, RoPE | 4096(36000) | 110M | T2T 拆分 v1 | T2T 拆分 v1 |
bigbird-base-sparse-t2t | BERT-12L, DeepSpeed稀疏操作, RoPE | 4096(36000) | 110M | T2T+1000G SNPs+多物种 | T2T+1000G SNPs |
bigbird-base-t2t | BERT-12L, HF BigBird | 4096(36000) | 110M | T2T+1000G SNPs+多物种 | T2T+1000G SNPs |
T2T 拆分 v1 指的是使用未增强的T2T人类基因组组装拆分的初步模型。基于BERT的模型采用预层归一化,而lastln明确表示层归一化也应用于最后一层。RoPE表示使用旋转位置嵌入代替BERT类型的绝对位置嵌入。
对于我们的首批模型(gena-lm-bert-base
和gena-lm-bigbird-base-sparse
),我们保留人类染色体22和Y(CP068256.2和CP086569.2)作为掩码语言建模任务的测试数据集。对于所有其他模型,我们保留人类染色体7和10(CP068271.2和CP068268.2);这些模型的名称中带有后缀"t2t"。其他数据用于训练。仅人类模型在预处理的Human T2T v2基因组组装及其1000基因组SNP增强数据上进行训练,总计约480 x 10^9碱基对。多物种模型在仅人类和多物种数据上进行训练,总计约1072 x 10^9碱基对。
下游任务上的预训练模型
模型 | 任务 | 任务序列长度 | 指标 | HF 分支名称 |
---|---|---|---|---|
gena-lm-bert-base-t2t | 启动子 | 300bp | 74.56+-0.36 F1 | promoters_300_run_1 |
gena-lm-bert-large-t2t | 启动子 | 300bp | 76.44+-0.16 F1 | promoters_300_run_1 |
gena-lm-bert-large-t2t | 启动子 | 2000bp | 93.70+-0.44 F1 | promoters_2000_run_1 |
gena-lm-bert-base-t2t | 剪接位点 | 15000bp | 92.63+-0.09 PR AUC | spliceai_run_1 |
gena-lm-bert-large-t2t | 剪接位点 | 15000bp | 93.59+-0.11 PR AUC | spliceai_run_1 |
要获取下游任务上的预训练模型,请用表格中的值替换model_name
和branch_name
。表格中的指标是多次运行的平均值。因此,每个检查点的值可能与此处报告的有所不同。
from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained(f'AIRI-Institute/{model_name}')
model = AutoModel.from_pretrained(f'AIRI-Institute/{model_name}', revision=branch_name, trust_remote_code=True)
示例
如何加载预训练的GENA-LM用于掩码语言建模
from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained('AIRI-Institute/gena-lm-bert-base-t2t')
model = AutoModel.from_pretrained('AIRI-Institute/gena-lm-bert-base-t2t', trust_remote_code=True)
如何加载预训练的GENA-LM以在分类任务上进行微调
从GENA-LM仓库获取模型类:
git clone https://github.com/AIRI-Institute/GENA_LM.git
from GENA_LM.src.gena_lm.modeling_bert import BertForSequenceClassification
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained('AIRI-Institute/gena-lm-bert-base-t2t')
model = BertForSequenceClassification.from_pretrained('AIRI-Institute/gena-lm-bert-base-t2t')
或者你可以直接下载modeling_bert.py并将其放在你的代码附近。
或者你可以从HuggingFace AutoModel获取模型类:
from transformers import AutoTokenizer, AutoModel
model = AutoModel.from_pretrained('AIRI-Institute/gena-lm-bert-base-t2t', trust_remote_code=True)
gena_module_name = model.__class__.__module__
print(gena_module_name)
import importlib
# 可用的类名:
# - BertModel, BertForPreTraining, BertForMaskedLM, BertForNextSentencePrediction,
# - BertForSequenceClassification, BertForMultipleChoice, BertForTokenClassification,
# - BertForQuestionAnswering
# 查看 https://huggingface.co/docs/transformers/model_doc/bert
cls = getattr(importlib.import_module(gena_module_name), 'BertForSequenceClassification')
print(cls)
model = cls.from_pretrained('AIRI-Institute/gena-lm-bert-base-t2t', num_labels=2)
GENA-LM bigbird-base-t2t
模型使用HuggingFace BigBird实现。因此,可以使用Transformers库中的默认类:
from transformers import AutoTokenizer, BigBirdForSequenceClassification
tokenizer = AutoTokenizer.from_pretrained('AIRI-Institute/gena-lm-bigbird-base-t2t')
model = BigBirdForSequenceClassification.from_pretrained('AIRI-Institute/gena-lm-bigbird-base-t2t')
笔记本
引用
@article {GENA_LM,
author = {Veniamin Fishman and Yuri Kuratov and Maxim Petrov and Aleksei Shmelev and Denis Shepelin and Nikolay Chekanov and Olga Kardymon and Mikhail Burtsev},
title = {GENA-LM: A Family of Open-Source Foundational Models for Long DNA Sequences},
elocation-id = {2023.06.12.544594},
year = {2023},
doi = {10.1101/2023.06.12.544594},
publisher = {Cold Spring Harbor Laboratory},
URL = {https://www.biorxiv.org/content/early/2023/06/13/2023.06.12.544594},
eprint = {https://www.biorxiv.org/content/early/2023/06/13/2023.06.12.544594.full.pdf},
journal = {bioRxiv}
}
下游任务
用于模型评估的下游任务包括预测启动子和增强子活性、剪接位点、染色质谱和多聚腺苷酸化位点强度。
查看downstream_tasks
文件夹以获取我们使用的代码和数据预处理脚本:
- [启动子预测](https://github.com/AIRI-Institute/GENA_LM/blob/main/./downstream_tasks/promoter_prediction/
- [剪接位点预测 (SpliceAI)](https://github.com/AIRI-Institute/GENA_LM/blob/main/./downstream_tasks/SpliceAI/
- [果蝇增强子预测 (DeepSTARR)](https://github.com/AIRI-Institute/GENA_LM/blob/main/./downstream_tasks/DeepSTARR/
- [染色质谱分析 (DeepSea)](https://github.com/AIRI-Institute/GENA_LM/blob/main/./downstream_tasks/DeepSea/
- 多聚腺苷酸化位点预测 (APARENT)
预训练数据
下载和预处理数据
要下载人类基因组,请运行以下脚本:
./download_data.sh human
对于预处理,执行以下脚本:
python src/gena_lm/genome_tools/create_corpus.py --input_file data/ncbi_dataset/data/GCA_009914755.4/GCA_009914755.4_T2T-CHM13v2.0_genomic.fna --output_dir data/processed/human/
安装
对于具有稀疏注意力的模型(gena-lm-bigbird-base-sparse
、gena-lm-bigbird-base-sparse-t2t
),需要FP16支持和DeepSpeed。
用于FP16的APEX
安装APEX https://github.com/NVIDIA/apex#quick-start
git clone https://github.com/NVIDIA/apex
cd apex
# 最新的提交可能会构建失败
git checkout 2386a912164b0c5cfcd8be7a2b890fbac5607c82
# 如果pip版本 >= 23.1(参考:https://pip.pypa.io/en/stable/news/#v23-1),支持多个相同键的`--config-settings`...
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
# 否则
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --global-option="--cpp_ext" --global-option="--cuda_ext" ./
用于稀疏操作的DeepSpeed
需要安装DeepSpeed以使用语言模型的稀疏注意力版本。DeepSpeed稀疏注意力仅支持计算能力 >= 7的GPU(V100、T4、A100),CUDA 10.1、10.2、11.0或11.1,并且仅在FP16模式下运行(截至DeepSpeed 0.6.0)。
可以使用pytorch.org提供的CUDA 10.2/11.0/11.1版本的PyTorch>=1.7.1,<=1.10.1轮子。 但是,使用CUDA 11.1 PyTorch轮子的稀疏操作需要在系统上安装CUDA 11.3/11.4。 稀疏操作也可以与PyTorch==1.12.1 CUDA 11.3轮子一起使用,但运行DeepSpeed稀疏操作测试需要修改它们,因为它们检查Torch CUDA版本 <=11.1。 用于Triton 1.1.1的DeepSpeed分支已经更新了测试。
Triton 1.0.0和1.1.1需要python<=3.9。
pip install triton==1.0.0
DS_BUILD_SPARSE_ATTN=1 pip install deepspeed==0.6.0 --global-option="build_ext" --global-option="-j8" --no-cache
并使用以下命令检查安装:
ds_report
Triton 1.1.1
Triton 1.1.1为A100上的稀疏操作带来了2倍的速度提升,但DeepSpeed(0.6.5)目前仅支持triton 1.0.0。 在需要这种速度提升的情况下,可以使用支持triton 1.1.1的DeepSpeed分支:
pip install triton==1.1.1
git clone https://github.com/yurakuratov/DeepSpeed.git
cd DeepSpeed
DS_BUILD_SPARSE_ATTN=1 pip install -e . --global-option="build_ext" --global-option="-j8" --no-cache
并使用以下命令运行稀疏操作测试:
cd tests/unit
pytest -v test_sparse_attention.py
使用lm-experiments-tools进行微调
我们使用lm-experiments-tools仓库中的Trainer和多GPU训练作为我们微调脚本的基础。但是,您也可以使用HF Transformers Trainer、PyTorch Lightning或Accelerate和PyTorch与自定义训练循环来代替。
git clone https://github.com/yurakuratov/t5-experiments
cd t5-experiments
pip install -e .