DOCSEARCHbyAlgolia
Submit Docs

enformer-pytorch

repository·main·Indexed 20 days ago

https://github.com/lucidrains/enformer-pytorch

A PyTorch implementation of DeepMind's Enformer model, an attention-based network for predicting gene expression from DNA sequences. The library supports loading pretrained weights from HuggingFace, fine-tuning via HeadAdapterWrapper, ContextAdapterWrapper, and ContextAttentionAdapterWrapper, and provides genomic data handling tools including GenomeIntervalDataset and FastaInterval for processing .bed and .fasta files.

Tokens
9.4K
Snippets
34
Records
34
Agent score
68%

What's inside enformer-pytorch

  1. Initialize and use the Enformer model

    main

    You can instantiate an Enformer model using Enformer.from_hparams. The model accepts DNA sequences as integer indices (representing ACGTN, where -1 is padding) or as one-hot encoded float tensors.

    To convert integer indices to one-hot encodings, use the seq_indices_to_one_hot utility function.

    import torch
from enformer_pytorch import Enformer, seq_indices_to_one_hot

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

# Using integer indices (ACGTN order, -1 for padding)
seq = torch.randint(0, 5, (1, 196_608))
output = model(seq)

# Using one-hot encodings
one_hot = seq_indices_to_one_hot(seq)
output = model(one_hot)

  2. Load pretrained Enformer models

    main

    Load the official ported weights from HuggingFace using from_pretrained.

    Important Note on Numerical Stability: As of version 0.8.0, from_pretrained automatically handles xlogy differences between TensorFlow and PyTorch. If you are manually instantiating a model via .from_hparams for fine-tuning, you must set use_tf_gamma = True to ensure numerical consistency.

    You can override target_length or enable use_checkpointing to save memory during fine-tuning.

    from enformer_pytorch import from_pretrained

# Basic loading
enformer = from_pretrained('EleutherAI/enformer-official-rough')

# Loading with custom parameters and memory optimization
model = from_pretrained(
    'EleutherAI/enformer-official-rough', 
    target_length = 128, 
    dropout_rate = 0.1,
    use_checkpointing = True
)

  4. How to evaluate Enformer correlation on validation/test sets

    main

    The following pattern demonstrates how to wrap the evaluation logic into a function to test correlation on different subsets (train, valid, or test) for a specific organism.

    def compute_correlation(model, organism:str="human", subset:str="valid", max_steps=-1):
  # ... setup dataset and loader ...
  corr_coef = MeanPearsonCorrCoefPerChannel(n_channels=ds.num_channels)
  # ... loop through loader ...
  # pred = model(sequence)[organism]
  # ...
  return corr_coef.compute().mean()

# Usage examples:
compute_correlation(model, organism="human", subset="valid", max_steps=100)
compute_correlation(model, organism="human", subset="test", max_steps=-1)
    def compute_correlation(model, organism:str="human", subset:str="valid", max_steps=-1):
  fasta_path = human_fasta_path if organism == "human" else mouse_fasta_path
  ds = BasenjiDataSet(organism, subset, SEQUENCE_LENGTH, fasta_path)
  total = len(ds.region_df)
  dl = torch.utils.data.DataLoader(ds, num_workers=0, batch_size=1)
  corr_coef = MeanPearsonCorrCoefPerChannel(n_channels=ds.num_channels)
  n_steps = total if max_steps <= 0 else max_steps
  for i,batch in enumerate(tqdm(dl, total=n_steps)):
    if max_steps > 0 and i >= max_steps:
      break
    batch_gpu = {k:v.to(model.device) for k,v in batch.items()}
    sequence = batch_gpu['sequence']
    target = batch_gpu['target']
    with torch.no_grad():
      pred = model(sequence)[organism]
      corr_coef(preds=pred.cpu(), target=target.cpu())
  return corr_coef.compute().mean()

compute_correlation(model, organism="human", subset="valid", max_steps=100)

  5. Install enformer-pytorch and dependencies

    main

    To use the Enformer model, you need to clone the repository, install the package, and install additional dependencies for genomic data processing and evaluation.

    # Clone the repository
!git clone https://github.com/lucidrains/enformer-pytorch.git

# Install the package
!cd enformer-pytorch && pip install .

# Install evaluation and genomic dependencies
!pip install torchmetrics kipoiseq==0.5.2 BioPython --quiet
    !git clone https://github.com/lucidrains/enformer-pytorch.git
!cd enformer-pytorch && pip install .
!pip install torchmetrics kipoiseq==0.5.2 BioPython --quiet > /dev/null

  6. Train Enformer with Poisson loss

    main

    For training, you can pass a specific head and target directly to the model to compute the Poisson loss. You can also compute the Pearson correlation coefficient (as used in the original paper) by setting return_corr_coef = True.

    # Standard training step with Poisson loss
loss = model(
    seq,
    head = 'human',
    target = target
)
loss.backward()

# Evaluation using Pearson R
corr_coef = model(
    seq,
    head = 'human',
    target = target,
    return_corr_coef = True
)

  7. Prepare genomic data with GenomeIntervalDataset

    main

    The GenomeIntervalDataset fetches sequences from .bed and .fasta files. It handles dynamic context lengthening and padding automatically.

    Key Parameters:

    • bed_file: Path to the .bed file (columns 0, 1, 2 must be <chromosome>, <start>, <end>).
    • fasta_file: Path to the reference genome FASTA file.
    • filter_df_fn: A function to filter the BED dataframe (e.g., for train/val splits).
    • return_seq_indices: If True, returns nucleotide indices (ACGTN); if False, returns one-hot encodings.
    • shift_augs: Tuple defining random shift augmentations (e.g., (-2, 2)).
    • rc_aug: If True, applies reverse complement augmentation with 50% probability.
    • context_length: The desired sequence length.
    • return_augs: If True, returns augmentation metadata (shift value and RC boolean).
    • chr_bed_to_fasta_map: Dictionary to map BED chromosome names to FASTA keys.
    import polars as pl
from enformer_pytorch import GenomeIntervalDataset

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

ds = GenomeIntervalDataset(
    bed_file = './sequences.bed',
    fasta_file = './hg38.ml.fa',
    filter_df_fn = filter_train,
    return_seq_indices = True,
    shift_augs = (-2, 2),
    context_length = 196_608,
    rc_aug = True,
    return_augs = True
)

# Returns: (sequence, random_shift_value, reverse_complement_bool)
seq, rand_shift_val, rc_bool = ds[0]

  8. Fine-tune Enformer with contextual data using ContextAdapterWrapper

    main

    Use ContextAdapterWrapper to incorporate contextual information (like cell type or transcription factors) into the model.

    • context_dim: The dimensionality of the context embeddings.
    from enformer_pytorch.finetune import ContextAdapterWrapper

model = ContextAdapterWrapper(
    enformer = enformer,
    context_dim = 1024
).cuda()

# context shape: (num_contexts, context_dim)
loss = model(seq, context = context, target = target)
loss.backward()

  9. Fetch model embeddings

    main

    To retrieve the model's embeddings (useful for fine-tuning or downstream tasks), set the return_embeddings flag to True during the forward pass.

    output, embeddings = model(one_hot, return_embeddings = True)
# embeddings shape: (1, 896, 3072)

  10. Fine-tune Enformer on new tracks with HeadAdapterWrapper

    main

    Use HeadAdapterWrapper to adapt the Enformer model to new tracks.

    • num_tracks: The number of tracks in your new target.
    • post_transformer_embed: If True, embeddings are taken from after the final pointwise block (conv -> gelu). If False (default), they are taken from after the transformer block with a learned layernorm.
    from enformer_pytorch.finetune import HeadAdapterWrapper

model = HeadAdapterWrapper(
    enformer = enformer,
    num_tracks = 128,
    post_transformer_embed = False
).cuda()

loss = model(seq, target = target)
loss.backward()

  11. Fine-tune Enformer with attention aggregation using ContextAttentionAdapterWrapper

    main

    Use ContextAttentionAdapterWrapper to perform attention aggregation from a set of context embeddings. This allows for multiple context tokens per track.

    • context_dim: Dimension of the context embeddings.
    • heads: Number of heads in the cross attention.
    • dim_head: Dimension per head.
    • context_mask: (Optional) A boolean mask for the context tokens.
    from enformer_pytorch.finetune import ContextAttentionAdapterWrapper

model = ContextAttentionAdapterWrapper(
    enformer = enformer,
    context_dim = 1024,
    heads = 8,
    dim_head = 64
).cuda()

# context shape: (num_tracks, num_context_tokens, context_dim)
loss = model(
    seq, 
    context = context, 
    context_mask = context_mask, 
    target = target
)
loss.backward()

  12. Use checkpointing to save memory

    main

    The Enformer model supports gradient checkpointing on the transformer trunk to reduce memory usage during training. This is controlled by the use_checkpointing flag in the EnformerConfig.

    When enabled, the model uses trunk_checkpointed during the forward pass, which applies torch.utils.checkpoint.checkpoint_sequential to the transformer layers.

    config = EnformerConfig(
    ...,
    use_checkpointing=True
)
model = Enformer(config)