• Stars
    star
    3,295
  • Rank 13,620 (Top 0.3 %)
  • Language
    Python
  • License
    Apache License 2.0
  • Created over 3 years ago
  • Updated about 1 month ago

Reviews

There are no reviews yet. Be the first to send feedback to the community and the maintainers!

Repository Details

Scenic: A Jax Library for Computer Vision Research and Beyond

Scenic

scenic logo

Scenic is a codebase with a focus on research around attention-based models for computer vision. Scenic has been successfully used to develop classification, segmentation, and detection models for multiple modalities including images, video, audio, and multimodal combinations of them.

More precisely, Scenic is a (i) set of shared light-weight libraries solving tasks commonly encountered tasks when training large-scale (i.e. multi-device, multi-host) vision models; and (ii) several projects containing fully fleshed out problem-specific training and evaluation loops using these libraries.

Scenic is developed in JAX and uses Flax.

Contents

What we offer

Among others Scenic provides

  • Boilerplate code for launching experiments, summary writing, logging, profiling, etc;
  • Optimized training and evaluation loops, losses, metrics, bi-partite matchers, etc;
  • Input-pipelines for popular vision datasets;
  • Baseline models, including strong non-attentional baselines.

SOTA models and baselines in Scenic

There are some SOTA models and baselines in Scenic which were either developed using Scenic, or have been reimplemented in Scenic:

Projects that were developed in Scenic or used it for their experiments:

More information can be found in projects.

Baselines that were reproduced in Scenic:

More information can be found in baseline models.

Philosophy

Scenic aims to facilitate rapid prototyping of large-scale vision models. To keep the code simple to understand and extend we prefer forking and copy-pasting over adding complexity or increasing abstraction. Only when functionality proves to be widely useful across many models and tasks it may be upstreamed to Scenic's shared libraries.

Getting started

  • See projects/baselines/README.md for a walk-through baseline models and instructions on how to run the code.
  • If you would like to contribute to Scenic, please check out the Philisophy, Code structure and Contributing sections. Should your contribution be a part of the shared libraries, please send us a pull request!

Quickstart

You will need Python 3.9 or later. Download the code from GitHub

$ git clone https://github.com/google-research/scenic.git
$ cd scenic
$ pip install .

and run training for ViT on ImageNet:

$ python scenic/main.py -- \
  --config=scenic/projects/baselines/configs/imagenet/imagenet_vit_config.py \
  --workdir=./

Note that for specific projects and baselines, you might need to install extra packages that are mentioned in their README.md or requirements.txt files.

Here is also a minimal colab to train a simple feed-forward model using Scenic.

Scenic component design

Scenic is designed to propose different levels of abstraction, to support hosting projects that only require changing hyper-parameters by defining config files, to those that need customization on the input pipeline, model architecture, losses and metrics, and the training loop. To make this happen, the code in Scenic is organized as either project-level code, which refers to customized code for specific projects or baselines or library-level code, which refers to common functionalities and general patterns that are adapted by the majority of projects. The project-level code lives in the projects directory.

scenic design

Library-level code

The goal is to keep the library-level code minimal and well-tested and to avoid introducing extra abstractions to support minor use-cases. Shared libraries provided by Scenic are split into:

  • dataset_lib: Implements IO pipelines for loading and pre-processing data for common Computer Vision tasks and benchmarks (see "Tasks and Datasets" section). All pipelines are designed to be scalable and support multi-host and multi-device setups, taking care dividing data among multiple hosts, incomplete batches, caching, pre-fetching, etc.
  • model_lib : Provides
    • several abstract model interfaces (e.g. ClassificationModel or SegmentationModel in model_lib.base_models) with task-specific losses and metrics;
    • neural network layers in model_lib.layers, focusing on efficient implementation of attention and transformer layers;
    • accelerator-friendly implementations of bipartite matching algorithms in model_lib.matchers.
  • train_lib: Provides tools for constructing training loops and implements several optimized trainers (classification trainer and segmentation trainer) that can be forked for customization.
  • common_lib: General utilities, like logging and debugging modules, functionalities for processing raw data, etc.

Project-level code

Scenic supports the development of customized solutions for customized tasks and data via the concept of "project". There is no one-fits-all recipe for how much code should be re-used by a project. Projects can consist of only configs and use the common models, trainers, task/data that live in library-level code, or they can simply fork any of the mentioned functionalities and redefine, layers, losses, metrics, logging methods, tasks, architectures, as well as training and evaluation loops. The modularity of library-level code makes it flexible for projects to fall placed on any spot in the "run-as-is" to "fully customized" spectrum.

Common baselines such as a ResNet and Vision Transformer (ViT) are implemented in the projects/baselines project. Forking models in this directory is a good starting point for new projects.

Citing Scenic

If you use Scenic, you can cite our white paper. Here is an example BibTeX entry:

@InProceedings{dehghani2021scenic,
    author    = {Dehghani, Mostafa and Gritsenko, Alexey and Arnab, Anurag and Minderer, Matthias and Tay, Yi},
    title     = {Scenic: A JAX Library for Computer Vision Research and Beyond},
    booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    year      = {2022},
    pages     = {21393-21398}
}

Disclaimer: This is not an official Google product.

More Repositories

1

bert

TensorFlow code and pre-trained models for BERT
Python
37,769
star
2

google-research

Google Research
Jupyter Notebook
33,759
star
3

tuning_playbook

A playbook for systematically maximizing the performance of deep learning models.
26,593
star
4

vision_transformer

Jupyter Notebook
10,251
star
5

text-to-text-transfer-transformer

Code for the paper "Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer"
Python
6,099
star
6

arxiv-latex-cleaner

arXiv LaTeX Cleaner: Easily clean the LaTeX code of your paper to submit to arXiv
Python
5,233
star
7

simclr

SimCLRv2 - Big Self-Supervised Models are Strong Semi-Supervised Learners
Jupyter Notebook
3,937
star
8

multinerf

A Code Release for Mip-NeRF 360, Ref-NeRF, and RawNeRF
Python
3,612
star
9

timesfm

TimesFM (Time Series Foundation Model) is a pretrained time-series foundation model developed by Google Research for time-series forecasting.
Python
3,576
star
10

football

Check out the new game server:
Python
3,260
star
11

albert

ALBERT: A Lite BERT for Self-supervised Learning of Language Representations
Python
3,209
star
12

frame-interpolation

FILM: Frame Interpolation for Large Motion, In ECCV 2022.
Python
2,818
star
13

t5x

Python
2,656
star
14

electra

ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators
Python
2,325
star
15

kubric

A data generation pipeline for creating semi-realistic synthetic multi-object videos with rich annotations such as instance segmentation masks, depth maps, and optical flow.
Jupyter Notebook
2,312
star
16

big_vision

Official codebase used to develop Vision Transformer, SigLIP, MLP-Mixer, LiT and more.
Jupyter Notebook
2,219
star
17

uda

Unsupervised Data Augmentation (UDA)
Python
2,131
star
18

language

Shared repository for open-sourced projects from the Google AI Language team.
Python
1,605
star
19

pegasus

Python
1,600
star
20

dex-lang

Research language for array processing in the Haskell/ML family
Haskell
1,581
star
21

torchsde

Differentiable SDE solvers with GPU support and efficient sensitivity analysis.
Python
1,548
star
22

parti

1,538
star
23

big_transfer

Official repository for the "Big Transfer (BiT): General Visual Representation Learning" paper.
Python
1,504
star
24

FLAN

Python
1,460
star
25

robotics_transformer

Python
1,337
star
26

disentanglement_lib

disentanglement_lib is an open-source library for research on learning disentangled representations.
Python
1,311
star
27

multilingual-t5

Python
1,197
star
28

circuit_training

Python
1,151
star
29

tapas

End-to-end neural table-text understanding models.
Python
1,143
star
30

planet

Learning Latent Dynamics for Planning from Pixels
Python
1,134
star
31

mixmatch

Python
1,130
star
32

deduplicate-text-datasets

Rust
1,104
star
33

fixmatch

A simple method to perform semi-supervised learning with limited data.
Python
1,094
star
34

morph-net

Fast & Simple Resource-Constrained Learning of Deep Network Structure
Python
1,016
star
35

maxim

[CVPR 2022 Oral] Official repository for "MAXIM: Multi-Axis MLP for Image Processing". SOTA for denoising, deblurring, deraining, dehazing, and enhancement.
Python
996
star
36

deeplab2

DeepLab2 is a TensorFlow library for deep labeling, aiming to provide a unified and state-of-the-art TensorFlow codebase for dense pixel labeling tasks.
Python
995
star
37

batch-ppo

Efficient Batched Reinforcement Learning in TensorFlow
Python
963
star
38

augmix

AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty
Python
951
star
39

magvit

Official JAX implementation of MAGVIT: Masked Generative Video Transformer
Python
947
star
40

pix2seq

Pix2Seq codebase: multi-tasks with generative modeling (autoregressive and diffusion)
Jupyter Notebook
865
star
41

seed_rl

SEED RL: Scalable and Efficient Deep-RL with Accelerated Central Inference. Implements IMPALA and R2D2 algorithms in TF2 with SEED's architecture.
Python
793
star
42

meta-dataset

A dataset of datasets for learning to learn from few examples
Jupyter Notebook
762
star
43

noisystudent

Code for Noisy Student Training. https://arxiv.org/abs/1911.04252
Python
751
star
44

rliable

[NeurIPS'21 Outstanding Paper] Library for reliable evaluation on RL and ML benchmarks, even with only a handful of seeds.
Jupyter Notebook
747
star
45

recsim

A Configurable Recommender Systems Simulation Platform
Python
739
star
46

jax3d

Python
733
star
47

long-range-arena

Long Range Arena for Benchmarking Efficient Transformers
Python
719
star
48

lottery-ticket-hypothesis

A reimplementation of "The Lottery Ticket Hypothesis" (Frankle and Carbin) on MNIST.
Python
706
star
49

federated

A collection of Google research projects related to Federated Learning and Federated Analytics.
Python
675
star
50

bleurt

BLEURT is a metric for Natural Language Generation based on transfer learning.
Python
651
star
51

prompt-tuning

Original Implementation of Prompt Tuning from Lester, et al, 2021
Python
642
star
52

nasbench

NASBench: A Neural Architecture Search Dataset and Benchmark
Python
641
star
53

neuralgcm

Hybrid ML + physics model of the Earth's atmosphere
Python
641
star
54

xtreme

XTREME is a benchmark for the evaluation of the cross-lingual generalization ability of pre-trained multilingual models that covers 40 typologically diverse languages and includes nine tasks.
Python
631
star
55

lasertagger

Python
606
star
56

sound-separation

Python
603
star
57

pix2struct

Python
587
star
58

vmoe

Jupyter Notebook
569
star
59

dreamer

Dream to Control: Learning Behaviors by Latent Imagination
Python
568
star
60

robopianist

[CoRL '23] Dexterous piano playing with deep reinforcement learning.
Python
562
star
61

omniglue

Code release for CVPR'24 submission 'OmniGlue'
Python
561
star
62

fast-soft-sort

Fast Differentiable Sorting and Ranking
Python
561
star
63

ravens

Train robotic agents to learn pick and place with deep learning for vision-based manipulation in PyBullet. Transporter Nets, CoRL 2020.
Python
560
star
64

sam

Python
551
star
65

batch_rl

Offline Reinforcement Learning (aka Batch Reinforcement Learning) on Atari 2600 games
Python
521
star
66

bigbird

Transformers for Longer Sequences
Python
518
star
67

tensor2robot

Distributed machine learning infrastructure for large-scale robotics research
Python
483
star
68

byt5

Python
477
star
69

adapter-bert

Python
476
star
70

mint

Multi-modal Content Creation Model Training Infrastructure including the FACT model (AI Choreographer) implementation.
Python
465
star
71

leaf-audio

LEAF is a learnable alternative to audio features such as mel-filterbanks, that can be initialized as an approximation of mel-filterbanks, and then be trained for the task at hand, while using a very small number of parameters.
Python
446
star
72

robustness_metrics

Jupyter Notebook
442
star
73

maxvit

[ECCV 2022] Official repository for "MaxViT: Multi-Axis Vision Transformer". SOTA foundation models for classification, detection, segmentation, image quality, and generative modeling...
Jupyter Notebook
436
star
74

receptive_field

Compute receptive fields of your favorite convnets
Python
434
star
75

maskgit

Official Jax Implementation of MaskGIT
Jupyter Notebook
429
star
76

weatherbench2

A benchmark for the next generation of data-driven global weather models.
Python
420
star
77

l2p

Learning to Prompt (L2P) for Continual Learning @ CVPR22 and DualPrompt: Complementary Prompting for Rehearsal-free Continual Learning @ ECCV22
Python
408
star
78

distilling-step-by-step

Python
407
star
79

ssl_detection

Semi-supervised learning for object detection
Python
398
star
80

nerf-from-image

Shape, Pose, and Appearance from a Single Image via Bootstrapped Radiance Field Inversion
Python
377
star
81

computation-thru-dynamics

Understanding computation in artificial and biological recurrent networks through the lens of dynamical systems.
Jupyter Notebook
369
star
82

tf-slim

Python
368
star
83

realworldrl_suite

Real-World RL Benchmark Suite
Python
341
star
84

python-graphs

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

rigl

End-to-end training of sparse deep neural networks with little-to-no performance loss.
Python
314
star
86

task_adaptation

Python
310
star
87

self-organising-systems

Jupyter Notebook
308
star
88

ibc

Official implementation of Implicit Behavioral Cloning, as described in our CoRL 2021 paper, see more at https://implicitbc.github.io/
Python
306
star
89

tensorflow_constrained_optimization

Python
300
star
90

syn-rep-learn

Learning from synthetic data - code and models
Python
294
star
91

arco-era5

Recipes for reproducing Analysis-Ready & Cloud Optimized (ARCO) ERA5 datasets.
Python
291
star
92

vdm

Jupyter Notebook
291
star
93

rlds

Jupyter Notebook
284
star
94

exoplanet-ml

Machine learning models and utilities for exoplanet science.
Python
283
star
95

retvec

RETVec is an efficient, multilingual, and adversarially-robust text vectorizer.
Jupyter Notebook
281
star
96

sparf

This is the official code release for SPARF: Neural Radiance Fields from Sparse and Noisy Poses [CVPR 2023-Highlight]
Python
279
star
97

tensorflow-coder

Python
275
star
98

lm-extraction-benchmark

Python
270
star
99

language-table

Suite of human-collected datasets and a multi-task continuous control benchmark for open vocabulary visuolinguomotor learning.
Jupyter Notebook
260
star
100

falken

Falken provides developers with a service that allows them to train AI that can play their games
Python
254
star