From de8cf677d9e59dbec8630fd9d6dc1dc386241307 Mon Sep 17 00:00:00 2001 From: D-X-Y <280835372@qq.com> Date: Mon, 17 May 2021 04:33:40 +0000 Subject: [PATCH] Update LFNA with resume --- exps/LFNA/lfna.py | 61 ++++++++++++++++++++++++++++-------- exps/LFNA/lfna_meta_model.py | 12 ++++--- lib/xlayers/super_module.py | 7 ++++- 3 files changed, 62 insertions(+), 18 deletions(-) diff --git a/exps/LFNA/lfna.py b/exps/LFNA/lfna.py index 6d498bc..f13c853 100644 --- a/exps/LFNA/lfna.py +++ b/exps/LFNA/lfna.py @@ -101,21 +101,49 @@ def main(args): ) lr_scheduler = torch.optim.lr_scheduler.MultiStepLR( optimizer, - milestones=[ - int(args.epochs * 0.8), - int(args.epochs * 0.9), - ], + milestones=[1, 2], gamma=0.1, ) logger.log("The base-model is\n{:}".format(base_model)) logger.log("The meta-model is\n{:}".format(meta_model)) logger.log("The optimizer is\n{:}".format(optimizer)) + logger.log("The scheduler is\n{:}".format(lr_scheduler)) logger.log("Per epoch iterations = {:}".format(len(env_loader))) - # LFNA meta-training + if logger.path("model").exists(): + ckp_data = torch.load(logger.path("model")) + base_model.load_state_dict(ckp_data["base_model"]) + meta_model.load_state_dict(ckp_data["meta_model"]) + optimizer.load_state_dict(ckp_data["optimizer"]) + lr_scheduler.load_state_dict(ckp_data["lr_scheduler"]) + last_success_epoch = ckp_data["last_success_epoch"] + start_epoch = ckp_data["iepoch"] + 1 + check_strs = [ + "epochs", + "env_version", + "hidden_dim", + "init_lr", + "layer_dim", + "time_dim", + "seq_length", + ] + for xstr in check_strs: + cx = getattr(args, xstr) + px = getattr(ckp_data["args"], xstr) + assert cx == px, "[{:}] {:} vs {:}".format(xstr, cx, ps) + success, _ = meta_model.save_best(ckp_data["cur_score"]) + logger.log("Load ckp from {:}".format(logger.path("model"))) + if success: + logger.log( + "Re-save the best model with score={:}".format(ckp_data["cur_score"]) + ) + else: + start_epoch, last_success_epoch = 0, 0 + + # LFNA meta-train + meta_model.set_best_dir(logger.path(None) / "checkpoint") per_epoch_time, start_time = AverageMeter(), time.time() - last_success_epoch = 0 - for iepoch in range(args.epochs): + for iepoch in range(start_epoch, args.epochs): head_str = "[{:}] [{:04d}/{:04d}] ".format( time_string(), iepoch, args.epochs @@ -132,11 +160,11 @@ def main(args): args.device, logger, ) - lr_scheduler.step() logger.log( head_str + " meta-loss: {meter.avg:.4f} ({meter.count:.0f})".format(meter=loss_meter) + " :: lr={:.5f}".format(min(lr_scheduler.get_last_lr())) + + " :: last-success={:}".format(last_success_epoch) ) success, best_score = meta_model.save_best(-loss_meter.avg) if success: @@ -145,8 +173,11 @@ def main(args): save_checkpoint( { "meta_model": meta_model.state_dict(), + "base_model": base_model.state_dict(), "optimizer": optimizer.state_dict(), "lr_scheduler": lr_scheduler.state_dict(), + "last_success_epoch": last_success_epoch, + "cur_score": -loss_meter.avg, "iepoch": iepoch, "args": args, }, @@ -154,8 +185,12 @@ def main(args): logger, ) if iepoch - last_success_epoch >= args.early_stop_thresh: - logger.log("Early stop at {:}".format(iepoch)) - break + if lr_scheduler.last_epoch > 2: + logger.log("Early stop at {:}".format(iepoch)) + break + else: + last_epoch.step() + logger.log("Decay the lr [{:}]".format(lr_scheduler.last_epoch)) per_epoch_time.update(time.time() - start_time) start_time = time.time() @@ -199,7 +234,7 @@ def main(args): [new_param], lr=args.init_lr, weight_decay=1e-5, amsgrad=True ) meta_model.replace_append_learnt( - torch.Tensor([future_time], device=args.device), new_param + torch.Tensor([future_time]).to(args.device), new_param ) meta_model.eval() base_model.train() @@ -289,8 +324,8 @@ if __name__ == "__main__": parser.add_argument( "--early_stop_thresh", type=int, - default=100, - help="The maximum epochs for early stop.", + default=50, + help="The #epochs for early stop.", ) parser.add_argument( "--seq_length", type=int, default=5, help="The sequence length." diff --git a/exps/LFNA/lfna_meta_model.py b/exps/LFNA/lfna_meta_model.py index c25e01e..79f6bd0 100644 --- a/exps/LFNA/lfna_meta_model.py +++ b/exps/LFNA/lfna_meta_model.py @@ -102,9 +102,11 @@ class LFNA_Meta(super_core.SuperModule): return torch.cat(meta_embed) def create_meta_embed(self): - param = torch.nn.Parameter(torch.Tensor(1, self._time_embed_dim)) + param = torch.Tensor(1, self._time_embed_dim) trunc_normal_(param, std=0.02) - return param.to(self._super_meta_embed.device) + param = param.to(self._super_meta_embed.device) + param = torch.nn.Parameter(param, True) + return param def get_closest_meta_distance(self, timestamp): with torch.no_grad(): @@ -112,12 +114,14 @@ class LFNA_Meta(super_core.SuperModule): return torch.min(distances).item() def replace_append_learnt(self, timestamp, meta_embed): - self._append_meta_embed["learnt"] = meta_embed self._append_meta_timestamps["learnt"] = timestamp + self._append_meta_embed["learnt"] = meta_embed def append_fixed(self, timestamp, meta_embed): with torch.no_grad(): - timestamp, meta_embed = timestamp.clone(), meta_embed.clone() + device = self._super_meta_embed.device + timestamp = timestamp.detach().clone().to(device) + meta_embed = meta_embed.detach().clone().to(device) if self._append_meta_timestamps["fixed"] is None: self._append_meta_timestamps["fixed"] = timestamp else: diff --git a/lib/xlayers/super_module.py b/lib/xlayers/super_module.py index 9beeb23..8007881 100644 --- a/lib/xlayers/super_module.py +++ b/lib/xlayers/super_module.py @@ -3,6 +3,7 @@ ##################################################### import os +from pathlib import Path import abc import tempfile import warnings @@ -90,6 +91,10 @@ class SuperModule(abc.ABC, nn.Module): total += buf.numel() return total + def set_best_dir(self, xdir): + self._meta_info[BEST_DIR_KEY] = str(xdir) + Path(xdir).mkdir(parents=True, exist_ok=True) + def save_best(self, score): if BEST_DIR_KEY not in self._meta_info: tempdir = tempfile.mkdtemp("-xlayers") @@ -97,7 +102,7 @@ class SuperModule(abc.ABC, nn.Module): if BEST_SCORE_KEY not in self._meta_info: self._meta_info[BEST_SCORE_KEY] = None best_score = self._meta_info[BEST_SCORE_KEY] - if best_score is None or best_score < score: + if best_score is None or best_score <= score: best_save_path = os.path.join( self._meta_info[BEST_DIR_KEY], "best-{:}.pth".format(self.__class__.__name__),