Add files for reproducibility on xBD data

xBD_code/adamw.py

@@ -0,0 +1,89 @@
+# Based on https://github.com/pytorch/pytorch/pull/3740
+import torch
+import math
+
+
+class AdamW(torch.optim.Optimizer):
+    """Implements AdamW algorithm.
+
+    It has been proposed in `Fixing Weight Decay Regularization in Adam`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+
+    .. Fixing Weight Decay Regularization in Adam:
+    https://arxiv.org/abs/1711.05101
+    """
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
+                 weight_decay=0):
+        defaults = dict(lr=lr, betas=betas, eps=eps,
+                        weight_decay=weight_decay)
+        super(AdamW, self).__init__(params, defaults)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('AdamW does not support sparse gradients, please consider SparseAdam instead')
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+
+                # according to the paper, this penalty should come after the bias correction
+                # if group['weight_decay'] != 0:
+                #     grad = grad.add(group['weight_decay'], p.data)
+
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(1 - beta1, grad)
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+
+                denom = exp_avg_sq.sqrt().add_(group['eps'])
+
+                bias_correction1 = 1 - beta1 ** state['step']
+                bias_correction2 = 1 - beta2 ** state['step']
+                step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1
+
+                # w = w - wd * lr * w
+                if group['weight_decay'] != 0:
+                    p.data.add_(-group['weight_decay'] * group['lr'], p.data)
+
+                # w = w - lr * w.grad
+                p.data.addcdiv_(-step_size, exp_avg, denom)
+
+                # w = w - wd * lr * w - lr * w.grad
+                # See http://www.fast.ai/2018/07/02/adam-weight-decay/
+
+        return loss

xBD_code/dual_hrnet_config.yaml

@@ -0,0 +1,118 @@
+OUTPUT_DIR: ''
+LOG_DIR: ''
+GPUS: [0,]
+WORKERS: 4
+PRINT_FREQ: 20
+AUTO_RESUME: False
+PIN_MEMORY: True
+RANK: 0
+
+# Cudnn related params
+CUDNN:
+  BENCHMARK: True
+  DETERMINISTIC: False
+  ENABLED: True
+
+# common params for NETWORK
+MODEL:
+  NAME: 'dual-hrnet'
+  PRETRAINED: './Checkpoints/HRNet/hrnetv2_w32_imagenet_pretrained.pth'
+  USE_FPN: False
+  IS_DISASTER_PRED: False
+  IS_SPLIT_LOSS: True
+  FUSE_CONV_K_SIZE: 1
+
+  # high_resoluton_net related params for segmentation
+  EXTRA:
+    PRETRAINED_LAYERS: ['*']
+    STEM_INPLANES: 64
+    FINAL_CONV_KERNEL: 1
+    WITH_HEAD: True
+
+    STAGE1:
+      NUM_MODULES: 1
+      NUM_BRANCHES: 1
+      NUM_BLOCKS: [4]
+      NUM_CHANNELS: [64]
+      BLOCK: 'BOTTLENECK'
+      FUSE_METHOD: 'SUM'
+
+    STAGE2:
+      NUM_MODULES: 1
+      NUM_BRANCHES: 2
+      NUM_BLOCKS: [4, 4]
+      NUM_CHANNELS: [32, 64]
+      BLOCK: 'BASIC'
+      FUSE_METHOD: 'SUM'
+
+    STAGE3:
+      NUM_MODULES: 4
+      NUM_BRANCHES: 3
+      NUM_BLOCKS: [4, 4, 4]
+      NUM_CHANNELS: [32, 64, 128]
+      BLOCK: 'BASIC'
+      FUSE_METHOD: 'SUM'
+
+    STAGE4:
+      NUM_MODULES: 3
+      NUM_BRANCHES: 4
+      NUM_BLOCKS: [4, 4, 4, 4]
+      NUM_CHANNELS: [32, 64, 128, 256]
+      BLOCK: 'BASIC'
+      FUSE_METHOD: 'SUM'
+
+#_C.MODEL.EXTRA= CN(new_allowed=True)
+
+LOSS:
+  CLASS_BALANCE: True
+
+# DATASET related params
+DATASET:
+  NUM_CLASSES: 4
+
+# training
+TRAIN:
+  # Augmentation
+  FLIP: True
+  MULTI_SCALE: [0.8, 1.2]
+  CROP_SIZE: [512, 512]
+
+  LR_FACTOR: 0.1
+  LR_STEP: [90, 110]
+  LR: 0.05
+  EXTRA_LR: 0.001
+
+  OPTIMIZER: 'sgd'
+  MOMENTUM: 0.9
+  WD: 0.0001
+  NESTEROV: False
+  IGNORE_LABEL: -1
+
+  NUM_EPOCHS: 500
+  RESUME: False
+
+  BATCH_SIZE_PER_GPU: 16
+  SHUFFLE: True
+# only using some training samples
+  NUM_SAMPLES: 0
+  CLASS_WEIGHTS: [0.4, 1.2, 1.2, 1.2]
+
+# testing
+TEST:
+  BATCH_SIZE_PER_GPU: 32
+# only testing some samples
+  NUM_SAMPLES: 0
+
+  MODEL_FILE: ''
+  FLIP_TEST: False
+  MULTI_SCALE: False
+  CENTER_CROP_TEST: False
+  SCALE_LIST: [1]
+
+# debug
+DEBUG:
+  DEBUG: False
+  SAVE_BATCH_IMAGES_GT: False
+  SAVE_BATCH_IMAGES_PRED: False
+  SAVE_HEATMAPS_GT: False
+  SAVE_HEATMAPS_PRED: False