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

Object recognition using Azure Custom Vision AI and Azure Functions

This repository contains small projects related to Neural Networks and Deep Learning in general.

Implementation of SSMF: Shifting Seasonal Matrix Factorization

Plover-tapey-tape: an alternative to Plover’s built-in paper tape

ICCV2021, Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet

Interactive Visualization to empower domain experts to align ML model behaviors with their knowledge.

Sentiment analysis translations of the Bhagavad Gita

All public open-source implementations of convnets benchmarks

tensorflow implementation of 'YOLO : Real-Time Object Detection'

[SIGGRAPH 2021 Asia] DeepVecFont: Synthesizing High-quality Vector Fonts via Dual-modality Learning

Asterisk is a framework to generate high-quality training datasets at scale

Deep Learning to Create StepMania SM FIles

Discord bot-CTFD-Thread-Parser - Discord bot CTFD-Thread-Parser

A more easy-to-use implementation of KPConv based on PyTorch.

Introducing neural networks to predict stock prices

This is the source code for the experiments related to the paper Unsupervised Audio Source Separation Using Differentiable Parametric Source Models

Very Deep Convolutional Networks for Large-Scale Image Recognition

Simple image captioning model - CLIP prefix captioning.

An implementation for the loss function proposed in Decoupled Contrastive Loss paper.

Self-supervised Label Augmentation via Input Transformations (ICML 2020)