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

Code for "Sparse Steerable Convolutions: An Efficient Learning of SE(3)-Equivariant Features for Estimation and Tracking of Object Poses in 3D Space"

x-transformers-paddle 2.x version

[内测中]前向式Python环境快捷封装工具，快速将Python打包为EXE并添加CUDA、NoAVX等支持。

Laplace Redux -- Effortless Bayesian Deep Learning

Pytorch implementation of "Forward Thinking: Building and Training Neural Networks One Layer at a Time"

Reinforcement learning for self-driving in a 3D simulation

DEEPAGÉ: Answering Questions in Portuguese about the Brazilian Environment

AdamW optimizer and cosine learning rate annealing with restarts

Object-aware Contrastive Learning for Debiased Scene Representation

PURE: End-to-End Relation Extraction

LSTMs (Long Short Term Memory) RNN for prediction of price trends

TinyML Cookbook, published by Packt

A pytorch implementation of Pytorch-Sketch-RNN

Differentiable architecture search for convolutional and recurrent networks

Learning Multiresolution Matrix Factorization and its Wavelet Networks on Graphs

MIMIC Code Repository: Code shared by the research community for the MIMIC-III database

Predictive Maintenance LSTM

Code and data for "Broaden the Vision: Geo-Diverse Visual Commonsense Reasoning" (EMNLP 2021).

Multi-Person Extreme Motion Prediction

Code for the Paper: Alexandra Lindt and Emiel Hoogeboom.