Implementation of the Borzoi model from Calico in Pytorch! Original implementation and weights are here.
borzoi-pytorch is avaiable 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.
[1]
Predicting RNA-seq coverage from DNA sequence as a unifying model of gene regulation
Johannes Linder, Divyanshi Srivastava, Han Yuan, Vikram Agarwal, David R. Kelley
bioRxiv 2023.08.30.555582; doi: https://doi.org/10.1101/2023.08.30.555582
[2]
enformer-pytorch github,
Phil Wang lucidrains
Please cite this repository if you used borzoi-pytorch.