Git Product home page Git Product logo

scrl's Introduction

KakaoBrain pytorch pytorch

Spatially Consistent Representation Learning (CVPR'21)

Abstract

SCRL is a self-supervised learning method that allows you to obtain a spatially consistent dense representation, especially useful for localization tasks such as object detection and segmentation. You might be able to improve your own localization task by simply initializing the backbone model with those parameters trained using our method. Please refer to the paper for more details.

Requirements

We have tested the code on the following environments:

  • Python 3.7.7 / Pytorch 1.6.0 / torchvision 0.7.0 / CUDA 10.1 / Ubuntu 18.04
  • Python 3.8.3 / Pytorch 1.7.1 / torchvision 0.8.2 / CUDA 11.1 / Ubuntu 18.04

Run the following command to install dependencies:

pip install -r requirements.txt

Configuration

The config directory contains predefined YAML configuration files for different learning methods and schedules.

There are two alternative ways to change training configuration:

  • Change the value of the field of interest directly in config/*.yaml files.
  • Pass the value as a command argument to main.py with the name that represents the field's hierarchy using the delimiter /.

Note that the values given by the command argument take precedence over the ones stored in the YAML file.

Refer to the Dataset subsection below for an example.

Model

We officially support ResNet-50 and ResNet-101 backbones. Note that all the hyperparameters have been tuned based on ResNet-50.

Dataset

Currently, only ImageNet dataset is supported. You must specify the dataset path in one of the following ways:

# (option 1) set field `dataset.root` in the YAML configuration file.
data:
  root: your_own_path
  ...
# (option 2) pass `--dataset/root` as an argument of the shell command.
$ python main.py --data/root your_own_path ...

How to Run

Overview

The code consists of two main parts: self-supervised pre-training(upstream) and linear evaluation protocol(downstream). After the pre-training is done, the evaluation protocol is automatically run by the default configurations. Note that, for simplicity, this code does not hold out validation set from train set as in BYOL paper. Although this may not be a rigorous implementation, training the linear classifier itself is not in our main interest here and the protocol still does its job well enough as an evaluation metric. Also note that the protocol may not be exactly the same with the details in other literatures in terms of batch size and learning rate schedule.

Resource & Batchsize

We recommend you run the code in a multi-node environment in order to reproduce the results reported in the paper. We used 4 V100 x 8 nodes = 32 GPUs for training. Under this circumstance, our upstream batch size is 8192 and downstream batch size is 4096. When the number of GPUs dwindled to half, we observed performance degradation. Although BYOL-like methods do not use negative examples in an explicit way, they can still suffer performance drops when their batch size is reduced, as illustrated in Figure 3 in BYOL paper.

Single-node Training

If your YAML configuration file is ./confing/scrl_200ep.yaml, run the command as follows:

$ python main.py --config config/scrl_200_ep.yaml

Multi-node Training

To train a single model with 2 nodes, for instance, run the commands below in sequence:

# on the machine #0
$ python main.py --config config/scrl_200_ep.yaml \
                 --dist_url tcp://{machine_0_ip}:{available_port} \
                 --num_machines 2 \
                 --machine_rank 0
# on the machine #1
$ python main.py --config config/scrl_200_ep.yaml \
                 --dist_url tcp://{machine_1_ip}:{available_port} \
                 --num_machines 2 \
                 --machine_rank 1

If IP address and port number are not known in advance, you can first run the command with --dist_url=auto in the master node. Then check the IP address and available port number that are printed on the command line, to which you should refer to launch the other nodes.

Linear Evaluation Only

You can also run the protocol using any given checkpoint on a stand-alone basis. You can evaluate the latest checkpoint anytime after the very first checkpoint has been dumped. Refer to the following command:

$ python main.py --config config/scrl_200ep.yaml --train/enabled=False --load_dir ...

Saving & Loading Checkpoints

Saved Filenames

  • save_dir will be automatically determined(with sequential number suffixes) unless otherwise designated.
  • Model's checkpoints are saved in ./{save_dir}/checkpoint_{epoch}.pth.
  • Symlinks of the last checkpoints are saved in ./{save_dir}/checkpoint_last.pth.

Automatic Loading

  • SCRLTraniner will automatically load checkpoint_last.pth if it exists in save_dir.
  • By default, save_dir is identical to load_dir. However, you can also set load_dir seperately.

Results

Method Epoch Linear Eval (Acc.) COCO BBox (AP) COCO Segm (AP) Checkpoint
Random -- --  /   --   29.77 / 30.95 28.70 --
IN-sup. 90 --  /  74.3 38.52 / 39.00 35.44 --
BYOL 200 72.90 / 73.14 38.35 / 38.86 35.96 download
BYOL 1000 74.47 / 74.51 40.10 / 40.19 37.16 download
SCRL 200 66.78 / 68.27 40.49 / 41.02 37.50 download
SCRL 1000 70.67 / 70.66 40.92 / 41.40 37.92 download
  • Epoch: self-supervised pre-training epochs
  • Linear Eval(linear evaluation): online / offline
  • COCO Bbox(Object Detection): Faster R-CNN w.FPN / Mask R-CNN w.FPN
  • COCO Segm(Instance Segmentation): Mask R-CNN w.FPN

On/offline Linear Evaluation

The online evaluation is done with a linear classifier attached to the top of the backbone, which is trained simultaneously during pre-training under its own objective but does NOT backpropagate the gradients to the main model. This facilitates monitoring of the learning progress while there is no target performance measure for self-supervised learning. The offline evaluation refers to the commonly known standard protocol.

COCO Localization Tasks

As you can see in the table, our method significantly boosts the performance of the localization downstream tasks. Note that the values in the table can be slightly different from the paper because of different random seeds. For the COCO localization tasks, we used Detectron2 repository publicly available, which is not included in this code. We simply initialized the ResNet-50 backbone with pre-trained parameters by our method and finetuned it under the downstream objective. Note that we used synchronized BN for all the configurable layers and retained the default hyperparameters for everything else including the training schedule(x1).

Downloading Pretrained Checkpoints

You can download the backbone checkpoints via those links in the table. To load them in our code and run the linear evaluation, change the filename to checkpoint_last.pth (or make symlink) and pass the parent directory path to either save_dir or load_dir.

Hyperparameters

For the 1000 epoch checkpoint of SCRL, we simply used the hyperparameters adopted in the official BYOL implementation, which is different from the description in the paper, but still matches the reported performance. As described in the Appendix in the paper, there could be potential room for improvement by extensive hyperparameter search.

Citation

If you use this code for your research, please cite our paper.

@inproceedings{roh2021scrl,
  title={Spatilly Consistent Representation Learning},
  author    = {Roh, Byungseok and
               Shin, Wuhyun and
               Kim, Ildoo and
               Kim, Sungwoong},
  booktitle = {CVPR},
  publisher = {IEEE},
  year      = {2021}
}

Contact for Issues

License

This project is licensed under the terms of the Apache License 2.0.

Copyright 2021 Kakao Brain Corp. https://www.kakaobrain.com All Rights Reserved.

scrl's People

Contributors

prayforher avatar ricoshin avatar

Stargazers

 avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar

Watchers

 avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar  avatar

scrl's Issues

problem about generating bboxes in intersection area

Thanks for sharing such wonderful work! it's so amazing! Now I am trying to reproduce your work, but I meet some problem about generating bboxes in intersection area. According to your paper, you use some operation similar to RandomResizedCrop to generate two views and compute intersection between these two views, but how to ensure these two views always have overlaps. In addition, I found that even if two views have overlaps, there are some image which can not generate enough bboxes, for example, 10. Is there some other operation to solve this problem? Can you give me some suggestion?

Questions about the feature maps

Thanks for your great repo! I have the following questions about the code:

  1. In both online and target network, the feature map is learned before finding ROIs. I would like to know how we can extract these feature maps?
  2. In your studies you mentioned that you used this work in localization tasks (object detection with Faster RCNN). Do you mind to share some sample scripts that how you used your method for object detection?

Thanks,

Do ur codes using Sync_Batchnorm in default?

First of all , i am very interested in your paper.
Since your model is based on BYOL, I think the BN layer is very important in projector and predictor, and it's beneficial from large batch-size.
Since your code is working with DDP, i didn't find the Sync_Batchnorm.....i think it's weird.... If u don't use Sync-BN, so the BN is calculated only in each single gpu but not share with other gpus.
I think maybe i make a mistake, so do ur codes using Sync_Batchnorm in default? Should we use SyncBN?
Hope anyone can answer my question, Thx!

error loading "checkpoint_best.pth" model

After the upstream training~
when i load the "checkpoint_best.pth", I got an error like [enforce fail at inline_container.cc:209]. file not found: archive/data/3259494704. I don't know exactly why the file was somehow corrupted.

How did you use unlabeled dataset?

Paper p. 12
Table A3. COCO detection using Faster R-CNN, ResNet-50-FPN.Upstreams are trained with the unlabeled COCO dataset with 2000epochs.

In your code, it have to use labels like classification task, but in the paper, you described you used unlabeled COCO dataset.
Doesn't it care about labels? Then, can I use like this?

data/mydataset/0/image_0.jpg
data/mydataset/0/image_1.jpg
data/mydataset/0/...
data/mydataset/1/image_0.jpg
data/mydataset/1/image_1.jpg
data/mydataset/1/...

(I changed imageNet to imageFolder.)

Try to pretrain on detection dataset

Hello, Thanks for sharing your codes ~
I try to train on VOC2012 without label to get the pretrained model, and then finetune with VOC2012 detection label. The result is very poor, but the result should be better Since pretrain dataset and finetune dataset is the same. Did you tried this or do you have any idea that what's the reason?

Recommend Projects

  • React photo React

    A declarative, efficient, and flexible JavaScript library for building user interfaces.

  • Vue.js photo Vue.js

    🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.

  • Typescript photo Typescript

    TypeScript is a superset of JavaScript that compiles to clean JavaScript output.

  • TensorFlow photo TensorFlow

    An Open Source Machine Learning Framework for Everyone

  • Django photo Django

    The Web framework for perfectionists with deadlines.

  • D3 photo D3

    Bring data to life with SVG, Canvas and HTML. 📊📈🎉

Recommend Topics

  • javascript

    JavaScript (JS) is a lightweight interpreted programming language with first-class functions.

  • web

    Some thing interesting about web. New door for the world.

  • server

    A server is a program made to process requests and deliver data to clients.

  • Machine learning

    Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.

  • Game

    Some thing interesting about game, make everyone happy.

Recommend Org

  • Facebook photo Facebook

    We are working to build community through open source technology. NB: members must have two-factor auth.

  • Microsoft photo Microsoft

    Open source projects and samples from Microsoft.

  • Google photo Google

    Google ❤️ Open Source for everyone.

  • D3 photo D3

    Data-Driven Documents codes.