The full training script for Enformer (Tensorflow Sonnet) on TPU clusters

Last update: Oct 19, 2022

Overview

Enformer TPU training script (wip)

The full training script for Enformer (Tensorflow Sonnet) on TPU clusters, in an effort to migrate the model to pytorch.

This was pieced together from the Deepmind Enformer repository, the colab training notebook, as well as Basenji sequence augmentation code

It accounts for:

distributed TPU training
distributed datasets
distributed validation
gradient clipping
cross replica batchnorms
dataset augmentation

Training takes about 3 days on v3-64

Todo

fix script for differences in sequence length in basenji training data, which is ~130k vs ~190k bp as in paper

Citations

@article {Avsec2021.04.07.438649,
    author  = {Avsec, {\v Z}iga and Agarwal, Vikram and Visentin, Daniel and Ledsam, Joseph R. and Grabska-Barwinska, Agnieszka and Taylor, Kyle R. and Assael, Yannis and Jumper, John and Kohli, Pushmeet and Kelley, David R.},
    title   = {Effective gene expression prediction from sequence by integrating long-range interactions},
    elocation-id = {2021.04.07.438649},
    year    = {2021},
    doi     = {10.1101/2021.04.07.438649},
    publisher = {Cold Spring Harbor Laboratory},
    URL     = {https://www.biorxiv.org/content/early/2021/04/08/2021.04.07.438649},
    eprint  = {https://www.biorxiv.org/content/early/2021/04/08/2021.04.07.438649.full.pdf},
    journal = {bioRxiv}
}

The full training script for Enformer (Tensorflow Sonnet) on TPU clusters

Related tags

Overview

Enformer TPU training script (wip)

Todo

Citations

Owner

Phil Wang

Aalto-cs-msc-theses - Listing of M.Sc. Theses of the Department of Computer Science at Aalto University

Probabilistic Programming and Statistical Inference in PyTorch

Code for Subgraph Federated Learning with Missing Neighbor Generation (NeurIPS 2021)

Gradient Step Denoiser for convergent Plug-and-Play

Source code for the GPT-2 story generation models in the EMNLP 2020 paper "STORIUM: A Dataset and Evaluation Platform for Human-in-the-Loop Story Generation"

AdaFocus V2: End-to-End Training of Spatial Dynamic Networks for Video Recognition

Reverse engineering Rosetta 2 in M1 Mac

BARF: Bundle-Adjusting Neural Radiance Fields 🤮 (ICCV 2021 oral)

Providing the solutions for high-frequency trading (HFT) strategies using data science approaches (Machine Learning) on Full Orderbook Tick Data.

MIMO-UNet - Official Pytorch Implementation

PhysCap: Physically Plausible Monocular 3D Motion Capture in Real Time

Official repository of "Investigating Tradeoffs in Real-World Video Super-Resolution"

A Python library for unevenly-spaced time series analysis

CasualHealthcare's Pneumonia detection with Artificial Intelligence (Convolutional Neural Network)

Semi-Supervised Learning for Fine-Grained Classification

[NeurIPS'21 Spotlight] PyTorch code for our paper "Aligned Structured Sparsity Learning for Efficient Image Super-Resolution"

MARS: Learning Modality-Agnostic Representation for Scalable Cross-media Retrieva

Retrieve and analysis data from SDSS (Sloan Digital Sky Survey)

[CVPR 2021] Semi-Supervised Semantic Segmentation with Cross Pseudo Supervision