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

Efficient Deep Learning Systems course

This is the code of NeurIPS'21 paper "Towards Enabling Meta-Learning from Target Models".

MultiLexNorm 2021 competition system from ÚFAL

Vector Neurons: A General Framework for SO(3)-Equivariant Networks

Pcos-prediction - Predicts the likelihood of Polycystic Ovary Syndrome based on patient attributes and symptoms

Video Background Music Generation with Controllable Music Transformer (ACM MM 2021 Oral)

Continual World is a benchmark for continual reinforcement learning

PyTorch implementation of "Learn to Dance with AIST++: Music Conditioned 3D Dance Generation."

A denoising autoencoder + adversarial losses and attention mechanisms for face swapping.

Official Repo for Ground-aware Monocular 3D Object Detection for Autonomous Driving

Open Source Differentiable Computer Vision Library for PyTorch

Syllabus del curso IIC2115 - Programación como Herramienta para la Ingeniería 2022/I

BEAS: Blockchain Enabled Asynchronous & Secure Federated Machine Learning

ResNEsts and DenseNEsts: Block-based DNN Models with Improved Representation Guarantees

DCSL - Generalizable Crowd Counting via Diverse Context Style Learning

Collision risk estimation using stochastic motion models

AI-UPV at IberLEF-2021 DETOXIS task: Toxicity Detection in Immigration-Related Web News Comments Using Transformers and Statistical Models

Unofficial PyTorch Implementation of "DOLG: Single-Stage Image Retrieval with Deep Orthogonal Fusion of Local and Global Features"

[Arxiv preprint] Causality-inspired Single-source Domain Generalization for Medical Image Segmentation (code&data-processing pipeline)

Synthetic Scene Text from 3D Engines