[飞桨论文复现挑战赛(第六期)] (99) Exploring Simple Siamese Representation Learning (#117)

* add simsiam

* add no_bias arg in nonlinearneckv2

* add tipc

* update tipc

* update simsiam optimize builder

* add lars apx support

Author: fuqianya

README.md

@@ -46,6 +46,7 @@ PASSL implements a series of self-supervised learning algorithms, See **Document
 | MoCo-BYOL | 300    | 71.56            | 72.10         | ResNet-50 | [download](https://passl.bj.bcebos.com/models/mocobyol_r50_ep300_ckpt.pdparams) | [Train MoCo-BYOL](docs/Train_MoCo-BYOL_model.md) |
 | BYOL      | 300    | 72.50            | 71.62         | ResNet-50 | [download](https://passl.bj.bcebos.com/models/byol_r50_300.pdparams) | [Train BYOL](docs/Train_BYOL_model.md)           |
 | PixPro    | 100    | 55.1(fp16)       | 57.2(fp32)    | ResNet-50 | [download](https://passl.bj.bcebos.com/models/pixpro_r50_ep100_no_instance_with_linear.pdparams) | [Train PixPro](docs/Train_PixPro_model.md)       |
+| SimSiam   | 100    | 68.3             | 68.4          | ResNet-50 | [download](https://drive.google.com/file/d/1kaAm8-tlvB570kzI4fo9h4dwGQFf_4FE/view?usp=sharing) | [Train SimSiam](docs/Train_SimSiam_model.md)      |
 
 > Benchmark Linear Image Classification on ImageNet-1K.
 

README_cn.md

@@ -46,6 +46,7 @@ PASSL 实现了一系列自监督学习算法，更具体的使用文档请参
 | MoCo-BYOL | 300    | 71.56            | 72.10         | ResNet-50 | [download](https://passl.bj.bcebos.com/models/mocobyol_r50_ep300_ckpt.pdparams) | [Train MoCo-BYOL](docs/Train_MoCo-BYOL_model.md) |
 | BYOL      | 300    | 72.50            | 71.62         | ResNet-50 | [download](https://passl.bj.bcebos.com/models/byol_r50_300.pdparams) | [Train BYOL](docs/Train_BYOL_model.md)           |
 | PixPro    | 100    | 55.1(fp16)       | 57.2(fp32)    | ResNet-50 | [download](https://passl.bj.bcebos.com/models/pixpro_r50_ep100_no_instance_with_linear.pdparams) | [Train PixPro](docs/Train_PixPro_model.md)       |
+| SimSiam   | 100    | 68.3             | 68.4          | ResNet-50 | [download](https://drive.google.com/file/d/1kaAm8-tlvB570kzI4fo9h4dwGQFf_4FE/view?usp=sharing) | [Train SimSiam](docs/Train_SimSiam_model.md)      |
 
 > Benchmark Linear Image Classification on ImageNet-1K.
 

configs/simsiam/simsiam_clas_r50.yaml

@@ -0,0 +1,87 @@
+epochs: 90
+use_simclr_iters: True
+global_batch_size: 4096
+output_dir: output_dir
+seed: 0
+device: gpu
+
+# used for static mode and model export
+image_shape: [3, 224, 224]
+save_inference_dir: ./inference
+
+model:
+  name: Classification
+  backbone:
+    name: ResNet
+    depth: 50
+    frozen_stages: 4
+  head:
+    name: ClasHead
+    with_avg_pool: true
+    in_channels: 2048
+
+dataloader:
+  train:
+    loader:
+      num_workers: 16
+      use_shared_memory: True
+    sampler:
+      batch_size: 512
+      shuffle: true
+      drop_last: true
+    dataset:
+      name: ImageNet
+      dataroot: data/ILSVRC2012/train
+      return_label: True
+      transforms:
+        - name: RandomResizedCrop
+          size: 224
+        - name: RandomHorizontalFlip
+        - name: Transpose
+        - name: NormalizeImage
+          scale: 1.0/255.0
+          mean: [0.485, 0.456, 0.406]
+          std: [0.229, 0.224, 0.225]
+  val:
+    loader:
+      num_workers: 16
+      use_shared_memory: True
+    sampler:
+      batch_size: 512
+      shuffle: false
+      drop_last: false
+    dataset:
+      name: ImageNet
+      dataroot: data/ILSVRC2012/val
+      return_label: True
+      transforms:
+        - name: Resize
+          size: 256
+        - name: CenterCrop
+          size: 224
+        - name: Transpose
+        - name: NormalizeImage
+          scale: 1.0/255.0
+          mean: [0.485, 0.456, 0.406]
+          std: [0.229, 0.224, 0.225]
+
+lr_scheduler:
+  name: Cosinesimclr
+  learning_rate: 1.6
+  T_max: 90
+
+optimizer:
+  name: LarsMomentumOptimizer
+  momentum: 0.9
+  lars_weight_decay: 0.0
+
+log_config:
+    name: LogHook
+    interval: 50
+
+lr_config:
+  name: LRSchedulerHook
+  unit: epoch
+
+custom_config:
+  - name: EvaluateHook

configs/simsiam/simsiam_r50.yaml

@@ -0,0 +1,107 @@
+epochs: 100
+use_simclr_iters: True
+global_batch_size: 512
+output_dir: output_dir
+seed: 0
+device: gpu
+
+model:
+  name: SimSiam
+  backbone:
+    name: ResNet
+    depth: 50
+    with_pool: True
+    num_classes: 2048
+    zero_init_residual: True
+  predictor:
+    name: NonLinearNeckV2
+    in_channels: 2048
+    hid_channels: 512
+    out_channels: 2048
+    with_bias: False
+    with_avg_pool: False
+  head:
+    name: SimSiamContrastiveHead
+
+dataloader:
+  train:
+    loader:
+      num_workers: 16
+      use_shared_memory: True
+    sampler:
+      batch_size: 64
+      shuffle: true
+      drop_last: true
+    dataset:
+      name: ImageNet
+      dataroot: data/ILSVRC2012/train
+      return_label: False
+      return_two_sample: True
+      transforms:
+        - name: RandomResizedCrop
+          size: 224
+          scale: [0.2, 1.]
+      view_trans1:
+        - name: RandomApply
+          transforms:
+          - name: ColorJitter
+            brightness: 0.4
+            contrast: 0.4
+            saturation: 0.4
+            hue: 0.1
+          p: 0.8
+        - name: RandomGrayscale
+          p: 0.2
+        - name: RandomApply
+          transforms:
+          - name: GaussianBlur
+            sigma: [0.1, 2.0]
+          p: 0.5
+        - name: RandomHorizontalFlip
+        - name: Transpose
+        - name: NormalizeImage
+          scale: 1.0/255.0
+          mean: [0.485, 0.456, 0.406]
+          std: [0.229, 0.224, 0.225]
+      view_trans2:
+        - name: RandomApply
+          transforms:
+          - name: ColorJitter
+            brightness: 0.4
+            contrast: 0.4
+            saturation: 0.4
+            hue: 0.1
+          p: 0.8
+        - name: RandomGrayscale
+          p: 0.2
+        - name: RandomApply
+          transforms:
+          - name: GaussianBlur
+            sigma: [0.1, 2.0]
+          p: 0.5
+        - name: RandomHorizontalFlip
+        - name: Transpose
+        - name: NormalizeImage
+          scale: 1.0/255.0
+          mean: [0.485, 0.456, 0.406]
+          std: [0.229, 0.224, 0.225]
+
+lr_scheduler:
+  name: Cosinesimclr
+  learning_rate: 0.1
+  T_max: 100
+
+optimizer:
+  name: Momentum
+  weight_decay: 0.0001
+
+optimizer_config:
+  name: SimsiamOptimizerHook
+
+log_config:
+    name: LogHook
+    interval: 50
+
+lr_config:
+  name: LRSchedulerHook
+  unit: epoch

docs/Train_SimSiam_model.md

@@ -0,0 +1,57 @@
+# Train SimSiam Model
+
+## Introduction
+
+PASSL reproduces [SimSiam](https://arxiv.org/abs/2011.10566), which is a simsiam network for unsupervised visual representation learning.
+
+## Installation
+- See [INSTALL.md](INSTALL.md)
+
+## Data Preparation
+- See [GETTING_STARTED.md](GETTING_STARTED.md)
+
+## Implemented Models
+Models are all trained with ResNet-50 backbone.
+|  | epochs |official results | passl results | Backbone| Model |
+| ---|--- | ----  | ---- | ----| ---- |
+| SimSiam | 100 | 68.3 | 68.4          | ResNet-50 | [download](https://drive.google.com/file/d/1kaAm8-tlvB570kzI4fo9h4dwGQFf_4FE/view?usp=sharing) |
+
+
+## Getting Started
+
+### 1. Train SimSiam
+
+#### single gpu
+```
+python tools/train.py -c configs/simsiam/simsiam_r50.yaml
+```
+
+#### multiple gpus
+
+```
+python -m paddle.distributed.launch --gpus="0,1,2,3,4,5,6,7" tools/train.py -c configs/simsiam/simsiam_r50.yaml
+```
+
+Pretraining models with 100 epochs can be found at [simsiam](https://drive.google.com/file/d/1kaAm8-tlvB570kzI4fo9h4dwGQFf_4FE/view?usp=sharing)
+
+Note: The default learning rate in config files is for 8 GPUs. If using differnt number GPUs, the total batch size will change in proportion, you have to scale the learning rate following ```new_lr = old_lr * new_ngpus / old_ngpus```.
+
+### 2. Extract backbone weights
+
+```
+python tools/extract_weight.py ${CHECKPOINT} --output ${WEIGHT_FILE} --prefix encoder --remove_prefix
+```
+
+### 3. Evaluation on ImageNet Linear Classification
+
+#### Train:
+```
+python -m paddle.distributed.launch --gpus="0,1,2,3,4,5,6,7" tools/train.py -c configs/simsiam/simsiam_clas_r50.yaml --pretrained ${WEIGHT_FILE}
+```
+
+#### Evaluate:
+```
+python -m paddle.distributed.launch --gpus="0,1,2,3,4,5,6,7" tools/train.py -c configs/simsiam/simsiam_clas_r50.yaml --load ${CLS_WEGHT_FILE} --evaluate-only
+```
+
+The trained linear weights in conjuction with the backbone weights can be found at [simsiam linear](https://drive.google.com/file/d/19smHZGhBEPWeyLjKIGhM7KPngr-8BOUl/view?usp=sharing)

passl/engine/trainer.py

@@ -135,6 +135,7 @@ def __init__(self, cfg):
 
         n_parameters = sum(p.numel() for p in self.model.parameters()
                            if not p.stop_gradient).item()
+
         i = int(math.log(n_parameters, 10) // 3)
         size_unit = ['', 'K', 'M', 'B', 'T', 'Q']
         self.logger.info("Number of Parameters is {:.2f}{}.".format(
@@ -163,6 +164,7 @@ def __init__(self, cfg):
         else:
             self.lr_scheduler = build_lr_scheduler(cfg.lr_scheduler,
                                                    self.iters_per_epoch)
+
         self.optimizer = build_optimizer(cfg.optimizer, self.lr_scheduler,
                                          [self.model])
 
@@ -439,7 +441,6 @@ def load(self, weight_path, export=False):
             for k, v in state_dict.items():
                 state_dict_['model.' + k] = v
             state_dict = state_dict_
-
         self.model.set_state_dict(state_dict)
 
     def export(self, ckpt):

passl/hooks/optimizer_hook.py

@@ -14,6 +14,7 @@
 
 from .hook import Hook
 from .builder import HOOKS
+from ..solver.builder import build_optimizer
 
 
 @HOOKS.register()
@@ -33,15 +34,68 @@ def train_iter_end(self, trainer):
         if trainer.use_amp:
             scaled_loss = trainer.scaler.scale(loss)
             scaled_loss.backward()
-            trainer.scaler.step(trainer.optimizer)
-            trainer.scaler.update()
+            if 'lars' in trainer.optimizer.type:
+                trainer.scaler.minimize(trainer.optimizer, scaled_loss)
+            else:
+                trainer.scaler.step(trainer.optimizer)
+                trainer.scaler.update()
+        else:
+            loss.backward()
+            if 'lars' in trainer.optimizer.type:
+                trainer.optimizer.minimize(loss)
+            else:
+                trainer.optimizer.step()
 
+        if 'loss' not in trainer.outputs:
+            trainer.outputs['loss'] = loss
+
+
+@HOOKS.register()
+class SimsiamOptimizerHook(Hook):
+    def __init__(self, priority=1):
+        self.priority = priority
+
+    def run_begin(self, trainer):
+        if hasattr(trainer.model, '_layers'):
+            model = trainer.model._layers
+        else:
+            model = trainer.model
+
+        # build simsiam optimizer
+        trainer.optimizer = build_optimizer(
+            trainer.cfg.optimizer, trainer.lr_scheduler, [model.encoder])
+        trainer.predictor_optimizer = build_optimizer(
+            trainer.cfg.optimizer, trainer.lr_scheduler.get_lr(), [model.predictor])
+        
+    def train_iter_end(self, trainer):
+        if 'Lars' in trainer.cfg['optimizer']['name']:
+            trainer.optimizer.clear_gradients()
+            trainer.predictor_optimizer.clear_gradients()
+        else:
+            trainer.optimizer.clear_grad()
+            trainer.predictor_optimizer.clear_grad()
+
+        loss = 0
+        loss = trainer.outputs['loss']
+        
+        if trainer.use_amp:
+            scaled_loss = trainer.scaler.scale(loss)
+            scaled_loss.backward()
+            if 'lars' in trainer.optimizer.type:
+                trainer.scaler.minimize(trainer.optimizer, scaled_loss)
+                trainer.scaler.minimize(trainer.predictor_optimizer, scaled_loss)
+            else:
+                trainer.scaler.step(trainer.optimizer)
+                trainer.scaler.step(trainer.predictor_optimizer)
+                trainer.scaler.update()
         else:
             loss.backward()
             if 'lars' in trainer.optimizer.type:
                 trainer.optimizer.minimize(loss)
+                trainer.predictor_optimizer.minimize(loss)
             else:
                 trainer.optimizer.step()
+                trainer.predictor_optimizer.step()
 
         if 'loss' not in trainer.outputs:
             trainer.outputs['loss'] = loss

passl/modeling/architectures/__init__.py

@@ -17,6 +17,7 @@
 from .clas import Classification
 from .moco import MoCo
 from .simclr import SimCLR
+from .simsiam import SimSiam
 from .pixpro import PixPro
 
 from .BEiTWrapper import BEiTWrapper, BEiTPTWrapper, BEiTFTWrapper