enformer-pytorch: Running DeepMind's Gene Expression Model in PyTorch

What enformer-pytorch replaces, and for whom

DeepMind's Enformer predicts gene expression from DNA sequence using an attention network, and the original code is TensorFlow with Sonnet. That is a real constraint if your training loop, data loaders and experiment tracking already live in PyTorch: you either maintain two stacks or rewrite the model. This repository is the port that removes that split. The README states it is an implementation of Enformer in PyTorch and that it "also contains the means to fine tune pretrained models for your downstream tasks." The intended reader is a computational biology or ML engineer who wants Enformer's architecture and released weights inside a normal PyTorch workflow, not a drop-in replacement for the DeepMind training pipeline. The pyproject.toml declares Python >= 3.6 and a dependency set that includes torch>=1.6, einops, torchmetrics, polars, pyfaidx, pyyaml and transformers[torch]==4.56.2, so the package assumes a modern scientific Python environment rather than a minimal one.

The architecture exposed through from_hparams

The model is constructed through a classmethod rather than a config file. The README's example passes dim = 1536, depth = 11, heads = 8, output_heads = dict(human = 5313, mouse = 1643) and target_length = 896. Input is a tensor of shape (1, 196_608) holding integer tokens for ACGTN in that order, with -1 for padding, and the output is a dictionary keyed by species: human gives (1, 896, 5313) and mouse gives (1, 896, 1643). Two details matter for anyone wiring this into a pipeline. First, you can pass one-hot encodings instead of indices, but the README says they "must be float values," and seq_indices_to_one_hot is provided for the conversion. Second, setting return_embeddings = True on forward returns a second tensor of shape (1, 896, 3072), which is the hook the fine-tuning wrappers use. The same forward call can compute a Poisson loss when you pass head and target, and return_corr_coef = True switches the output to a Pearson R, described in the README as the metric used in the paper.

Installing enformer-pytorch and running a first forward pass

Installation is a single pip command. The README gives no conda channel, no source build and no CUDA notes, so treat pip as the only documented path.

$ pip install enformer-pytorch

The simplest way to get a working model is from_pretrained, which pulls the ported TensorFlow weights from the Hugging Face repository EleutherAI/enformer-official-rough. The README notes the download is roughly 1GB.

from enformer_pytorch import from_pretrained

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

To check that the port behaves as expected on your machine, the repository ships test_pretrained.py. Running it prints a correlation coefficient on a single human validation sample; the README shows 0.5963 for that sample. If your number is far off, the problem is in your environment, not in your downstream code. If you would rather build the model yourself rather than load the pretrained one, the README warns that you must set use_tf_gamma = True in from_hparams to reproduce the pretrained behaviour, because of a difference between the TensorFlow and PyTorch implementations of xlogy.

Fine-tuning through adapter wrappers, not full retraining

The repository does not ask you to retrain Enformer. It provides three wrappers in enformer_pytorch.finetune that sit on top of the pretrained model. HeadAdapterWrapper takes num_tracks and an enformer instance, and a post_transformer_embed flag that selects whether embeddings come from after the final pointwise block (conv then gelu, the default) or right after the transformer block with a learned layernorm. ContextAdapterWrapper takes a context_dim and accepts a context tensor alongside the sequence, which is how you condition on cell type or transcription factor; the README's example uses context_dim = 1024 with 4 contexts and 4 target tracks. ContextAttentionAdapterWrapper aggregates attention over a set of context embeddings instead of a single one. For memory, from_pretrained accepts use_checkpointing = True, and target_length can be overridden, for example target_length = 128 with dropout_rate = 0.1 for shorter sequences. That last option is the one to think about: it changes the model's output resolution, so it is a fine-tuning convenience, not a way to run the pretrained model on truncated input and expect the published numbers.

The numerical accuracy caveat is the main limitation

The README is unusually direct about this, and it should shape how you use the package. The ported weights were released as enformer-official-rough because "there are still some rounding errors that seem to be accruing across the layers, resulting in an absolute error as high as 0.5." The author suspected the attention pooling module, noting that attention logits are high. A later update reports that John St. John found the rough model hits the paper's marks, with human Pearson R of 0.625 on validation and 0.65 on test. From version 0.8.0, from_pretrained uses precomputed gamma positions to address the TensorFlow versus PyTorch xlogy difference. That history is not a reason to avoid the package, but it is a reason to treat correlation as something you verify rather than assume. The bundled test reports 0.5963 on one validation sample, which is below the 0.625 figure quoted for validation overall, so a single-sample check is a sanity test, not a benchmark. If your work depends on absolute predicted values rather than rank correlation, the 0.5 absolute error note deserves attention before you build on top of it.

When the TensorFlow original is still the right choice

The honest alternative is DeepMind's own TensorFlow/Sonnet implementation, which the README links to under deepmind-research. The difference is not just language. The original is the reference: it is what the paper's numbers were produced with, and it is where any future official weights or corrections would land first. This port is a reimplementation whose accuracy depends on matching TensorFlow numerics, which is exactly the problem the gamma-position fix and the use_tf_gamma flag exist to solve. If your team already runs TensorFlow, or if you need to reproduce published results exactly rather than approximately, the port adds a translation layer you would then have to validate. The port's advantage is the surrounding ecosystem: PyTorch training loops, the adapter wrappers, and Hugging Face hosting. A second point of comparison is Basenji2, an earlier chromatin accessibility model that appears in related searches; it is a different architecture and not a substitute for Enformer's attention-based sequence-to-expression mapping.

Maintenance, licence and upgrade cost

The repository is not archived, and the last push was on 2026-06-26. The most recent release listed is 0.8.11 from 2025-07-07, while pyproject.toml in the repository declares version 0.8.12, so the working tree is ahead of the last tagged release. The licence is MIT, declared both in the LICENSE file and in the pyproject classifiers, which permits commercial and academic use with attribution; that is a statement about the repository's terms, not legal advice, and the ported model weights on Hugging Face are a separate artifact whose terms you should check yourself. Upgrade cost is dominated by one pinned dependency: transformers[torch]==4.56.2 is an exact pin, not a range. Loading the pretrained model depends on Hugging Face's custom model feature, so a transformers upgrade is not something you can do casually without checking that path still works. The rest of the dependency list (einops, torchmetrics, polars, pyfaidx, pyyaml) is unpinned and lower risk. The declared minimum is Python >= 3.6, which is old enough that most users will be well above it.

Editorial conclusion

Adopt enformer-pytorch if you already work in PyTorch and want Enformer's pretrained weights plus the HeadAdapterWrapper or ContextAdapterWrapper for fine-tuning on your own tracks. Do not adopt it if you need the reference TensorFlow/Sonnet implementation, or if you cannot afford to re-verify outputs against the paper's reported correlation. Before committing, run the bundled test_pretrained.py to confirm the port's correlation on your hardware, and check whether you need use_tf_gamma = True when instantiating from hparams rather than from_pretrained.

Frequently asked questions

Official sources

1. Issues
2. License: MIT
3. lucidrains/enformer-pytorch on GitHub
4. README
5. Releases