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

Implementation of the paper NAST: Non-Autoregressive Spatial-Temporal Transformer for Time Series Forecasting.

Low Complexity Channel estimation with Neural Network Solutions

The official PyTorch code implementation of "Personalized Trajectory Prediction via Distribution Discrimination" in ICCV 2021.

Convert BART models to ONNX with quantization. 3X reduction in size, and upto 3X boost in inference speed

Exact Pareto Optimal solutions for preference based Multi-Objective Optimization

[AI6101] Introduction to AI & AI Ethics is a core course of MSAI, SCSE, NTU, Singapore

Torch-based tool for quantizing high-dimensional vectors using additive codebooks

OpenMMLab Computer Vision Foundation

A lightweight library designed to accelerate the process of training PyTorch models by providing a minimal

IOT: Instance-wise Layer Reordering for Transformer Structures

Selecting Parallel In-domain Sentences for Neural Machine Translation Using Monolingual Texts

The official implementation for ACL 2021 "Challenges in Information Seeking QA: Unanswerable Questions and Paragraph Retrieval".

Customizable RecSys Simulator for OpenAI Gym

This is the code related to "Sparse-to-dense Feature Matching: Intra and Inter domain Cross-modal Learning in Domain Adaptation for 3D Semantic Segmentation" (ICCV 2021).

Pytorch implementation for the Temporal and Object Quantification Networks (TOQ-Nets).

Combinatorial model of ligand-receptor binding

An implementation of Deep Forest 2021.2.1.

PyTorch Implementation of "Non-Autoregressive Neural Machine Translation"

A simple baseline for 3d human pose estimation in tensorflow. Presented at ICCV 17.

🦕 NanoSaur is a little tracked robot ROS2 enabled, made for an NVIDIA Jetson Nano