Training

The Bolero model and its trainers. Borzoi is the reimplemented Flashzoi backbone; BorzoiLoRA ("Bolero") adds the conditional-LoRA adapters and output heads. The trainers drive the wandb setup → fit → test loop for single- and multi-dataset training.

Model

Borzoi

Bases: Module

create_from_config classmethod

create_from_config(config: dict)

Create the model from a configuration dictionary.

get_default_config classmethod

get_default_config()

Get default config.

from_checkpoint classmethod

from_checkpoint(checkpoint_path, weights_only=False)

Load model from checkpoint.

forward

forward(x, crop=True, *args, **kwargs)

Borzoi forward pass.

loss

loss(y_pred, y_true, reduce=True, position_weights=None, loss_type='poisson_multinomial')

Compute the loss for the Borzoi model.

Parameters:

Name	Type	Description	Default
y_pred	Tensor	Predicted values, shape (batch_size, out_channels, seq_len).	required
y_true	Tensor	True values, shape (batch_size, out_channels, seq_len).	required
reduce	bool	Whether to reduce the loss to a scalar.	True
position_weights	Tensor	Per-position weights applied to the loss.	None
loss_type	str	Which loss to compute (e.g. "poisson_multinomial").	'poisson_multinomial'

model_summary

model_summary(model, row_settings=('var_names',), input_size=None, input_data=None, depth=3, cache_forward_pass=True, col_names=('num_params',))

Print model summary.

BorzoiLoRA

Bases: Borzoi

setup_output_head

setup_output_head(out_channels, activation='softplus')

Setup a single output head

setup_dual_output_head

setup_dual_output_head(mc_channels=1, atac_channels=1)

Setup dual output for mc and atac dual modality training

setup_scooby_head

setup_scooby_head(embedding_dim, out_channels)

Setup a single Scooby output head.

setup_gene_count

setup_gene_count(out_channels)

Setup a single gene count output head.

setup_eqtl

setup_eqtl(out_channels)

Setup a single QTL output head.

setup_signal_model

setup_signal_model(out_channels, cell_emb_dim, cond_emb_dim, **cond_emb_kwargs)

Setup special modules for a signal model

freeze_batchnorms

freeze_batchnorms()

Freeze batchnorms in the base model.

https://github.com/lucidrains/tf-bind-transformer/blob/main/tf_bind_transformer/tf_bind_transformer.py#L468-L470

When finetune Enformer or Borzoi, we have to freeze batch norm due to small batch size used.

freeze_all_parameter_except_output_head

freeze_all_parameter_except_output_head()

Freeze all parameters except the final output head.

load_checkpoint_from_path

load_checkpoint_from_path(checkpoint_path, strict=True)

Load the pre-trained LoRA model.

make_lora_config

make_lora_config()

Make LoRA configuration for the Borzoi model.

collapse_lora

collapse_lora(embedding)

Collapse the LoRA model into base form given an embedding.

convert_to_lora

convert_to_lora()

Convert the model to LoRA.

forward

forward(x, embedding=None, crop=True, return_dna_embedding=False, **kwargs)

Borzoi forward pass to get final output.

forward_dna_only_model

forward_dna_only_model(x, embedding=None, crop=True, return_dna_embedding=False, gene_mask=None)

Forward pass with DNA input only

forward_signal_model

forward_signal_model(x, signal, embedding, crop=True, return_dna_embedding=False, gene_mask=None)

Forward pass with DNA and signal input.

forward_gene

forward_gene(x, gene_mask, embedding=None, crop=True, return_dna_embedding=False, **kwargs)

For gene count model (output_head_type="gene_count"), do a forward pass to get both final output and gene count output.

forward_qtl_with_dna_emb_and_mask

forward_qtl_with_dna_emb_and_mask(dna_embedding, gene_mask)

Forward pass with DNA embedding and gene mask input.

Parameters:

Name	Type	Description	Default
dna_embedding	Tensor	DNA embedding, shape (batch_size, 1920, seq_len).	required
gene_mask	Tensor	Gene mask, shape (batch_size, 5, seq_len). Gene mask has 5 channels: 1 for gene mask, 4 for variant alt allele one-hot encoding.	required

Returns:

Name	Type	Description
pval_and_pip_output	Tensor	Pval and pip output, shape (batch_size, 2).
slope_output	Tensor	Slope output, shape (batch_size, 1).

forward_qtl

forward_qtl(x, gene_mask, embedding=None, crop=True, return_dna_embedding=False, **kwargs)

Forward pass with DNA and gene mask input. For QTL model (output_head_type="qtl"), do a forward pass to get both final output and QTL output.

Parameters:

Name	Type	Description	Default
x	Tensor	DNA input, shape (batch_size, 4, seq_len).	required
gene_mask	Tensor	Gene mask, shape (batch_size, 5, seq_len). Gene mask has 5 channels: 1 for gene mask, 4 for variant alt allele one-hot encoding.	required
embedding	Tensor	Embedding, shape (batch_size, embedding_dim).	None
crop	bool	Whether to crop the track output.	True
return_dna_embedding	bool	Whether to return the DNA embedding.	False
**kwargs	dict	Additional keyword arguments.	{}

loss

loss(y_pred, y_true, reduce=True, position_weights=None)

Compute the loss for the Borzoi model.

Parameters:

Name	Type	Description	Default
y_pred	Tensor	Predicted values, shape (batch_size, out_channels, seq_len).	required
y_true	Tensor	True values, shape (batch_size, out_channels, seq_len).	required
reduce	bool	Whether to reduce the loss on seq_len dim	True
position_weights	Tensor	Whether to weight loss on each position along seq_len	None

delta_mse_loss

delta_mse_loss(y_pred, y_true, reduce=True)

Compute the delta MSE loss.

Parameters:

Name	Type	Description	Default
y_pred	Tensor	Predicted delta values, shape (batch_size, out_channels, seq_len).	required
y_true	Tensor	True delta values, shape (batch_size, out_channels, seq_len).	required

gene_count_loss

gene_count_loss(y_pred, y_true, reduce=True)

Compute the gene count loss.

qtl_loss

qtl_loss(y_pred_pval_and_pip, y_pred_slope, y_true_pval_and_pip, y_true_slope)

Compute the QTL loss.

Parameters:

Name	Type	Description	Default
y_pred_pval_and_pip	Tensor	Predicted pval and pip output, shape (batch_size, 2).	required
y_pred_slope	Tensor	Predicted slope output, shape (batch_size, 1).	required
y_true_pval_and_pip	Tensor	True pval and pip output, shape (batch_size, 2).	required
y_true_slope	Tensor	True slope output, shape (batch_size, 1).	required

Trainers

BorzoiLoRATrainer

Bases: BorzoiTrainerMixin

Train LoRA model on pseudobulk single-cell ATAC data.

train

train() -> None

Train the Borzoi LoRA model.

MultiBorzoiLoRATrainer

Bases: BorzoiLoRATrainer

Borzoi trainer for training on multiple datasets

Cross-validation splits

bolero.hg38_splits and bolero.mm10_splits (re-exported from bolero.tl.generic.train_helper) are 5-element lists giving the Borzoi fold-0…4 chromosome splits. Each element is a dict with "train", "valid", and "test" chromosome lists:

import bolero

fold0 = bolero.hg38_splits[0]
fold0["test"]   # test chromosomes for fold 0
fold0["valid"]  # validation chromosomes for fold 0