add wandb

Author: dmMaze

.gitignore

@@ -1,10 +1,12 @@
 *.ckpt
 *.pt
 *.tar
+*.zip
 temp.ipynb
 
 checkpoints
 logs
+wandb
 VOC2012
 VOCdevkit
 __pycache__

config/config.py

@@ -14,7 +14,8 @@
 # DATA
 cfg.data = CN()
 cfg.data.type = 'voc2012_aug'
-cfg.data_root = './data'
+cfg.data.data_dir = './data'
+cfg.data.crop_size = 512
 cfg.data.num_classes = 21
 cfg.data.batch_size = 1
 cfg.data.num_workers = 0
@@ -50,5 +51,15 @@
 # LOGGING
 cfg.train.logger = CN()
 cfg.train.logger.log_dir = './logs'
-cfg.train.logger.tensorboard = True
+
+# tensorboard setting
+cfg.train.logger.use_tensorboard = True
+cfg.train.logger.tensorboard = CN()
+
+# wandb setting
+cfg.train.logger.use_wandb = False
+cfg.train.logger.wandb = CN()
+cfg.train.logger.wandb.project = 'UNet3Plus'
+cfg.train.logger.wandb.run_id = ''
+
 

config/test_voc.yaml

@@ -23,5 +23,5 @@ train:
   loss_type: u3p
 
   logger:
-    tensorboard: True
+    use_tensorboard: True
     log_dir: ./logs/

config/test_voc_cpu.yaml

@@ -0,0 +1,28 @@
+model:
+  encoder: resnet18
+  skip_ch: 16
+  aux_losses: -1
+  pretrained: True
+
+data:
+  type: voc2012_aug
+  num_classes: 21
+  num_workers: 2
+  batch_size: 1
+  max_training_samples: 10
+  
+train:
+  seed: 42
+  num_epochs: 20
+  lr: 0.001
+  weight_decay: 0.0001
+  optimizer: adamw
+  accum_steps: 2
+  resume: ''
+  device: cpu
+  loss_type: focal
+
+  logger:
+    use_tensorboard: False
+    use_wandb: False
+    log_dir: ./logs/

config/u3p_resnet18_voc.yaml

@@ -25,7 +25,7 @@ train:
   warmup_iters: 1
 
   logger:
-    tensorboard: True
+    use_tensorboard: True
     log_dir: ./logs
 
 

config/u3p_resnet34_voc.yaml

@@ -25,7 +25,7 @@ train:
   warmup_iters: 1000
 
   logger:
-    tensorboard: True
+    use_tensorboard: True
     log_dir: ./logs
 
 

datasets/__init__.py

@@ -40,8 +40,8 @@ def get_voc(data_root='./data', crop_size=SIZE, crop_val=SIZE, year='2012_aug',
 
     return train_dst, val_dst
 
-def build_data_loader(batch_size=1, num_workers=0, max_training_samples=-1) -> Tuple[DataLoader, DataLoader]:
-    train_dataset, val_dataset = get_voc()
+def build_data_loader(data_root='./data', batch_size=1, num_workers=0, max_training_samples=-1, crop_size=512) -> Tuple[DataLoader, DataLoader]:
+    train_dataset, val_dataset = get_voc(data_root, crop_size=crop_size, crop_val=crop_size)
     if max_training_samples > 0:    # for testing
         num_samples = len(train_dataset)
         train_dataset.image_set

model/unet3plus.py

@@ -26,6 +26,7 @@ def en2dec_layer(in_ch, out_ch, scale):
 def dec2dec_layer(in_ch, out_ch, scale, efficient=False):
     up = [nn.Upsample(scale_factor=scale, mode='bilinear', align_corners=True) if scale != 1 else nn.Identity()]
     m = [u3pblock(in_ch, out_ch, num_block=1)]
+    efficient = True
     if efficient:
         m = m + up
     else:

train.py

@@ -16,7 +16,7 @@
 from datasets import build_data_loader
 from config.config import cfg
 from utils.loss import build_u3p_loss
-from utils.log import AverageMeter
+from utils.logging import AverageMeter, SummaryLogger
 from utils.metrics import StreamSegMetrics
 
 def one_cycle(y1=0.0, y2=1.0, steps=100):
@@ -50,7 +50,7 @@ def __init__(self, cfg, model, train_loader, val_loader):
 
         # build loss
         self.criterion = build_u3p_loss(cfg.loss_type, cfg.aux_weight)
-        self.scaler = amp.GradScaler(enabled=True)  # mixed precision training
+        self.scaler = amp.GradScaler(enabled=cfg.device == 'cuda')  # mixed precision training
 
         # build optimizer
         if cfg.optimizer == 'sgd':
@@ -71,10 +71,7 @@ def __init__(self, cfg, model, train_loader, val_loader):
         # build scheduler
         self.scheduler = lr_scheduler.LambdaLR(self.optimizer, lr_lambda=self.lr_func)
 
-        if cfg.logger.tensorboard:
-            self.writer = SummaryWriter(log_dir=cfg.logger.log_dir)
-        else:
-            self.writer = None
+        self.logger = SummaryLogger(self.cfg_all)
 
         self.model.to(cfg.device)
         if cfg.resume:
@@ -174,22 +171,27 @@ def update_loss_dict(self, loss_dict, batch_loss_dict=None):
                 loss_dict[k].update(v)
 
     def log_results(self):
-        if self.writer is not None:
-            for k, v in self.loss_dict.items():
-                self.writer.add_scalars('Train_metrics/' + k, {"Train": v.avg}, self.global_iter)
-            self.update_loss_dict(self.loss_dict, None)     # clean loss meters
-            lr = self.optimizer.param_groups[0]['lr']
-            self.writer.add_scalars('Train_metrics/lr', {"lr": lr}, self.global_iter)
-
-            for k, v in self.val_loss_dict.items():
-                self.writer.add_scalars('Val_metrics/' + k, {"Val": v.avg}, self.global_iter)
-            self.update_loss_dict(self.val_loss_dict, None)
-            
-            for k, v in self.val_score_dict.items():
-                if k == 'Class IoU':
-                    continue
-                self.writer.add_scalars('Val_metrics/' + k, {"Val": v}, self.global_iter)
-            self.writer.flush()
+        log_dict = {
+            'Train': {},
+            'Val': {}
+        }
+
+        for k, v in self.loss_dict.items():
+            log_dict['Train'][k] = v.avg
+        self.update_loss_dict(self.loss_dict, None)
+        log_dict['Train']['lr'] = self.optimizer.param_groups[0]['lr']
+
+        for k, v in self.val_loss_dict.items():
+            log_dict['Val'][k] = v.avg
+        self.update_loss_dict(self.val_loss_dict, None)
+        
+        for k, v in self.val_score_dict.items():
+            if k == 'Class IoU':
+                print(v)
+                # self.logger.cmd_logger.info(v)
+                continue
+            log_dict['Val'][k] = v
+        self.logger.summary(log_dict, self.global_iter)
 
 
     def validate(self):
@@ -226,14 +228,15 @@ def main(args):
         cfg.train.seed = int(args.seed)
     if args.resume:
         cfg.train.resume = args.resume
+    cfg.data.data_dir = args.data_dir
 
     cfg.freeze()
     print(cfg)
     model, data = cfg.model, cfg.data
     model = build_unet3plus(data.num_classes, model.encoder, model.skip_ch, model.aux_losses, model.use_cgm, model.pretrained)
     # model = UNet_3Plus_DeepSup()
     if data.type in ['voc2012', 'voc2012_aug']:
-        train_loader, val_loader = build_data_loader(data.batch_size, data.num_workers, data.max_training_samples)
+        train_loader, val_loader = build_data_loader(data.data_dir, data.batch_size, data.num_workers, data.max_training_samples)
     else:
         raise NotImplementedError
 
@@ -254,6 +257,9 @@ def main(args):
                         help='resume from checkpoint',
                         default='',
                         type=str)
+    parser.add_argument('--data_dir',
+                        default="./data",
+                        type=str)
 
     args = parser.parse_args()
     main(args)