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

Automated Hyperparameter Optimization Competition

Plenoxels: Radiance Fields without Neural Networks, Code release WIP

Global-Local Attention for Emotion Recognition

HDR Video Reconstruction: A Coarse-to-fine Network and A Real-world Benchmark Dataset (ICCV 2021)

Offical implementation of Shunted Self-Attention via Multi-Scale Token Aggregation

code for EMNLP 2019 paper Text Summarization with Pretrained Encoders

A FAIR dataset of TCV experimental results for validating edge/divertor turbulence models.

Discriminative Condition-Aware PLDA

AI-generated-characters for Learning and Wellbeing

Dynamic Head: Unifying Object Detection Heads with Attentions

(ICCV 2021) ProHMR - Probabilistic Modeling for Human Mesh Recovery

Wind Speed Prediction using LSTMs in PyTorch

Style-based Neural Drum Synthesis with GAN inversion

Denoising Normalizing Flow

Dynamic hair modeling from monocular videos using deep neural networks

Code for our NeurIPS 2021 paper Mining the Benefits of Two-stage and One-stage HOI Detection

Using LSTM to detect spoofing attacks in an Air-Ground network

Implementation of algorithms for continuous control (DDPG and NAF).

Official implementation of "Robust channel-wise illumination estimation"

Code for CVPR 2021 paper TransNAS-Bench-101: Improving Transferrability and Generalizability of Cross-Task Neural Architecture Search.