Pytorch implementation of the Borzoi model from Calico, and Flashzoi, an up to 3x faster Borzoi enhancement!
borzoi-pytorch is available on PyPI and can be installed with
pip install borzoi-pytorch
Ported weights (with permission) are uploaded to Huggingface, the model (human or mouse heads) can be loaded with:
from borzoi_pytorch import Borzoi
borzoi = Borzoi.from_pretrained('johahi/borzoi-replicate-0') # 'johahi/borzoi-replicate-[0-3][-mouse]'The Pytorch version produces the same predictions as the original implementation, see for instance in this notebook.
After installation of FlashAttention-2, Flashzoi offers 3x the speed of Borzoi at comparable or slightly better predictive performance and can be loaded with:
from borzoi_pytorch import Borzoi
borzoi = Borzoi.from_pretrained('johahi/flashzoi-replicate-0') # 'johahi/flashzoi-replicate-[0-3]'Note that this model should/must be run in autocast, and requires a modern Nvidia GPU.
The relative shift operation should be faster than in enformer_pytorch or other implementations.
Original Borzoi implementation and weights are here.
[1]
Linder, Johannes, et al. "Predicting RNA-seq coverage from DNA sequence as a unifying model of gene regulation." Nature Genetics (2025): 1-13; doi: https://doi.org/10.1101/2023.08.30.555582
[2]
Flashzoi: An enhanced Borzoi model for accelerated genomic analysis
Johannes C. Hingerl, Alexander Karollus, Julien Gagneur
bioRxiv 2024.12.18.629121; doi: https://doi.org/10.1101/2024.12.18.629121
[3]
enformer-pytorch github,
Phil Wang
Please cite the Borzoi paper [1], along with Flashzoi [2], if you used this repository or the models.
