add HRDA

brdav · Sep 2, 2022 · 56a8625 · 56a8625
1 parent d90bd8e
commit 56a8625
Show file tree

Hide file tree

Showing 25 changed files with 906 additions and 77 deletions.
diff --git a/README.md b/README.md
@@ -15,6 +15,8 @@
 
 This repository provides the official code for the WACV 2023 paper [Refign: Align and Refine for Adaptation of Semantic Segmentation to Adverse Conditions](https://arxiv.org/abs/2207.06825). The code is organized using [PyTorch Lightning](https://github.com/Lightning-AI/lightning). 
 
+🔥 Applied on top of [HRDA](https://github.com/lhoyer/HRDA), Refign ranks #1 on both the [ACDC leaderboard](https://acdc.vision.ee.ethz.ch/benchmarks#semanticSegmentation)—**72.05 mIoU**—and the [Dark Zurich leaderboard](https://codalab.lisn.upsaclay.fr/competitions/3783#results)—**63.91 mIoU**. See below for training configurations.
+
 <img src="./docs/method.png" width="900"/>
 
 ## Abstract
@@ -26,7 +28,7 @@ Refign consists of two steps:
 (1) aligning the normal-condition image to the corresponding adverse-condition image using an uncertainty-aware dense matching network, and
 (2) refining the adverse prediction with the normal prediction using an adaptive label correction mechanism.
 We design custom modules to streamline both steps and set the new state of the art for domain-adaptive semantic segmentation on several adverse-condition benchmarks, including ACDC and Dark Zurich.
-The approach introduces no extra training parameters, minimal computational overhead---during training only---and can be used as a drop-in extension to improve any given self-training-based UDA method.
+The approach introduces no extra training parameters, minimal computational overhead—during training only—and can be used as a drop-in extension to improve any given self-training-based UDA method.
 
 
 ## Usage
@@ -217,26 +219,29 @@ Before running the code, download and extract the corresponding datasets to the
 </details>
 
 
-### Pretrained Models
+### Pretrained Models and Results
 
 We provide pretrained models of both UDA and alignment networks.
+Note that the UAWarpC checkpoint is needed to train Refign. To avoid config file edits, save it to `./pretrained_models/`.
+To facilitate qualitative segmentation comparisons, validation set predictions of Refign can be directly downloaded. Starred models use Cityscapes pretrained weights in the backbone, the others ImageNet pretrained.
 
 #### UDA
-- [Refign](https://data.vision.ee.ethz.ch/brdavid/refign/refign_acdc.ckpt) -- Cityscapes-->ACDC, 65.5 mIoU
-- [Refign](https://data.vision.ee.ethz.ch/brdavid/refign/refign_darkzurich.ckpt) -- Cityscapes-->DarkZurich, 56.2 mIoU
-- [Refign](https://data.vision.ee.ethz.ch/brdavid/refign/refign_robotcar.ckpt) -- Cityscapes-->RobotCar, 60.5 mIoU
-
-#### Alignment
-- [UAWarpC](https://data.vision.ee.ethz.ch/brdavid/refign/uawarpc_megadepth.ckpt)  -- MegaDepth
 
-Note that the UAWarpC checkpoint is needed to train Refign. To avoid config file edits, save it to `./pretrained_models/`.
+| Model         | Task           | Test Set       | Test Score    | Config      |  Checkpoint    |   Predictions  |
+|---------------|----------------|-----------------|-----------------|------------|----------------|------------|
+| Refign-DAFormer | Cityscapes→ACDC | ACDC test | 65.5 mIoU | [config](https://github.com/brdav/refign/blob/main/configs/cityscapes_acdc/refign_daformer.yaml) | [model](https://data.vision.ee.ethz.ch/brdavid/refign/refign_daformer_acdc.ckpt) | [ACDC val](https://data.vision.ee.ethz.ch/brdavid/refign/colored_preds_val_ACDC_refign_daformer.zip) 
+| Refign-HRDA* | Cityscapes→ACDC | ACDC test | 72.1 mIoU | [config](https://github.com/brdav/refign/blob/main/configs/cityscapes_acdc/refign_hrda_star.yaml) | [model](https://data.vision.ee.ethz.ch/brdavid/refign/refign_hrda_star_acdc.ckpt) | [ACDC val](https://data.vision.ee.ethz.ch/brdavid/refign/colored_preds_val_ACDC_refign_hrda_star.zip)
+|||||||
+| Refign-DAFormer | Cityscapes→Dark Zurich | Dark Zurich test | 56.2 mIoU | [config](https://github.com/brdav/refign/blob/main/configs/cityscapes_darkzurich/refign_daformer.yaml) | [model](https://data.vision.ee.ethz.ch/brdavid/refign/refign_daformer_darkzurich.ckpt) | [Dark Zurich val](https://data.vision.ee.ethz.ch/brdavid/refign/colored_preds_val_DarkZurich_refign_daformer.zip)
+| Refign-HRDA* | Cityscapes→Dark Zurich | Dark Zurich test | 63.9 mIoU | [config](https://github.com/brdav/refign/blob/main/configs/cityscapes_darkzurich/refign_hrda_star.yaml) | [model](https://data.vision.ee.ethz.ch/brdavid/refign/refign_hrda_star_darkzurich.ckpt) | [Dark Zurich val](https://data.vision.ee.ethz.ch/brdavid/refign/colored_preds_val_DarkZurich_refign_hrda_star.zip)
+|||||||
+| Refign-DAFormer | Cityscapes→RobotCar | RobotCar Seg. test       | 60.5 mIoU | [config](https://github.com/brdav/refign/blob/main/configs/cityscapes_robotcar/refign_daformer.yaml) | [model](https://data.vision.ee.ethz.ch/brdavid/refign/refign_daformer_robotcar.ckpt)  | [RobotCar val](https://data.vision.ee.ethz.ch/brdavid/refign/colored_preds_val_RobotCar_refign_daformer.zip)
 
-### Qualitative Refign Predictions
+#### Alignment
+| Model         | Task           | Test Set       | Score    | Config      |  Checkpoint    |
+|---------------|----------------|-----------------|-----------------|------------|----------------|
+| UAWarpC | MegaDepth Dense Matching | RobotCar Matching test | 36.8 PCK-5 | [stage1](https://github.com/brdav/refign/blob/main/configs/megadepth/uawarpc_stage1.yaml), [stage2](https://github.com/brdav/refign/blob/main/configs/megadepth/uawarpc_stage2.yaml) | [model](https://data.vision.ee.ethz.ch/brdavid/refign/uawarpc_megadepth.ckpt) | 
 
-To facilitate qualitative comparisons, validation set predictions of Refign can be directly downloaded:
-- [Refign on ACDC val](https://data.vision.ee.ethz.ch/brdavid/refign/colored_preds_val_ACDC.zip)
-- [Refign on Dark Zurich val](https://data.vision.ee.ethz.ch/brdavid/refign/colored_preds_val_DarkZurich.zip)
-- [Refign on RobotCar val](https://data.vision.ee.ethz.ch/brdavid/refign/colored_preds_val_RobotCar.zip)
 
 ### Refign Training
 
@@ -246,7 +251,7 @@ Enter the path to the UAWarpC model for `model.init_args.alignment_head.init_arg
 To train Refign on ACDC (single GPU, with AMP) use the following command:
 
 ```bash
-python tools/run.py fit --config configs/cityscapes_acdc/refign_daformer.yaml --trainer.gpus 1 --trainer.precision 16
+python tools/run.py fit --config configs/cityscapes_acdc/refign_hrda_star.yaml --trainer.gpus 1 --trainer.precision 16
 ```
 
 Similar config files are available for Dark Zurich and RobotCar.
@@ -258,15 +263,15 @@ As mentioned in the previous section, modify the config file by adding the UAWar
 To evaluate Refign e.g. on the ACDC validation set, use the following command:
 
 ```bash
-python tools/run.py test --config configs/cityscapes_acdc/refign_daformer.yaml --ckpt_path /path/to/trained/model --trainer.gpus 1
+python tools/run.py test --config configs/cityscapes_acdc/refign_hrda_star.yaml --ckpt_path /path/to/trained/model --trainer.gpus 1
 ```
 
 We also provide pretrained models, which can be downloaded from the link above. To evaluate them, simply provide them as the argument `--ckpt_path`.
 
 To get test set scores for ACDC and DarkZurich, predictions are evaluated on the respective evaluation servers: [ACDC](https://acdc.vision.ee.ethz.ch/submit) and [DarkZurich](https://codalab.lisn.upsaclay.fr/competitions/3783).
 To create and save test predictions for e.g. ACDC, use this command:
 ```bash
-python tools/run.py predict --config configs/cityscapes_acdc/refign_daformer.yaml --ckpt_path /path/to/trained/model --trainer.gpus 1
+python tools/run.py predict --config configs/cityscapes_acdc/refign_hrda_star.yaml --ckpt_path /path/to/trained/model --trainer.gpus 1
 ```
 
 ### UAWarpC Training

diff --git a/configs/cityscapes_acdc/ablations/refign_daformer_row1.yaml b/configs/cityscapes_acdc/ablations/refign_daformer_row1.yaml
@@ -87,8 +87,6 @@ model:
   class_path: models.DomainAdaptationSegmentationModel
   init_args:
     backbone_lr_factor: 0.1
-    psweight_ignore_top: 15
-    psweight_ignore_bottom: 15
     enable_fdist: True
     use_refign: True
     adapt_to_ref: False

diff --git a/configs/cityscapes_acdc/ablations/refign_daformer_row2.yaml b/configs/cityscapes_acdc/ablations/refign_daformer_row2.yaml
@@ -87,8 +87,6 @@ model:
   class_path: models.DomainAdaptationSegmentationModel
   init_args:
     backbone_lr_factor: 0.1
-    psweight_ignore_top: 15
-    psweight_ignore_bottom: 15
     enable_fdist: True
     use_refign: True
     adapt_to_ref: False

diff --git a/configs/cityscapes_acdc/ablations/refign_daformer_row3.yaml b/configs/cityscapes_acdc/ablations/refign_daformer_row3.yaml
@@ -87,8 +87,6 @@ model:
   class_path: models.DomainAdaptationSegmentationModel
   init_args:
     backbone_lr_factor: 0.1
-    psweight_ignore_top: 15
-    psweight_ignore_bottom: 15
     enable_fdist: True
     use_refign: True
     adapt_to_ref: False

diff --git a/configs/cityscapes_acdc/ablations/refign_daformer_row4.yaml b/configs/cityscapes_acdc/ablations/refign_daformer_row4.yaml
@@ -87,8 +87,6 @@ model:
   class_path: models.DomainAdaptationSegmentationModel
   init_args:
     backbone_lr_factor: 0.1
-    psweight_ignore_top: 15
-    psweight_ignore_bottom: 15
     enable_fdist: True
     use_refign: True
     adapt_to_ref: False

diff --git a/configs/cityscapes_acdc/ablations/refign_daformer_row5.yaml b/configs/cityscapes_acdc/ablations/refign_daformer_row5.yaml
@@ -87,8 +87,6 @@ model:
   class_path: models.DomainAdaptationSegmentationModel
   init_args:
     backbone_lr_factor: 0.1
-    psweight_ignore_top: 15
-    psweight_ignore_bottom: 15
     enable_fdist: True
     use_refign: True
     adapt_to_ref: False

diff --git a/configs/cityscapes_acdc/ablations/refign_daformer_row6.yaml b/configs/cityscapes_acdc/ablations/refign_daformer_row6.yaml
@@ -87,8 +87,6 @@ model:
   class_path: models.DomainAdaptationSegmentationModel
   init_args:
     backbone_lr_factor: 0.1
-    psweight_ignore_top: 15
-    psweight_ignore_bottom: 15
     enable_fdist: True
     use_refign: True
     adapt_to_ref: True

diff --git a/configs/cityscapes_acdc/refign_daformer.yaml b/configs/cityscapes_acdc/refign_daformer.yaml
@@ -87,8 +87,6 @@ model:
   class_path: models.DomainAdaptationSegmentationModel
   init_args:
     backbone_lr_factor: 0.1
-    psweight_ignore_top: 15
-    psweight_ignore_bottom: 15
     enable_fdist: True
     use_refign: True
     adapt_to_ref: True

diff --git a/...capes_acdc/refign_daformer_deeplabv2.yaml → ...igs/cityscapes_acdc/refign_deeplabv2.yaml b/...capes_acdc/refign_daformer_deeplabv2.yaml → ...igs/cityscapes_acdc/refign_deeplabv2.yaml
@@ -87,8 +87,6 @@ model:
   class_path: models.DomainAdaptationSegmentationModel
   init_args:
     backbone_lr_factor: 0.1
-    psweight_ignore_top: 15
-    psweight_ignore_bottom: 15
     enable_fdist: True
     use_refign: True
     adapt_to_ref: True
@@ -170,7 +168,7 @@ trainer:
     class_path: pytorch_lightning.loggers.TensorBoardLogger
     init_args:
       save_dir: lightning_logs
-      name: refign_daformer_deeplabv2_acdc
+      name: refign_deeplabv2_acdc
   callbacks:
     - class_path: pytorch_lightning.callbacks.LearningRateMonitor
     - class_path: helpers.callbacks.ValEveryNSteps

diff --git a/configs/cityscapes_acdc/refign_hrda_star.yaml b/configs/cityscapes_acdc/refign_hrda_star.yaml
@@ -0,0 +1,177 @@
+seed_everything: 0
+data:
+  class_path: data_modules.CombinedDataModule
+  init_args:
+    batch_size: 4
+    num_workers: 4
+    load_config:
+      train:
+        Cityscapes:
+          rcs_enabled: True
+          rcs_min_crop_ratio: 2.0
+          load_keys:
+            - image
+            - semantic
+          transforms:
+            - class_path: data_modules.transforms.ToTensor
+            - class_path: data_modules.transforms.RandomCrop
+              init_args:
+                size: 
+                  - 1024
+                  - 1024
+                cat_max_ratio: 0.75
+            - class_path: data_modules.transforms.RandomHorizontalFlip
+            - class_path: data_modules.transforms.ConvertImageDtype
+            - class_path: data_modules.transforms.Normalize
+        ACDC:
+          load_keys:
+            - image
+            - image_ref
+          transforms:
+            - class_path: data_modules.transforms.ToTensor
+            - class_path: data_modules.transforms.RandomCrop
+              init_args:
+                size: 
+                  - 1024
+                  - 1024
+            - class_path: data_modules.transforms.RandomHorizontalFlip
+            - class_path: data_modules.transforms.ConvertImageDtype
+            - class_path: data_modules.transforms.Normalize
+      val:
+        ACDC:
+          load_keys:
+            - image
+            - semantic
+          transforms:
+            - class_path: data_modules.transforms.ToTensor
+            - class_path: data_modules.transforms.ConvertImageDtype
+            - class_path: data_modules.transforms.Normalize
+      test:
+        ACDC:
+          load_keys:
+            - image
+            - semantic
+          transforms:
+            - class_path: data_modules.transforms.ToTensor
+            - class_path: data_modules.transforms.ConvertImageDtype
+            - class_path: data_modules.transforms.Normalize
+      predict:
+        ACDC:
+          predict_on: test
+          load_keys:
+            - image
+          transforms:
+            - class_path: data_modules.transforms.ToTensor
+            - class_path: data_modules.transforms.ConvertImageDtype
+            - class_path: data_modules.transforms.Normalize
+model:
+  class_path: models.DomainAdaptationSegmentationModel
+  init_args:
+    backbone_lr_factor: 0.1
+    enable_fdist: True
+    use_hrda: True
+    hrda_output_stride: 4
+    use_slide_inference: True
+    use_refign: True
+    adapt_to_ref: True
+    gamma: 0.25
+    backbone:
+      class_path: models.backbones.MixVisionTransformer
+      init_args:
+        model_type: mit_b5
+        pretrained: cityscapes
+    head:
+      class_path: models.heads.DAFormerHead
+      init_args:
+        in_channels:
+          - 64
+          - 128
+          - 320
+          - 512
+        in_index:
+          - 0
+          - 1
+          - 2
+          - 3
+        num_classes: 19
+        input_transform: multiple_select
+    hrda_scale_attention:
+      class_path: models.heads.SegFormerHead
+      init_args:
+        in_channels:
+          - 64
+          - 128
+          - 320
+          - 512
+        in_index:
+          - 0
+          - 1
+          - 2
+          - 3
+        num_classes: 19
+        input_transform: multiple_select
+    alignment_backbone:
+      class_path: models.backbones.VGG
+      init_args:
+        model_type: vgg16
+        pretrained: imagenet
+        out_indices:
+          - 2
+          - 3
+          - 4
+    alignment_head:
+      class_path: models.heads.UAWarpCHead
+      init_args:
+        in_index:
+          - 0
+          - 1
+        input_transform: multiple_select
+        estimate_uncertainty: True
+        pretrained: pretrained_models/uawarpc_megadepth.ckpt
+    loss:
+      class_path: models.losses.PixelWeightedCrossEntropyLoss
+    metrics:
+      val:
+        ACDC:
+          - class_path: helpers.metrics.IoU
+            init_args:
+              ignore_index: 255
+              num_classes: 19
+              compute_on_step: False
+      test:
+        ACDC:
+          - class_path: helpers.metrics.IoU
+            init_args:
+              ignore_index: 255
+              num_classes: 19
+              compute_on_step: False
+optimizer:
+  class_path: torch.optim.AdamW
+  init_args:
+    lr: 0.0006
+    weight_decay: 0.01
+lr_scheduler:
+  class_path: helpers.lr_scheduler.LinearWarmupPolynomialLR
+  init_args:
+    warmup_iters: 1500
+    warmup_ratio: 0.000001
+    power: 1.0
+    max_steps: 40000
+trainer:
+  max_steps: 40000
+  check_val_every_n_epoch: 40000
+  sync_batchnorm: True
+  multiple_trainloader_mode: min_size
+  logger:
+    - class_path: pytorch_lightning.loggers.TensorBoardLogger
+      init_args:
+        save_dir: lightning_logs
+        name: refign_hrda_daformer_cs_acdc
+  callbacks:
+    - class_path: pytorch_lightning.callbacks.LearningRateMonitor
+    - class_path: helpers.callbacks.ValEveryNSteps
+      init_args:
+        every_n_steps: 4000
+    - class_path: pytorch_lightning.callbacks.ModelCheckpoint
+      init_args:
+        save_last: True
diff --git a/configs/cityscapes_darkzurich/refign_daformer.yaml b/configs/cityscapes_darkzurich/refign_daformer.yaml
@@ -115,8 +115,6 @@ model:
   class_path: models.DomainAdaptationSegmentationModel
   init_args:
     backbone_lr_factor: 0.1
-    psweight_ignore_top: 15
-    psweight_ignore_bottom: 15
     enable_fdist: True
     use_refign: True
     adapt_to_ref: True

diff --git a/...darkzurich/refign_daformer_deeplabv2.yaml → ...tyscapes_darkzurich/refign_deeplabv2.yaml b/...darkzurich/refign_daformer_deeplabv2.yaml → ...tyscapes_darkzurich/refign_deeplabv2.yaml
@@ -115,8 +115,6 @@ model:
   class_path: models.DomainAdaptationSegmentationModel
   init_args:
     backbone_lr_factor: 0.1
-    psweight_ignore_top: 15
-    psweight_ignore_bottom: 15
     enable_fdist: True
     use_refign: True
     adapt_to_ref: True