Refine Evaluation and Resume (#7)

* refine cfg and scheduler

* add auto resume

* add vl_pretrain_load

* rename test to eval

* remove old train and eval engine

* refine evaluation

* delete useless print

Author: rentainhe

config/common/train.py

@@ -4,35 +4,36 @@
 from simrec.scheduler.lr_scheduler import WarmupCosineLR
 
 train = dict(
+    output_dir = "./test",
+    warmup_epochs=3,
+    epochs = 25,
+    base_lr=1e-4,
+    warmup_lr=1e-7,
+    min_lr=1e-6,
     batch_size=8,
-    num_workers=8,
+    log_period=1,
+    data=dict(pin_memory=True, num_workers=8),
+    scheduler=dict(
+        name="cosine",
+        decay_epochs=[30, 35, 37],
+        lr_decay_rate=0.2,
+    ),
     amp=dict(enabled=False),
     ddp=dict(
         backend="nccl",
         init_method="env://",
     ),
     ema=dict(enabled=True, alpha=0.9997, buffer_ema=True),
-    epochs = 25,
-    output_dir = "./test",
-    log_period = 1,
-    resume=dict(enable=False, auto_resume=True, resume_path=""),
+    clip_grad_norm=0.15,
+    auto_resume=dict(enabled=True),
+    resume_path="",
     vl_pretrain_weight="",
-
-    scheduler = LazyCall(WarmupCosineLR)(
-        # optimizer and epochs and n_iter_per_epoch will be set in train.py
-        warmup_epochs = 3,
-        warmup_lr = 0.0000001,
-        base_lr = 0.0001,
-        min_lr = 0.000001,
-    ),
-
     multi_scale_training=dict(
         enabled=True,
         img_scales=[[224,224],[256,256],[288,288],[320,320],[352,352],
                     [384,384],[416,416],[448,448],[480,480],[512,512],
                     [544,544],[576,576],[608,608]]
     ),
-    clip_grad_norm=0.15,
-    log_image = False,
+    log_image=False,
     seed = 123456,
 )

config/simrec_refcoco_scratch.py

@@ -7,4 +7,6 @@
 
 dataset.ann_path["refcoco"] = "/home/rentianhe/dataset/rec/anns/refcoco.json"
 dataset.image_path["refcoco"] = "/home/rentianhe/dataset/rec/images/train2014"
-dataset.mask_path["refcoco"] = "/home/rentianhe/dataset/rec/masks/refcoco"
+dataset.mask_path["refcoco"] = "/home/rentianhe/dataset/rec/masks/refcoco"
+
+train.resume_path = "/home/rentianhe/code/SimREC/output/ckpt_epoch_3.pth"

simrec/datasets/dataloader.py

@@ -31,8 +31,8 @@ def build_loader(cfg, dataset: torch.utils.data.Dataset, rank: int, shuffle=True
                             dataset,
                             batch_size=cfg.train.batch_size,
                             sampler=dist_sampler,
-                            num_workers=cfg.train.num_workers,
-                            pin_memory=True,
+                            num_workers=cfg.train.data.num_workers,
+                            pin_memory=cfg.train.data.pin_memory,
                             drop_last=drop_last
                         )
     return data_loader

simrec/models/heads/rec_heads.py

@@ -33,7 +33,6 @@ def __init__(
         super().__init__()
         # same padding
         pad = (ksize - 1) // 2
-        print(in_channels,out_channels,ksize,stride)
         self.conv = nn.Conv2d(
             in_channels,
             out_channels,

simrec/models/mcn.py

@@ -18,7 +18,7 @@
 from simrec.models.heads.mcn_heads import MCNhead
 from simrec.models.backbones.build import build_visual_encoder
 from simrec.models.language_encoders.build import build_language_encoder
-from simrec.layers.fusion_layer import SimpleFusion,MultiScaleFusion,GaranAttention
+from simrec.layers.fusion_layer import SimpleFusion, MultiScaleFusion, GaranAttention
 
 
 class MCN(nn.Module):

simrec/models/simrec.py

@@ -36,23 +36,11 @@ def __init__(
         super(SimREC, self).__init__()
         self.visual_encoder=visual_backbone
         self.lang_encoder=language_encoder
-        # self.multi_scale_manner = MultiScaleFusion(v_planes=(512, 512, hidden_size), scaled=True)
         self.multi_scale_manner = multi_scale_manner
-        # self.fusion_manner=nn.ModuleList(
-        #     [
-        #         SimpleFusion(v_planes=256, out_planes=512, q_planes=512),
-        #         SimpleFusion(v_planes=512, out_planes=512, q_planes=512),
-        #         SimpleFusion(v_planes=1024, out_planes=512, q_planes=512)
-        #     ]
-        # )
         self.fusion_manner = fusion_manner
-        # self.attention_manner=GaranAttention(512,512)
         self.attention_manner = attention_manner
         self.head=head
 
-        total = sum([param.nelement() for param in self.lang_encoder.parameters()])
-        print('  + Number of lang enc params: %.2fM' % (total / 1e6))
-
 
     def frozen(self,module):
         if getattr(module,'module',False):

simrec/scheduler/build.py

@@ -13,10 +13,31 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from simrec.config import instantiate
+from .lr_scheduler import StepLR, WarmupCosineLR
 
-def build_lr_scheduler(cfg, optimizer):
-    """Build learning rate scheduler, defined by ``cfg.train.scheduler``."""
-    cfg.optimizer = optimizer
-    scheduler = instantiate(cfg)
-    return scheduler
+def build_lr_scheduler(cfg, optimizer, n_iter_per_epoch):
+    """Build learning rate scheduler."""
+    scheduler_name = cfg.train.scheduler.name.lower()
+    
+    scheduler = None
+    if scheduler_name == "cosine":
+        scheduler = WarmupCosineLR(
+            optimizer=optimizer,
+            warmup_epochs=cfg.train.warmup_epochs,
+            epochs=cfg.train.epochs,
+            warmup_lr=cfg.train.warmup_lr,
+            base_lr=cfg.train.base_lr,
+            min_lr=cfg.train.min_lr,
+            n_iter_per_epoch=n_iter_per_epoch
+        )
+    elif scheduler_name == "step":
+        scheduler = StepLR(
+            optimizer=optimizer,
+            warmup_epochs=cfg.train.warmup_epochs,
+            epochs=cfg.train.epochs,
+            decay_epochs=cfg.train.scheduler.decay_epochs,
+            lr_decay_rate=cfg.train.lr_decay_rate,
+            n_iter_per_epoch=n_iter_per_epoch,
+        )
+    
+    return scheduler

simrec/utils/checkpoint.py

@@ -14,64 +14,23 @@
 # limitations under the License.
 
 import os
-import warnings
 
 import torch
-from torch.nn import DataParallel as DP
-from torch.nn.parallel import DistributedDataParallel as DDP
 
+from simrec.utils.distributed import is_main_process
 
-def save_ckpt(net, optimizer,scheduler, misc, __C):
-    path = __C.CKPTs_PATH
-    if not os.path.exists(path):
-        os.mkdir(path)
-    path += '/' + __C.VERSION
-    if not os.path.exists(path):
-        os.mkdir(path)
-    assert isinstance(misc, dict)
-    if isinstance(net, DP) or isinstance(net, DDP):
-        path += '/' + 'dist_'
-    path += str(misc['epoch']) + '.pth.tar'
-    ckpt = {
-        'net_state_dict': net.state_dict(),
-        'optimizer_state_dict': optimizer.state_dict(),
-        'scheduler':scheduler.state_dict(),
-        'epoch':misc['epoch'],
-        'lr':optimizer.param_groups[0]["lr"],
-    }
-    torch.save(ckpt, path)
-
-
-def load_ckpt(net, optimizer,scheduler, path, rank=None):
-    loc = f'cuda:{rank}' if rank is not None else None
-    ckpt = torch.load(path, map_location=loc)
-
-    flag = isinstance(net, DP) or isinstance(net, DDP)
-    if '_dist' in path:
-        if not flag:
-            for name in ckpt['net_state_dict']:
-                assert name.startswith('module.')
-                ckpt['net_state_dict'][name.lstrip('module.')] = ckpt['net_state_dict'].pop(name)
-    else:
-        if flag:
-            for name in ckpt['net_state_dict']:
-                ckpt['net_state_dict']['module.' + name] = ckpt['net_state_dict'].pop(name)
-
-    optimizer.load_state_dict(ckpt['optimizer_state_dict'])
 
-    scheduler.load_state_dict(ckpt['scheduler'])
-
-    missing, unexpected = net.load_state_dict(ckpt['net_state_dict'], strict=False)
-    if unexpected.__len__ != 0:
-        warnings.warn(f'Current model misses {unexpected.__len__} parameters from checkpointing model')
-        for name in missing:
-            print('\n' + name + '\n')
-    if missing.__len__ != 0:
-        warnings.warn(f'Current model contains {missing.__len__} parameters that checkpointing model doesn\'t contain')
-        for name in unexpected:
-            print('\n' + name + '\n')
-
-    return ckpt
+def load_checkpoint(cfg, model, optimizer, scheduler, logger):
+    logger.info(f"==============> Resuming form {cfg.train.resume_path}....................")
+    checkpoint = torch.load(cfg.train.resume_path, map_location=lambda storage, loc: storage.cuda())
+    msg = model.load_state_dict(checkpoint['state_dict'], strict=False)
+    logger.info(msg)
+    optimizer.load_state_dict(checkpoint["optimizer"])
+    scheduler.load_state_dict(checkpoint["scheduler"])
+    start_epoch = checkpoint["epoch"]
+    logger.info("==> loaded checkpoint from {}\n".format(cfg.train.resume_path) +
+                "==> epoch: {} lr: {} ".format(checkpoint['epoch'],checkpoint['lr']))
+    return start_epoch + 1
 
 
 def save_checkpoint(cfg, epoch, model, optimizer, scheduler, logger, det_best=False, seg_best=False):
@@ -99,4 +58,16 @@ def save_checkpoint(cfg, epoch, model, optimizer, scheduler, logger, det_best=Fa
     if seg_best:
         seg_best_model_path = os.path.join(cfg.train.output_dir, f'seg_best_model.pth')
         torch.save(save_state, seg_best_model_path)
-    logger.info(f"checkpoints saved !!!")
+    logger.info(f"checkpoints saved !!!")
+
+
+def auto_resume_helper(output_dir):
+    checkpoints = os.listdir(output_dir)
+    checkpoints = [ckpt for ckpt in checkpoints if ckpt.endswith('pth')]
+    print(f"All checkpoints founded in {output_dir}: {checkpoints}")
+    if len(checkpoints) > 0:
+        resume_file = os.path.join(output_dir, "last_checkpoint.pth")
+    else:
+        resume_file = None
+
+    return resume_file

simrec/utils/distributed.py

@@ -75,11 +75,17 @@ def synchronize():
         dist.barrier()
 
 
-
 def cleanup_distributed():
     dist.destroy_process_group()
 
 
+def reduce_tensor(tensor):
+    rt = tensor.clone()
+    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
+    rt /= dist.get_world_size()
+    return rt
+
+
 def reduce_meters(meters, rank, cfg):
     """Sync and flush meters."""
     assert isinstance(meters, dict), "collect AverageMeters into a dict"
@@ -94,6 +100,7 @@ def reduce_meters(meters, rank, cfg):
             value = torch.mean(torch.cat(avg_reduce)).item()
             meter.update_reduce(value)
 
+
 def find_free_port():
     import socket
     sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)