{"id":13477965,"url":"https://github.com/lucidrains/enformer-pytorch","last_synced_at":"2025-05-15T23:08:50.878Z","repository":{"id":38817735,"uuid":"420505474","full_name":"lucidrains/enformer-pytorch","owner":"lucidrains","description":"Implementation of Enformer, Deepmind's attention network for predicting gene expression, in Pytorch","archived":false,"fork":false,"pushed_at":"2024-10-09T16:53:36.000Z","size":9183,"stargazers_count":487,"open_issues_count":15,"forks_count":88,"subscribers_count":16,"default_branch":"main","last_synced_at":"2025-05-15T23:08:38.272Z","etag":null,"topics":["artificial-intelligence","attention-mechanism","deep-learning","dna-sequences","gene-expression","genomics","transformer"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/lucidrains.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2021-10-23T19:33:01.000Z","updated_at":"2025-05-12T20:42:05.000Z","dependencies_parsed_at":"2024-06-21T05:48:15.031Z","dependency_job_id":"96c1cfb1-6849-4642-a97d-e47d6f7d97bb","html_url":"https://github.com/lucidrains/enformer-pytorch","commit_stats":{"total_commits":169,"total_committers":8,"mean_commits":21.125,"dds":0.08875739644970415,"last_synced_commit":"5a5974d2821c728f93294731c50b55f1f55fd86d"},"previous_names":[],"tags_count":107,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fenformer-pytorch","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fenformer-pytorch/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fenformer-pytorch/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/lucidrains%2Fenformer-pytorch/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/lucidrains","download_url":"https://codeload.github.com/lucidrains/enformer-pytorch/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":254436950,"owners_count":22070947,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["artificial-intelligence","attention-mechanism","deep-learning","dna-sequences","gene-expression","genomics","transformer"],"created_at":"2024-07-31T16:01:50.590Z","updated_at":"2025-05-15T23:08:45.185Z","avatar_url":"https://github.com/lucidrains.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":"\u003cimg src=\"./enformer.png\" width=\"450px\"\u003e\u003c/img\u003e\n\n## Enformer - Pytorch\n\nImplementation of \u003ca href=\"https://deepmind.com/blog/article/enformer\"\u003eEnformer\u003c/a\u003e, Deepmind's attention network for predicting gene expression, in Pytorch. This repository also contains the means to fine tune pretrained models for your downstream tasks. The original tensorflow sonnet code can be found \u003ca href=\"https://github.com/deepmind/deepmind-research/tree/master/enformer\"\u003ehere\u003c/a\u003e.\n\nUpdate: finetuned for predicting pseudobulk chromatin accessibility \u003ca href=\"https://www.biorxiv.org/content/10.1101/2023.11.27.568764v1\"\u003ehere\u003c/a\u003e\n\n## Install\n\n```bash\n$ pip install enformer-pytorch\n```\n\n## Usage\n\n```python\nimport torch\nfrom enformer_pytorch import Enformer\n\nmodel = Enformer.from_hparams(\n    dim = 1536,\n    depth = 11,\n    heads = 8,\n    output_heads = dict(human = 5313, mouse = 1643),\n    target_length = 896,\n)\n    \nseq = torch.randint(0, 5, (1, 196_608)) # for ACGTN, in that order (-1 for padding)\noutput = model(seq)\n\noutput['human'] # (1, 896, 5313)\noutput['mouse'] # (1, 896, 1643)\n```\n\nYou can also directly pass in the sequence as one-hot encodings, which must be float values\n\n```python\nimport torch\nfrom enformer_pytorch import Enformer, seq_indices_to_one_hot\n\nmodel = Enformer.from_hparams(\n    dim = 1536,\n    depth = 11,\n    heads = 8,\n    output_heads = dict(human = 5313, mouse = 1643),\n    target_length = 896,\n)\n\nseq = torch.randint(0, 5, (1, 196_608))\none_hot = seq_indices_to_one_hot(seq)\n\noutput = model(one_hot)\n\noutput['human'] # (1, 896, 5313)\noutput['mouse'] # (1, 896, 1643)\n```\n\nFinally, one can fetch the embeddings, for fine-tuning and otherwise, by setting the `return_embeddings` flag to be `True` on forward\n\n```python\nimport torch\nfrom enformer_pytorch import Enformer, seq_indices_to_one_hot\n\nmodel = Enformer.from_hparams(\n    dim = 1536,\n    depth = 11,\n    heads = 8,\n    output_heads = dict(human = 5313, mouse = 1643),\n    target_length = 896,\n)\n\nseq = torch.randint(0, 5, (1, 196_608))\none_hot = seq_indices_to_one_hot(seq)\n\noutput, embeddings = model(one_hot, return_embeddings = True)\n\nembeddings # (1, 896, 3072)\n```\n\nFor training, you can directly pass the head and target in to get the poisson loss\n\n```python\nimport torch\nfrom enformer_pytorch import Enformer, seq_indices_to_one_hot\n\nmodel = Enformer.from_hparams(\n    dim = 1536,\n    depth = 11,\n    heads = 8,\n    output_heads = dict(human = 5313, mouse = 1643),\n    target_length = 200,\n).cuda()\n\nseq = torch.randint(0, 5, (196_608 // 2,)).cuda()\ntarget = torch.randn(200, 5313).cuda()\n\nloss = model(\n    seq,\n    head = 'human',\n    target = target\n)\n\nloss.backward()\n\n# after much training\n\ncorr_coef = model(\n    seq,\n    head = 'human',\n    target = target,\n    return_corr_coef = True\n)\n\ncorr_coef # pearson R, used as a metric in the paper\n```\n\n## Pretrained Model\n\nDeepmind has released the weights for their tensorflow sonnet Enformer model! I have ported it over to Pytorch and uploaded it to \u003ca href=\"https://huggingface.co/EleutherAI/enformer-official-rough\"\u003e🤗 Huggingface\u003c/a\u003e (~1GB). There are still some rounding errors that seem to be accruing across the layers, resulting in an absolute error as high as `0.5`. However, correlation coefficient look good so I am releasing the 'rough'ly working version. Will keep working on figuring out where the numerical errors are happening (it may be the attention pooling module, as I noticed the attention logits are pretty high).\n\nUpdate: \u003ca href=\"https://github.com/jstjohn\"\u003eJohn St. John\u003c/a\u003e did some work and found that the `enformer-official-rough` model hits the reported marks in the paper - human pearson R of `0.625` for validation, and `0.65` for test.\n\nUpdate: As of version 0.8.0, if one were to use the `from_pretrained` function to load the pretrained model, it should automatically use precomputed gamma positions to address a difference between tensorflow and pytorch `xlogy`. This should resolve the numerical discrepancy above. If you were to further finetune and not be using the `from_pretrained` function, please make sure to set `use_tf_gamma = True` when using `.from_hparams` to instantiate the `Enformer`\n\n```bash\n$ pip install enformer-pytorch\u003e=0.5\n````\n\nLoading the model\n\n```python\nfrom enformer_pytorch import from_pretrained\n\nenformer = from_pretrained('EleutherAI/enformer-official-rough')\n```\n\nQuick sanity check on a single human validation point\n\n```python\n$ python test_pretrained.py\n# 0.5963 correlation coefficient on a validation sample\n```\n\nThis is all made possible thanks to HuggingFace's [custom model](https://huggingface.co/docs/transformers/master/en/custom_models) feature.\n\nYou can also load, with overriding of the `target_length` parameter, if you are working with shorter sequence lengths\n\n```python\nfrom enformer_pytorch import from_pretrained\n\nmodel = from_pretrained('EleutherAI/enformer-official-rough', target_length = 128, dropout_rate = 0.1)\n\n# do your fine-tuning\n```\n\nTo save on memory during fine-tuning a large Enformer model\n\n```python\nfrom enformer_pytorch import from_pretrained\n\nenformer = from_pretrained('EleutherAI/enformer-official-rough', use_checkpointing = True)\n\n# finetune enformer on a limited budget\n```\n\n## Fine-tuning\n\nThis repository will also allow for easy fine-tuning of Enformer.\n\nFine-tuning on new tracks\n\n```python\nimport torch\nfrom enformer_pytorch import from_pretrained\nfrom enformer_pytorch.finetune import HeadAdapterWrapper\n\nenformer = from_pretrained('EleutherAI/enformer-official-rough')\n\nmodel = HeadAdapterWrapper(\n    enformer = enformer,\n    num_tracks = 128,\n    post_transformer_embed = False   # by default, embeddings are taken from after the final pointwise block w/ conv -\u003e gelu - but if you'd like the embeddings right after the transformer block with a learned layernorm, set this to True\n).cuda()\n\nseq = torch.randint(0, 5, (1, 196_608 // 2,)).cuda()\ntarget = torch.randn(1, 200, 128).cuda()  # 128 tracks\n\nloss = model(seq, target = target)\nloss.backward()\n```\n\nFinetuning on contextual data (cell type, transcription factor, etc)\n\n```python\nimport torch\nfrom enformer_pytorch import from_pretrained\nfrom enformer_pytorch.finetune import ContextAdapterWrapper\n\nenformer = from_pretrained('EleutherAI/enformer-official-rough')\n    \nmodel = ContextAdapterWrapper(\n    enformer = enformer,\n    context_dim = 1024\n).cuda()\n\nseq = torch.randint(0, 5, (1, 196_608 // 2,)).cuda()\n\ntarget = torch.randn(1, 200, 4).cuda()  # 4 tracks\ncontext = torch.randn(4, 1024).cuda()   # 4 contexts for the different 'tracks'\n\nloss = model(\n    seq,\n    context = context,\n    target = target\n)\n\nloss.backward()\n```\n\nFinally, there is also a way to use attention aggregation from a set of context embeddings (or a single context embedding). Simply use the `ContextAttentionAdapterWrapper`\n\n```python\nimport torch\nfrom enformer_pytorch import from_pretrained\nfrom enformer_pytorch.finetune import ContextAttentionAdapterWrapper\n\nenformer = from_pretrained('EleutherAI/enformer-official-rough')\n    \nmodel = ContextAttentionAdapterWrapper(\n    enformer = enformer,\n    context_dim = 1024,\n    heads = 8,              # number of heads in the cross attention\n    dim_head = 64           # dimension per head\n).cuda()\n\nseq = torch.randint(0, 5, (1, 196_608 // 2,)).cuda()\n\ntarget = torch.randn(1, 200, 4).cuda()      # 4 tracks\ncontext = torch.randn(4, 16, 1024).cuda()   # 4 contexts for the different 'tracks', each with 16 tokens\n\ncontext_mask = torch.ones(4, 16).bool().cuda() # optional context mask, in example, include all context tokens\n\nloss = model(\n    seq,\n    context = context,\n    context_mask = context_mask,\n    target = target\n)\n\nloss.backward()\n```\n\n## Data\n\nYou can use the `GenomicIntervalDataset` to easily fetch sequences of any length from a \u003ca href=\"https://genome.ucsc.edu/FAQ/FAQformat.html#format1\"\u003e`.bed`\u003c/a\u003e file, with greater context length dynamically computed if specified\n\n```python\nimport torch\nimport polars as pl\nfrom enformer_pytorch import Enformer, GenomeIntervalDataset\n\nfilter_train = lambda df: df.filter(pl.col('column_4') == 'train')\n\nds = GenomeIntervalDataset(\n    bed_file = './sequences.bed',                       # bed file - columns 0, 1, 2 must be \u003cchromosome\u003e, \u003cstart position\u003e, \u003cend position\u003e\n    fasta_file = './hg38.ml.fa',                        # path to fasta file\n    filter_df_fn = filter_train,                        # filter dataframe function\n    return_seq_indices = True,                          # return nucleotide indices (ACGTN) or one hot encodings\n    shift_augs = (-2, 2),                               # random shift augmentations from -2 to +2 basepairs\n    context_length = 196_608,\n    # this can be longer than the interval designated in the .bed file,\n    # in which case it will take care of lengthening the interval on either sides\n    # as well as proper padding if at the end of the chromosomes\n    chr_bed_to_fasta_map = {\n        'chr1': 'chromosome1',  # if the chromosome name in the .bed file is different than the key name in the fasta file, you can rename them on the fly\n        'chr2': 'chromosome2',\n        'chr3': 'chromosome3',\n        # etc etc\n    }\n)\n\nmodel = Enformer.from_hparams(\n    dim = 1536,\n    depth = 11,\n    heads = 8,\n    output_heads = dict(human = 5313, mouse = 1643),\n    target_length = 896,\n)\n\nseq = ds[0] # (196608,)\npred = model(seq, head = 'human') # (896, 5313)\n```\n\nTo return the random shift value, as well as whether reverse complement was activated (in the case you need to reverse the corresponding chip-seq target data), just set `return_augs = True` when initializing the `GenomicIntervalDataset`\n\n```python\nimport torch\nimport polars as pl\nfrom enformer_pytorch import Enformer, GenomeIntervalDataset\n\nfilter_train = lambda df: df.filter(pl.col('column_4') == 'train')\n\nds = GenomeIntervalDataset(\n    bed_file = './sequences.bed',                       # bed file - columns 0, 1, 2 must be \u003cchromosome\u003e, \u003cstart position\u003e, \u003cend position\u003e\n    fasta_file = './hg38.ml.fa',                        # path to fasta file\n    filter_df_fn = filter_train,                        # filter dataframe function\n    return_seq_indices = True,                          # return nucleotide indices (ACGTN) or one hot encodings\n    shift_augs = (-2, 2),                               # random shift augmentations from -2 to +2 basepairs\n    rc_aug = True,                                      # use reverse complement augmentation with 50% probability\n    context_length = 196_608,\n    return_augs = True                                  # return the augmentation meta data\n)\n\nseq, rand_shift_val, rc_bool = ds[0] # (196608,), (1,), (1,)\n```\n\n## Appreciation\n\nSpecial thanks goes out to \u003ca href=\"https://www.eleuther.ai/\"\u003eEleutherAI\u003c/a\u003e for providing the resources to retrain the model, during a time when the official model from Deepmind had not been released yet.\n\nThanks also goes out to \u003ca href=\"johahi\"\u003e@johahi\u003c/a\u003e for finding out that there are numerical differences between the torch and tensorflow implementations of `xlogy`. He provided a fix for this difference, which is adopted in this repository in `v0.8.0`\n\n## Todo\n\n- [x] script to load weights from trained tensorflow enformer model to pytorch model\n- [x] add loss wrapper with poisson loss\n- [x] move the metrics code over to pytorch as well\n- [x] train enformer model\n- [x] build context manager for fine-tuning with unfrozen enformer but with frozen batchnorm\n- [x] allow for plain fine-tune with fixed static context\n- [x] allow for fine tuning with only unfrozen layernorms (technique from fine tuning transformers)\n- [x] fix handling of 'N' in sequence, figure out representation of N in basenji barnyard\n- [x] take care of shift augmentation in `GenomicIntervalDataset`\n- [x] speed up `str_to_seq_indices`\n- [x] add to EleutherAI huggingface (done thanks to \u003ca href=\"https://github.com/NielsRogge\"\u003eNiels\u003c/a\u003e)\n- [ ] offer some basic training utils, as gradient accumulation will be needed for fine tuning\n\n## Citations\n\n```bibtex\n@article {Avsec2021.04.07.438649,\n    author  = {Avsec, {\\v Z}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.},\n    title   = {Effective gene expression prediction from sequence by integrating long-range interactions},\n    elocation-id = {2021.04.07.438649},\n    year    = {2021},\n    doi     = {10.1101/2021.04.07.438649},\n    publisher = {Cold Spring Harbor Laboratory},\n    URL     = {https://www.biorxiv.org/content/early/2021/04/08/2021.04.07.438649},\n    eprint  = {https://www.biorxiv.org/content/early/2021/04/08/2021.04.07.438649.full.pdf},\n    journal = {bioRxiv}\n}\n```\n\n```bibtex\n@misc{liu2022convnet,\n    title   = {A ConvNet for the 2020s},\n    author  = {Zhuang Liu and Hanzi Mao and Chao-Yuan Wu and Christoph Feichtenhofer and Trevor Darrell and Saining Xie},\n    year    = {2022},\n    eprint  = {2201.03545},\n    archivePrefix = {arXiv},\n    primaryClass = {cs.CV}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fenformer-pytorch","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Flucidrains%2Fenformer-pytorch","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Flucidrains%2Fenformer-pytorch/lists"}