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

Tracking Progress in Question Answering over Knowledge Graphs

Efficient 6-DoF Grasp Generation in Cluttered Scenes

Implementation for the "Surface Reconstruction from 3D Line Segments" paper.

Visual odometry package based on hardware-accelerated NVIDIA Elbrus library with world class quality and performance.

Implementation of CrossViT: Cross-Attention Multi-Scale Vision Transformer for Image Classification

Explicable Reward Design for Reinforcement Learning Agents [NeurIPS'21]

we propose a novel deep network, named feature aggregation and refinement network (FARNet), for the automatic detection of anatomical landmarks.

NVIDIA Deep Learning Examples for Tensor Cores

Code in PyTorch for the convex combination linear IAF and the Householder Flow, J.M. Tomczak & M. Welling

Compact Bidirectional Transformer for Image Captioning

TUPÃ was developed to analyze electric field properties in molecular simulations

The official repo of the CVPR 2021 paper Group Collaborative Learning for Co-Salient Object Detection .

Speech recognition tool to convert audio to text transcripts, for Linux and Raspberry Pi.

NVTabular is a feature engineering and preprocessing library for tabular data designed to quickly and easily manipulate terabyte scale datasets used to train deep learning based recommender systems.

Deep Learning for Time Series Classification

[CVPR 2021] Released code for Counterfactual Zero-Shot and Open-Set Visual Recognition

Repository relating to the CVPR21 paper TimeLens: Event-based Video Frame Interpolation

Pytorch implementation of U-Net, R2U-Net, Attention U-Net, and Attention R2U-Net.

🐾 Semantic segmentation of paws from cute pet images (PyTorch)

Trainable PyTorch reproduction of AlphaFold 2