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

Implement object segmentation on images using HOG algorithm proposed in CVPR 2005

NeurIPS-2021: Neural Auto-Curricula in Two-Player Zero-Sum Games.

A static analysis library for computing graph representations of Python programs suitable for use with graph neural networks.

Multi-Anchor Active Domain Adaptation for Semantic Segmentation (ICCV 2021 Oral)

GNN-based Recommendation Benchma

Fast and Easy Infinite Neural Networks in Python

High accurate tool for automatic faces detection with landmarks

Make your own game in a font!

Code for the Paper "Diffusion Models for Handwriting Generation"

Official PyTorch implementation of the paper "TEMOS: Generating diverse human motions from textual descriptions"

Reinforcement learning library in JAX.

EqGAN - Improving GAN Equilibrium by Raising Spatial Awareness

Text-to-SQL in the Wild: A Naturally-Occurring Dataset Based on Stack Exchange Data

Código de um painel de auto atendimento feito em Python.

Code for "Multi-View Multi-Person 3D Pose Estimation with Plane Sweep Stereo"

Code release for "Transferable Semantic Augmentation for Domain Adaptation" (CVPR 2021)

Implementation of Neural Style Transfer in Pytorch

Algorithms for outlier, adversarial and drift detection

Centroid-UNet is deep neural network model to detect centroids from satellite images.

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