enformer-pytorch:基于 Pytorch 的基因表达预测网络项目

Enformer - PyTorch

在PyTorch中实现Enformer——DeepMind用于基因表达预测的注意力网络。此仓库还包含了对预训练模型进行微调以适应下游任务的方法。原始的TensorFlow Sonnet代码可以在这里找到。

更新：已针对伪批量染色质可及性预测进行微调：此处

安装

$ pip install enformer-pytorch

使用方法

import torch
from enformer_pytorch import Enformer

model = Enformer.from_hparams(
    dim = 1536,
    depth = 11,
    heads = 8,
    output_heads = {'human': 5313, 'mouse': 1643},
    target_length = 896,
)
    
seq = torch.randint(0, 5, (1, 196_608)) # ACGTN序列为范围，-1表示填充
output = model(seq)

output['human'] # (1, 896, 5313)
output['mouse'] # (1, 896, 1643)

您也可以直接传入一个热编码序列，该序列必须是浮点值

import torch
from enformer_pytorch import Enformer, seq_indices_to_one_hot

model = Enformer.from_hparams(
    dim = 1536,
    depth = 11,
    heads = 8,
    output_heads = {'human': 5313, 'mouse': 1643},
    target_length = 896,
)

seq = torch.randint(0, 5, (1, 196_608))
one_hot = seq_indices_to_one_hot(seq)

output = model(one_hot)

output['human'] # (1, 896, 5313)
output['mouse'] # (1, 896, 1643)

最后，可以通过设置前向传播中的return_embeddings标志为True来获取嵌入向量

import torch
from enformer_pytorch import Enformer, seq_indices_to_one_hot

model = Enformer.from_hparams(
    dim = 1536,
    depth = 11,
    heads = 8,
    output_heads = {'human': 5313, 'mouse': 1643},
    target_length = 896,
)

seq = torch.randint(0, 5, (1, 196_608))
one_hot = seq_indices_to_one_hot(seq)

output, embeddings = model(one_hot, return_embeddings=True)

embeddings # (1, 896, 3072)

对于训练，您可以直接传递头和目标来计算泊松损失

import torch
from enformer_pytorch import Enformer, seq_indices_to_one_hot

model = Enformer.from_hparams(
    dim = 1536,
    depth = 11,
    heads = 8,
    output_heads = {'human': 5313, 'mouse': 1643},
    target_length = 200,
).cuda()

seq = torch.randint(0, 5, (196_608 // 2,)).cuda()
target = torch.randn(200, 5313).cuda()

loss = model(
    seq,
    head='human',
    target=target
)

loss.backward()

# 训练很多次后

corr_coef = model(
    seq,
    head='human',
    target=target,
    return_corr_coef=True
)

corr_coef # 相关系数，论文中用作指标

预训练模型

DeepMind发布了他们的TensorFlow Sonnet Enformer模型的权重！我已经将其移植到PyTorch并上传到了🤗 Huggingface（约1GB）。不同层之间似乎存在一些舍入误差，导致绝对误差高达0.5。然而，相关系数看起来不错，所以我发布了“大致”正常运行的版本。我将继续研究数值错误发生在哪里（可能是注意力池化模块，因为我注意到注意力logits相当高）。

更新：John St. John做了一些工作，发现enformer-official-rough模型达到了论文中报告的指标——人类验证样本的相关系数为0.625，测试样本为0.65。

更新：自版本0.8.0起，如果使用from_pretrained函数加载预训练模型，它会自动使用预先计算的gamma位置来解决TensorFlow和PyTorch的xlogy之间的差异。这应解决上面提到的数值不匹配问题。如果您进一步微调且不使用from_pretrained函数，请确保在使用.from_hparams实例化Enformer时设置use_tf_gamma = True。

$ pip install enformer-pytorch>=0.5

加载模型

from enformer_pytorch import from_pretrained

enformer = from_pretrained('EleutherAI/enformer-official-rough')

快速检查单个人类验证点的合理性

$ python test_pretrained.py
# 0.5963 的相关系数在一个验证样本上

这一切都得益于HuggingFace的自定义模型功能。

您还可以根据需要覆盖target_length参数，如果您正在处理较短的序列长度

from enformer_pytorch import from_pretrained

model = from_pretrained('EleutherAI/enformer-official-rough', target_length=128, dropout_rate=0.1)

# 进行微调

为了在微调大型Enformer模型时节省内存

from enformer_pytorch import from_pretrained

enformer = from_pretrained('EleutherAI/enformer-official-rough', use_checkpointing=True)

# 在有限预算下微调enformer

微调

此仓库也允许对Enformer进行轻松微调。

对新轨道进行微调

import torch
from enformer_pytorch import from_pretrained
from enformer_pytorch.finetune

```markdown
## 数据处理
你可以使用`GenomicIntervalDataset`从<a href="https://genome.ucsc.edu/FAQ/FAQformat.html#format1">`.bed`</a>文件中轻松获取任意长度的序列，并根据需要动态计算更大的上下文长度。

```python
import torch
import polars as pl
from enformer_pytorch import Enformer, GenomeIntervalDataset

def filter_train(df):
    return df.filter(pl.col('column_4') == 'train')

dataset = GenomeIntervalDataset(
    bed_file = './sequences.bed',                     # bed 文件，第 0、1、2 列分别是染色体、起始位置和结束位置
    fasta_file = './hg38.ml.fa',                      # 快速傅里叶变换（fasta）文件路径
    filter_df_fn = filter_train,                      # 过滤数据框的函数
    return_seq_indices = True,                         # 返回核苷酸索引（ACGTN）或独热编码
    shift_augs = (-2, 2),                             # 随机位移增强，范围从 -2 到 +2 碱基对
    context_length = 196_608,
    # 可以比 .bed 文件指定的区间更长，此时它会在两侧扩展区间
    # 并在染色体末端适当地填充
    chr_bed_to_fasta_map = {
        'chr1': 'chromosome1',                    # 如果 .bed 文件中的染色体名与 fasta 文件中的键名不同，可以实时重命名
        'chr2': 'chromosome2',
        'chr3': 'chromosome3',
        # ...
    }
)

model = Enformer.from_hparams(
    dim = 1536,
    depth = 11,
    heads = 8,
    output_heads = dict(human = 5313, mouse = 1643),
    target_length = 896,
)

sequence = dataset[0]  # (196608,)
prediction = model(sequence, head = 'human')  # (896, 5313)

为了返回随机位移值以及是否启用了反向互补（如果你需要反转对应的芯片测序目标数据），只需在初始化GenomicIntervalDataset时设置return_augs = True。

import torch
import polars as pl
from enformer_pytorch import Enformer, GenomeIntervalDataset

def filter_train(df):
    return df.filter(pl.col('column_4') == 'train')

dataset = GenomeIntervalDataset(
    bed_file = './sequences.bed',
    fasta_file = './hg38.ml.fa',
    filter_df_fn = filter_train,
    return_seq_indices = True,
    shift_augs = (-2, 2),
    rc_aug = True,                                   # 50% 的概率使用反向互补增强
    context_length = 196_608,
    return_augs = True                               # 返回增强元数据
)

seq, rand_shift_val, rc_bool = dataset[0]  # (196608,), (1,), (1,)

致谢

特别感谢EleutherAI提供了资源来重新训练模型，当时Deepmind的官方模型尚未发布。

同样要感谢@johahi发现了PyTorch和TensorFlow实现xlogy之间的数值差异并提供了修复方法，该修复已在本仓库的v0.8.0版本中采用。

待办事项

• 脚本：将训练好的Tensorflow Enformer模型权重加载到PyTorch模型
• 添加带泊松损失的损失包装器
• 将度量代码迁移到PyTorch
• 训练Enformer模型
• 构建用于微调的上下文管理器（解冻Enformer但冻结批量归一化）
• 允许固定静态上下文的简单微调
• 允许仅解冻层标准化（来自微调变压器的技术）
• 解决序列中'N'的处理，确定Barnyard中'N'的表示方式
• 处理GenomicIntervalDataset中的位移增强
• 加快str_to_seq_indices的速度
• 添加到EleutherAI的HuggingFace（感谢Niels的贡献）
• 提供一些基础训练工具，因为微调可能需要梯度累积

引用

@article {Avsec2021.04.07.438649,
    author  = {Avsec, Žiga and Agarwal, Vikram and Visentin, Daniel and Ledsam, Joseph R. and Grabska-Barwinska, Agnieszka and Taylor, Kyle R. and Assael, Yannis and Jumper, John and Kohli, Pushmeet and Kelley, David R.},
    title   = {有效从序列预测基因表达，通过整合长程相互作用},
    elocation-id = {2021.04.07.438649},
    year    = {2021},
    doi     = {10.1101/2021.04.07.438649},
    publisher = {Cold Spring Harbor Laboratory},
    URL     = {https://www.biorxiv.org/content/early/2021/04/08/2021.04.07.438649},
    eprint  = {https://www.biorxiv.org/content/early/2021/04/08/2021.04.07.438649.full.pdf},
    journal = {bioRxiv}
}

@misc{liu2022convnet,
    title   = {2020年代的卷积网络},
    author  = {Liu, Zhuang and Mao, Hanzi and Wu, Chao-Yuan and Feichtenhofer, Christoph and Darrell, Trevor and Xie, Saining},
    year    = {2022},
    eprint  = {2201.03545},
    archivePrefix = {arXiv},
    primaryClass = {cs.CV}
}