• Stars
    star
    394
  • Rank 105,363 (Top 3 %)
  • Language
    Python
  • License
    Apache License 2.0
  • Created about 4 years ago
  • Updated about 1 year ago

Reviews

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

Repository Details

Semi-supervised learning for object detection

Source code for STAC: A Simple Semi-Supervised Learning Framework for Object Detection

STAC is a simple yet effective SSL framework for visual object detection along with a data augmentation strategy. STAC deploys highly confident pseudo labels of localized objects from an unlabeled image and updates the model by enforcing consistency via strong augmentation.

This code is only used for research. This is not an official Google product.

Instruction

Install dependencies

Set global enviroment variables.

export PRJROOT=/path/to/your/project/directory/STAC
export DATAROOT=/path/to/your/dataroot
export COCODIR=$DATAROOT/coco
export VOCDIR=$DATAROOT/voc
export PYTHONPATH=$PYTHONPATH:${PRJROOT}/third_party/FasterRCNN:${PRJROOT}/third_party/auto_augment:${PRJROOT}/third_party/tensorpack

Install virtual environment in the root folder of the project

cd ${PRJROOT}

sudo apt install python3-dev python3-virtualenv python3-tk imagemagick
virtualenv -p python3 --system-site-packages env3
. env3/bin/activate
pip install -r requirements.txt

# Make sure your tensorflow version is 1.14 not only in virtual environment but also in
# your machine, 1.15 can cause OOM issues.
python -c 'import tensorflow as tf; print(tf.__version__)'

# install coco apis
pip3 install 'git+https://github.com/cocodataset/cocoapi.git#subdirectory=PythonAPI'

(Optional) Install tensorpack

tensorpack with a compatible version is already included at third_party/tensorpack. bash cd ${PRJROOT}/third_party pip install --upgrade git+https://github.com/tensorpack/tensorpack.git

Download COCO/PASCAL VOC data and pre-trained models

Download data

See DATA.md

Download backbone model

cd ${COCODIR}
wget http://models.tensorpack.com/FasterRCNN/ImageNet-R50-AlignPadding.npz

Training

There are three steps:

  • 1. Train a standard detector on labeled data (detection/scripts/coco/train_stg1.sh).
  • 2. Predict pseudo boxes and labels of unlabeled data using the trained detector (detection/scripts/coco/eval_stg1.sh).
  • 3. Use labeled data and unlabeled data with pseudo labels to train a STAC detector (detection/scripts/coco/train_stg2.sh).

Besides instruction at here, detection/scripts/coco/train_stac.sh provides a combined script to train STAC.

detection/scripts/voc/train_stac.sh is a combined script to train STAC on PASCAL VOC.

The following example use labeled data as 10% train2017 and rest 90% train2017 data as unlabeled data.

Step 0: Set variables

cd ${PRJROOT}/detection

# Labeled and Unlabeled datasets
DATASET=coco_train2017.1@10
UNLABELED_DATASET=${DATASET}-unlabeled

# PATH to save trained models
CKPT_PATH=result/${DATASET}

# PATH to save pseudo labels for unlabeled data
PSEUDO_PATH=${CKPT_PATH}/PSEUDO_DATA

# Train with 8 GPUs
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

Step 1: Train FasterRCNN on labeled data

. scripts/coco/train_stg1.sh.

Set TRAIN.AUGTYPE_LAB=strong to apply strong data augmentation.

# --simple_path makes train_log/${DATASET}/${EXPNAME} as exact location to save
python3 train_stg1.py \
    --logdir ${CKPT_PATH} --simple_path --config \
    BACKBONE.WEIGHTS=${COCODIR}/ImageNet-R50-AlignPadding.npz \
    DATA.BASEDIR=${COCODIR} \
    DATA.TRAIN="('${DATASET}',)" \
    MODE_MASK=False \
    FRCNN.BATCH_PER_IM=64 \
    PREPROC.TRAIN_SHORT_EDGE_SIZE="[500,800]" \
    TRAIN.EVAL_PERIOD=20 \
    TRAIN.AUGTYPE_LAB='default'

Step 2: Generate pseudo labels of unlabeled data

. scripts/coco/eval_stg1.sh.

Evaluate using COCO metrics and save eval.json

# Check pseudo path
if [ ! -d ${PSEUDO_PATH} ]; then
    mkdir -p ${PSEUDO_PATH}
fi

# Evaluate the model for sanity check
# model-180000 is the last checkpoint
# save eval.json at $PSEUDO_PATH

python3 predict.py \
    --evaluate ${PSEUDO_PATH}/eval.json \
    --load "${CKPT_PATH}"/model-180000 \
    --config \
    DATA.BASEDIR=${COCODIR} \
    DATA.TRAIN="('${UNLABELED_DATASET}',)"

Generate pseudo labels for unlabeled data

Set EVAL.PSEUDO_INFERENCE=True to use original images rather than resized ones for inference.

# Extract pseudo label
python3 predict.py \
    --predict_unlabeled ${PSEUDO_PATH} \
    --load "${CKPT_PATH}"/model-180000 \
    --config \
    DATA.BASEDIR=${COCODIR} \
    DATA.TRAIN="('${UNLABELED_DATASET}',)" \
    EVAL.PSEUDO_INFERENCE=True

Step 3: Train STAC

. scripts/coco/train_stg2.sh.

The dataloader loads pseudo labels from ${PSEUDO_PATH}/pseudo_data.npy.

Apply default augmentation on labeled data and strong augmentation on unlabeled data.

TRAIN.CONFIDENCE and TRAIN.WU are two major parameters of the method.

python3 train_stg2.py \
    --logdir=${CKPT_PATH}/STAC --simple_path \
    --pseudo_path=${PSEUDO_PATH} \
    --config \
    BACKBONE.WEIGHTS=${COCODIR}/ImageNet-R50-AlignPadding.npz \
    DATA.BASEDIR=${COCODIR} \
    DATA.TRAIN="('${DATASET}',)" \
    DATA.UNLABEL="('${UNLABELED_DATASET}',)" \
    MODE_MASK=False \
    FRCNN.BATCH_PER_IM=64 \
    PREPROC.TRAIN_SHORT_EDGE_SIZE="[500,800]" \
    TRAIN.EVAL_PERIOD=20 \
    TRAIN.AUGTYPE_LAB='default' \
    TRAIN.AUGTYPE='strong' \
    TRAIN.CONFIDENCE=0.9 \
    TRAIN.WU=2

Tensorboard

All training logs and tensorboard info are under ${PRJROOT}/detection/train_log. Visualize using

tensorboard --logdir=${PRJROOT}/detection/train_log

Citation

@inproceedings{sohn2020detection,
  title={A Simple Semi-Supervised Learning Framework for Object Detection},
  author={Kihyuk Sohn and Zizhao Zhang and Chun-Liang Li and Han Zhang and Chen-Yu Lee and Tomas Pfister},
  year={2020},
  booktitle={arXiv:2005.04757}
}

Acknowledgement

More Repositories

1

bert

TensorFlow code and pre-trained models for BERT
Python
36,701
star
2

google-research

Google Research
Jupyter Notebook
32,494
star
3

tuning_playbook

A playbook for systematically maximizing the performance of deep learning models.
24,615
star
4

vision_transformer

Jupyter Notebook
9,288
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
5,820
star
6

arxiv-latex-cleaner

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

simclr

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

multinerf

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

football

Check out the new game server:
Python
3,230
star
10

albert

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

scenic

Scenic: A Jax Library for Computer Vision Research and Beyond
Python
2,999
star
12

frame-interpolation

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

t5x

Python
2,494
star
14

electra

ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators
Python
2,284
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,180
star
16

uda

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

pegasus

Python
1,578
star
18

big_vision

Official codebase used to develop Vision Transformer, SigLIP, MLP-Mixer, LiT and more.
Jupyter Notebook
1,555
star
19

language

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

dex-lang

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

parti

1,513
star
22

big_transfer

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

torchsde

Differentiable SDE solvers with GPU support and efficient sensitivity analysis.
Python
1,444
star
24

FLAN

Python
1,373
star
25

disentanglement_lib

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

multilingual-t5

Python
1,197
star
27

robotics_transformer

Python
1,192
star
28

planet

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

mixmatch

Python
1,126
star
30

tapas

End-to-end neural table-text understanding models.
Python
1,080
star
31

fixmatch

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

morph-net

Fast & Simple Resource-Constrained Learning of Deep Network Structure
Python
1,011
star
33

deduplicate-text-datasets

Rust
982
star
34

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
976
star
35

batch-ppo

Efficient Batched Reinforcement Learning in TensorFlow
Python
963
star
36

augmix

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

maxim

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

magvit

Official JAX implementation of MAGVIT: Masked Generative Video Transformer
Python
847
star
39

pix2seq

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

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
790
star
41

meta-dataset

A dataset of datasets for learning to learn from few examples
Python
740
star
42

noisystudent

Code for Noisy Student Training. https://arxiv.org/abs/1911.04252
Python
736
star
43

jax3d

Python
718
star
44

recsim

A Configurable Recommender Systems Simulation Platform
Python
717
star
45

lottery-ticket-hypothesis

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

rliable

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

circuit_training

Python
685
star
48

long-range-arena

Long Range Arena for Benchmarking Efficient Transformers
Python
681
star
49

federated

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

nasbench

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

prompt-tuning

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

bleurt

BLEURT is a metric for Natural Language Generation based on transfer learning.
Python
611
star
53

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
608
star
54

lasertagger

Python
603
star
55

sound-separation

Python
578
star
56

dreamer

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

robopianist

[CoRL '23] Dexterous piano playing with deep reinforcement learning.
Python
531
star
58

pix2struct

Python
530
star
59

fast-soft-sort

Fast Differentiable Sorting and Ranking
Python
527
star
60

bigbird

Transformers for Longer Sequences
Python
518
star
61

ravens

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

sam

Python
512
star
63

vmoe

Jupyter Notebook
507
star
64

batch_rl

Offline Reinforcement Learning (aka Batch Reinforcement Learning) on Atari 2600 games
Python
489
star
65

tensor2robot

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

mint

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

byt5

Python
464
star
68

adapter-bert

Python
459
star
69

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
70

robustness_metrics

Jupyter Notebook
442
star
71

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
417
star
72

receptive_field

Compute receptive fields of your favorite convnets
Python
412
star
73

maskgit

Official Jax Implementation of MaskGIT
Jupyter Notebook
376
star
74

l2p

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

nerf-from-image

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

computation-thru-dynamics

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

tf-slim

Python
360
star
78

realworldrl_suite

Real-World RL Benchmark Suite
Python
332
star
79

distilling-step-by-step

Python
325
star
80

rigl

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

python-graphs

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

weatherbench2

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

tensorflow_constrained_optimization

Python
301
star
84

task_adaptation

Python
295
star
85

exoplanet-ml

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

ibc

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

self-organising-systems

Jupyter Notebook
279
star
88

tensorflow-coder

Python
275
star
89

retvec

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

vdm

Jupyter Notebook
267
star
91

sparf

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

falken

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

syn-rep-learn

Learning from synthetic data - code and models
Python
246
star
94

lm-extraction-benchmark

Python
244
star
95

meliad

Python
231
star
96

3d-moments

Code for CVPR 2022 paper '3D Moments from Near-Duplicate Photos'
Python
229
star
97

perceiver-ar

Python
224
star
98

rlds

Jupyter Notebook
216
star
99

ott

Python
215
star
100

language-table

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