Robust segvit compatibility + new dataset.

experimental/robust_segvit/README.md

@@ -1,16 +1,25 @@
 # Robust segvit
 
-*Robust_segvit* is a codebase to evaluate the robustness of semantic segmentation models.
+**Robust_segvit** is a codebase to evaluate the robustness of semantic segmentation models.  The code is built on top of [uncertainty_baselines](https://github.com/google/uncertainty-baselines) and [Scenic](https://github.com/google-research/scenic). 
 
-Robust_segvit is developed in [JAX](https://github.com/google/jax) and uses [Flax](https://github.com/google/flax), [uncertainty_baselines](https://github.com/google/uncertainty-baselines) and [Scenic](https://github.com/google-research/scenic).
+## Installation
+Robust_segvit is developed in [JAX](https://github.com/google/jax)/[Flax](https://github.com/google/flax).
 
-## Code structure
-See uncertainty_baselines/google/experimental/cityscapes.
+To run the code: <br>
+1. Install [uncertainty_baselines](https://github.com/google/uncertainty-baselines). <br>
+2. Install [Scenic](https://github.com/google-research/scenic). <br>
+3. Follow the instructions for a toy run in [./run_deterministic_mac.sh]().
 
+## Datasets
+The experiment configurations for the different datasets are in:
 
-## Cityscapes
+- configs/cityscapes: Cityscapes dataset. <br>
+- configs/ade20k_ind: ADE20k_ind dataset. <br>
+- configs/street_hazards: Street Hazards dataset. <br>
 
-We investigate the performance of different reliability methods on image segmentation tasks. <br>
+## Comments:
+- The checkpoint used for finetuning is the same the original segmenter model: [vit_large_patch16_384](https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py)
 
-[x] configs/cityscapes: contains experiment configurations for the cityscapes dataset. <br>
+## Citing work:
 
+If you reference this code, please cite [our paper](https://github.com/google/uncertainty-baselines). <br>

experimental/robust_segvit/configs/ade20k_ind/be.py

@@ -22,6 +22,8 @@
 # pylint: enable=line-too-long
 
 import ml_collections
+import os
+import datetime
 
 _CITYSCAPES_FINE_TRAIN_SIZE = 2975
 _CITYSCAPES_COARSE_TRAIN_SIZE = 19998
@@ -40,21 +42,24 @@
 
 # Model specs.
 LOAD_PRETRAINED_BACKBONE = True
-BACKBONE_ORIGIN = 'big_vision'
+BACKBONE_ORIGIN = 'vision_transformer'
 VIT_SIZE = 'L'
 STRIDE = 16
 RESNET_SIZE = None
 CLASSIFIER = 'token'
 target_size = (640, 640)
-UPSTREAM_TASK = 'i21k+imagenet2012'
+UPSTREAM_TASK = 'augreg+i21k+imagenet2012'
 
 
 # Upstream
 MODEL_PATHS = {
 
     # Imagenet 21k + finetune in imagenet2012 with perf 0.85 adap_res 384
-    ('big_vision', 'L', 16, None, 'token', 'i21k+imagenet2012'):
-        'gs://vit_models/imagenet21k%2Bimagenet2012/ViT-L_16.npz',
+    ('vision_transformer', 'L', 16, None, 'token', 'i21k+imagenet2012'):
+        'gs://vit_models/imagenet21k+imagenet2012/ViT-L_16.npz',
+    ('vision_transformer', 'L', 16, None, 'token', 'augreg+i21k+imagenet2012'):
+        'gs://vit_models/augreg/L_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_384.npz',  # pylint: disable=g-long-lambda
+
 }
 
 
@@ -180,9 +185,25 @@ def get_config(runlocal=''):
   config.eval_label_shift = False
   config.model.input_shape = target_size
 
+  config.eval_robustness_configs = ml_collections.ConfigDict()
+  config.eval_robustness_configs.auc_online = True
+  config.eval_robustness_configs.method_name = 'mlogit'
+
+  # wandb.ai configurations.
+  config.use_wandb = False
+  config.wandb_dir = 'wandb'
+  config.wandb_project = 'rdl-debug'
+  config.wandb_entity = 'ekellbuch'
+  config.wandb_exp_name = None  # Give experiment a name.
+  config.wandb_exp_name = (
+          os.path.splitext(os.path.basename(__file__))[0] + '_' +
+          datetime.datetime.today().strftime('%Y-%m-%d-%H-%M-%S'))
+  config.wandb_exp_group = None  # Give experiment a group name.
+
   if runlocal:
     config.count_flops = False
     config.dataset_configs.train_target_size = (128, 128)
+    config.model.input_shape = config.dataset_configs.train_target_size
     config.batch_size = 8
     config.num_training_epochs = 5
     config.warmup_steps = 0

experimental/robust_segvit/configs/ade20k_ind/be_eval.py

@@ -20,6 +20,8 @@
 # pylint: enable=line-too-long
 
 import ml_collections
+import datetime
+import os
 
 _CITYSCAPES_FINE_TRAIN_SIZE = 2975
 _CITYSCAPES_COARSE_TRAIN_SIZE = 19998
@@ -42,14 +44,14 @@
 STRIDE = 16
 RESNET_SIZE = None
 CLASSIFIER = 'token'
-EXPERIMENTID = '43838358-2'
+EXPERIMENTID = '45349725-1'
 
 target_size = (640, 640)
 
 # Upstream
 CHECKPOINT_PATHS = {
-    ('ub', 'L', 16, None, 'token', '43838358-2'):
-        'gs://ub-ekb/checkpoints_to_upload/ade20k/43838358-2',
+    ('ub', 'L', 16, None, 'token', '45349725-1'):
+        'gs://ub-checkpoints/45349725-ade20k_ind_segmenter_be/1',
 }
 
 
@@ -162,17 +164,33 @@ def get_config(runlocal=''):
   config.eval_mode = True
   config.eval_configs = ml_collections.ConfigDict()
   config.eval_configs.mode = 'standard'
-  config.eval_covariate_shift = True
-  config.eval_label_shift = True
   config.model.input_shape = target_size
+  config.eval_configs.store_logits = False
 
   # Eval parameters for robustness
   config.eval_label_shift = True
   config.eval_covariate_shift = True
   config.eval_robustness_configs = ml_collections.ConfigDict()
   config.eval_robustness_configs.auc_online = True
-  config.eval_robustness_configs.method_name = 'msp'
-  config.eval_robustness_configs.num_top_k = 5
+  config.eval_robustness_configs.method_name = 'nmlogit'
+  config.eval_robustness_configs.num_top_k = 1
+
+  # Load checkpoint
+  config.checkpoint_configs = ml_collections.ConfigDict()
+  config.checkpoint_configs.checkpoint_format = CHECKPOINT_ORIGIN
+  config.checkpoint_configs.checkpoint_path = CHECKPOINT_PATH
+  config.checkpoint_configs.classifier = 'token'
+
+  # wandb.ai configurations.
+  config.use_wandb = False
+  config.wandb_dir = 'wandb'
+  config.wandb_project = 'rdl-debug'
+  config.wandb_entity = 'ekellbuch'
+  config.wandb_exp_name = None  # Give experiment a name.
+  config.wandb_exp_name = (
+          os.path.splitext(os.path.basename(__file__))[0] + '_' +
+          datetime.datetime.today().strftime('%Y-%m-%d-%H-%M-%S'))
+  config.wandb_exp_group = None  # Give experiment a group name.
 
   if runlocal:
     config.count_flops = False
@@ -183,12 +201,6 @@ def get_config(runlocal=''):
     config.dataset_configs.train_split = f'train[:{TRAIN_SAMPLES}]'
     config.dataset_configs.validation_split = f'validation[:{TRAIN_SAMPLES}]'
     config.num_train_examples = TRAIN_SAMPLES
-  else:
-    # Load checkpoint
-    config.checkpoint_configs = ml_collections.ConfigDict()
-    config.checkpoint_configs.checkpoint_format = CHECKPOINT_ORIGIN
-    config.checkpoint_configs.checkpoint_path = CHECKPOINT_PATH
-    config.checkpoint_configs.classifier = 'token'
   return config
 
 

experimental/robust_segvit/configs/ade20k_ind/deterministic.py

@@ -22,6 +22,8 @@
 # pylint: enable=line-too-long
 
 import ml_collections
+import os
+import datetime
 
 _CITYSCAPES_FINE_TRAIN_SIZE = 2975
 _CITYSCAPES_COARSE_TRAIN_SIZE = 19998
@@ -40,21 +42,23 @@
 
 # Model specs.
 LOAD_PRETRAINED_BACKBONE = True
-BACKBONE_ORIGIN = 'big_vision'
+BACKBONE_ORIGIN = 'vision_transformer'
 VIT_SIZE = 'L'
 STRIDE = 16
 RESNET_SIZE = None
 CLASSIFIER = 'token'
 target_size = (640, 640)
-UPSTREAM_TASK = 'i21k+imagenet2012'
+UPSTREAM_TASK = 'augreg+i21k+imagenet2012'
 
 
 # Upstream
 MODEL_PATHS = {
 
     # Imagenet 21k + finetune in imagenet2012 with perf 0.85 adap_res 384
-    ('big_vision', 'L', 16, None, 'token', 'i21k+imagenet2012'):
-        'gs://vit_models/imagenet21k%2Bimagenet2012/ViT-L_16.npz',
+    ('vision_transformer', 'L', 16, None, 'token', 'i21k+imagenet2012'):
+        'gs://vit_models/imagenet21k+imagenet2012/ViT-L_16.npz',
+    ('vision_transformer', 'L', 16, None, 'token', 'augreg+i21k+imagenet2012'):
+        'gs://vit_models/augreg/L_16-i21k-300ep-lr_0.001-aug_medium1-wd_0.1-do_0.1-sd_0.1--imagenet2012-steps_20k-lr_0.01-res_384.npz',
 }
 
 
@@ -174,9 +178,25 @@ def get_config(runlocal=''):
   config.eval_label_shift = False
   config.model.input_shape = target_size
 
+  config.eval_robustness_configs = ml_collections.ConfigDict()
+  config.eval_robustness_configs.auc_online = True
+  config.eval_robustness_configs.method_name = 'mlogit'
+
+  # wandb.ai configurations.
+  config.use_wandb = False
+  config.wandb_dir = 'wandb'
+  config.wandb_project = 'rdl-debug'
+  config.wandb_entity = 'ekellbuch'
+  config.wandb_exp_name = None  # Give experiment a name.
+  config.wandb_exp_name = (
+          os.path.splitext(os.path.basename(__file__))[0] + '_' +
+          datetime.datetime.today().strftime('%Y-%m-%d-%H-%M-%S'))
+  config.wandb_exp_group = None  # Give experiment a group name.
+
   if runlocal:
     config.count_flops = False
     config.dataset_configs.train_target_size = (128, 128)
+    config.model.input_shape = config.dataset_configs.train_target_size
     config.batch_size = 8
     config.num_training_epochs = 5
     config.warmup_steps = 0