Update LFNA with train/valid

exps/LFNA/lfna.py

@@ -2,7 +2,8 @@
 # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.04 #
 #####################################################
 # python exps/LFNA/lfna.py --env_version v1 --workers 0
-# python exps/LFNA/lfna.py --env_version v1 --device cuda
+# python exps/LFNA/lfna.py --env_version v1 --device cuda --lr 0.001
+# python exps/LFNA/lfna.py --env_version v1 --device cuda --lr 0.002
 #####################################################
 import sys, time, copy, torch, random, argparse
 from tqdm import tqdm
@@ -58,9 +59,40 @@ def epoch_train(loader, meta_model, base_model, optimizer, criterion, device, lo
     return loss_meter
 
 
+def epoch_evaluate(loader, meta_model, base_model, criterion, device, logger):
+    with torch.no_grad():
+        base_model.eval()
+        meta_model.eval()
+        loss_meter = AverageMeter()
+        for ibatch, batch_data in enumerate(loader):
+            timestamps, (batch_seq_inputs, batch_seq_targets) = batch_data
+            timestamps = timestamps.squeeze(dim=-1).to(device)
+            batch_seq_inputs = batch_seq_inputs.to(device)
+            batch_seq_targets = batch_seq_targets.to(device)
+
+            batch_seq_containers = meta_model(timestamps)
+            losses = []
+            for seq_containers, seq_inputs, seq_targets in zip(
+                batch_seq_containers, batch_seq_inputs, batch_seq_targets
+            ):
+                for container, inputs, targets in zip(
+                    seq_containers, seq_inputs, seq_targets
+                ):
+                    predictions = base_model.forward_with_container(inputs, container)
+                    loss = criterion(predictions, targets)
+                    losses.append(loss)
+            final_loss = torch.stack(losses).mean()
+            loss_meter.update(final_loss.item())
+    return loss_meter
+
+
 def main(args):
     logger, env_info, model_kwargs = lfna_setup(args)
-    dynamic_env = get_synthetic_env(mode="train", version=args.env_version)
+    train_env = get_synthetic_env(mode="train", version=args.env_version)
+    valid_env = get_synthetic_env(mode="valid", version=args.env_version)
+    logger.log("training enviornment: {:}".format(train_env))
+    logger.log("validation enviornment: {:}".format(valid_env))
+
     base_model = get_model(**model_kwargs)
     base_model = base_model.to(args.device)
     criterion = torch.nn.MSELoss()
@@ -68,26 +100,25 @@ def main(args):
     shape_container = base_model.get_w_container().to_shape_container()
 
     # pre-train the hypernetwork
-    timestamps = dynamic_env.get_timestamp(None)
+    timestamps = train_env.get_timestamp(None)
     meta_model = LFNA_Meta(shape_container, args.layer_dim, args.time_dim, timestamps)
     meta_model = meta_model.to(args.device)
 
     logger.log("The base-model has {:} weights.".format(base_model.numel()))
     logger.log("The meta-model has {:} weights.".format(meta_model.numel()))
 
-    batch_sampler = EnvSampler(dynamic_env, args.meta_batch, args.sampler_enlarge)
+    batch_sampler = EnvSampler(train_env, args.meta_batch, args.sampler_enlarge)
-    dynamic_env.reset_max_seq_length(args.seq_length)
+    train_env.reset_max_seq_length(args.seq_length)
-    """
+    valid_env.reset_max_seq_length(args.seq_length)
-    env_loader = torch.utils.data.DataLoader(
+    valid_env_loader = torch.utils.data.DataLoader(
-        dynamic_env,
+        valid_env,
         batch_size=args.meta_batch,
         shuffle=True,
         num_workers=args.workers,
         pin_memory=True,
     )
-    """
+    train_env_loader = torch.utils.data.DataLoader(
-    env_loader = torch.utils.data.DataLoader(
+        train_env,
-        dynamic_env,
         batch_sampler=batch_sampler,
         num_workers=args.workers,
         pin_memory=True,
@@ -95,7 +126,7 @@ def main(args):
 
     optimizer = torch.optim.Adam(
         meta_model.parameters(),
-        lr=args.init_lr,
+        lr=args.lr,
         weight_decay=args.weight_decay,
         amsgrad=True,
     )
@@ -108,7 +139,7 @@ def main(args):
     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)))
+    logger.log("Per epoch iterations = {:}".format(len(train_env_loader)))
 
     if logger.path("model").exists():
         ckp_data = torch.load(logger.path("model"))
@@ -122,7 +153,7 @@ def main(args):
             "epochs",
             "env_version",
             "hidden_dim",
-            "init_lr",
+            "lr",
             "layer_dim",
             "time_dim",
             "seq_length",
@@ -152,7 +183,7 @@ def main(args):
         )
 
         loss_meter = epoch_train(
-            env_loader,
+            train_env_loader,
             meta_model,
             base_model,
             optimizer,
@@ -160,9 +191,16 @@ def main(args):
             args.device,
             logger,
         )
+
+        valid_loss_meter = epoch_evaluate(
+            valid_env_loader, meta_model, base_model, criterion, args.device, logger
+        )
         logger.log(
             head_str
-            + " meta-loss: {meter.avg:.4f} ({meter.count:.0f})".format(meter=loss_meter)
+            + " meta-train-loss: {meter.avg:.4f} ({meter.count:.0f})".format(
+                meter=loss_meter
+            )
+            + " meta-valid-loss: {meter.val:.4f}".format(meter=valid_loss_meter)
             + " :: lr={:.5f}".format(min(lr_scheduler.get_last_lr()))
             + "  :: last-success={:}".format(last_success_epoch)
         )
@@ -231,14 +269,14 @@ def main(args):
         #
         new_param = meta_model.create_meta_embed()
         optimizer = torch.optim.Adam(
-            [new_param], lr=args.init_lr, weight_decay=1e-5, amsgrad=True
+            [new_param], lr=args.refine_lr, weight_decay=1e-5, amsgrad=True
         )
         meta_model.replace_append_learnt(
             torch.Tensor([future_time]).to(args.device), new_param
         )
         meta_model.eval()
         base_model.train()
-        for iepoch in range(args.epochs):
+        for iepoch in range(args.refine_epochs):
             optimizer.zero_grad()
             [seq_containers] = meta_model(time_seqs)
             future_container = seq_containers[-1]
@@ -297,7 +335,7 @@ if __name__ == "__main__":
     )
     #####
     parser.add_argument(
-        "--init_lr",
+        "--lr",
         type=float,
         default=0.005,
         help="The initial learning rate for the optimizer (default is Adam)",
@@ -321,10 +359,19 @@ if __name__ == "__main__":
         help="Enlarge the #iterations for an epoch",
     )
     parser.add_argument("--epochs", type=int, default=10000, help="The total #epochs.")
+    parser.add_argument(
+        "--refine_lr",
+        type=float,
+        default=0.005,
+        help="The learning rate for the optimizer, during refine",
+    )
+    parser.add_argument(
+        "--refine_epochs", type=int, default=1000, help="The final refine #epochs."
+    )
     parser.add_argument(
         "--early_stop_thresh",
         type=int,
-        default=50,
+        default=20,
         help="The #epochs for early stop.",
     )
     parser.add_argument(
@@ -350,7 +397,7 @@ if __name__ == "__main__":
         args.hidden_dim,
         args.layer_dim,
         args.time_dim,
-        args.init_lr,
+        args.lr,
         args.weight_decay,
         args.epochs,
         args.env_version,

exps/LFNA/lfna_meta_model.py

@@ -44,6 +44,7 @@ class LFNA_Meta(super_core.SuperModule):
         self._append_meta_embed = dict(fixed=None, learnt=None)
         self._append_meta_timestamps = dict(fixed=None, learnt=None)
 
+        self._time_prob_drop = super_core.SuperDrop(dropout, (-1, 1), recover=False)
         # build transformer
         layers = []
         for ilayer in range(mha_depth):
@@ -149,10 +150,12 @@ class LFNA_Meta(super_core.SuperModule):
         meta_match = meta_match.view(batch, seq, -1)
         # create the probability
         time_probs = (1 / torch.exp(time_match_v * 10)).view(batch, seq, 1)
-        if self.training:
+
-            time_probs[:, -1, :] = 0
+        x_time_probs = self._time_prob_drop(time_probs)
+        # if self.training:
+        #    time_probs[:, -1, :] = 0
         unknown_token = self._unknown_token.view(1, 1, -1)
-        raw_meta_embed = time_probs * meta_match + (1 - time_probs) * unknown_token
+        raw_meta_embed = x_time_probs * meta_match + (1 - x_time_probs) * unknown_token
 
         meta_embed = self.meta_corrector(raw_meta_embed)
         # create joint embed

lib/datasets/synthetic_env.py

@@ -151,12 +151,15 @@ class SyntheticDEnv(data.Dataset):
         return len(self._timestamp_generator)
 
     def __repr__(self):
-        return "{name}({cur_num:}/{total} elements, ndim={ndim}, num_per_task={num_per_task})".format(
+        return "{name}({cur_num:}/{total} elements, ndim={ndim}, num_per_task={num_per_task}, range=[{xrange_min:.5f}~{xrange_max:.5f}], mode={mode})".format(
             name=self.__class__.__name__,
             cur_num=len(self),
             total=len(self._timestamp_generator),
             ndim=self._ndim,
             num_per_task=self._num_per_task,
+            xrange_min=self.min_timestamp,
+            xrange_max=self.max_timestamp,
+            mode=self._timestamp_generator.mode,
         )
 
 

lib/xlayers/super_core.py

@@ -15,6 +15,7 @@ from .super_norm import SuperLayerNorm1D
 from .super_norm import SuperSimpleLearnableNorm
 from .super_norm import SuperIdentity
 from .super_dropout import SuperDropout
+from .super_dropout import SuperDrop
 
 super_name2norm = {
     "simple_norm": SuperSimpleNorm,

lib/xlayers/super_dropout.py

@@ -6,7 +6,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 import math
-from typing import Optional, Callable
+from typing import Optional, Callable, Tuple
 
 import spaces
 from .super_module import SuperModule
@@ -38,3 +38,46 @@ class SuperDropout(SuperModule):
     def extra_repr(self) -> str:
         xstr = "inplace=True" if self._inplace else ""
         return "p={:}".format(self._p) + ", " + xstr
+
+
+class SuperDrop(SuperModule):
+    """Applies a the drop-path function element-wise."""
+
+    def __init__(self, p: float, dims: Tuple[int], recover: bool = True) -> None:
+        super(SuperDrop, self).__init__()
+        self._p = p
+        self._dims = dims
+        self._recover = recover
+
+    @property
+    def abstract_search_space(self):
+        return spaces.VirtualNode(id(self))
+
+    def forward_candidate(self, input: torch.Tensor) -> torch.Tensor:
+        return self.forward_raw(input)
+
+    def forward_raw(self, input: torch.Tensor) -> torch.Tensor:
+        if not self.training or self._p <= 0:
+            return input
+        keep_prob = 1 - self._p
+        shape = [input.shape[0]] + [
+            x if y == -1 else y for x, y in zip(input.shape[1:], self._dims)
+        ]
+        random_tensor = keep_prob + torch.rand(
+            shape, dtype=input.dtype, device=input.device
+        )
+        random_tensor.floor_()  # binarize
+        if self._recover:
+            return input.div(keep_prob) * random_tensor
+        else:
+            return input * random_tensor  # as masks
+
+    def forward_with_container(self, input, container, prefix=[]):
+        return self.forward_raw(input)
+
+    def extra_repr(self) -> str:
+        return (
+            "p={:}".format(self._p)
+            + ", dims={:}".format(self._dims)
+            + ", recover={:}".format(self._recover)
+        )