Update LFNA

exps/LFNA/lfna-tall-hpnet.py

@@ -1,7 +1,7 @@
 #####################################################
 # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.04 #
 #####################################################
-# python exps/LFNA/lfna-tall-hpnet.py --env_version v1 --hidden_dim 16 --epochs 100000 --meta_batch 16
+# python exps/LFNA/lfna-tall-hpnet.py --env_version v1 --hidden_dim 16 --epochs 100000 --meta_batch 64
 #####################################################
 import sys, time, copy, torch, random, argparse
 from tqdm import tqdm
@@ -33,7 +33,7 @@ from lfna_models import HyperNet
 def main(args):
     logger, env_info, model_kwargs = lfna_setup(args)
     dynamic_env = env_info["dynamic_env"]
-    model = get_model(dict(model_type="simple_mlp"), **model_kwargs)
+    model = get_model(**model_kwargs)
     criterion = torch.nn.MSELoss()
 
     logger.log("There are {:} weights.".format(model.get_w_container().numel()))
@@ -72,7 +72,7 @@ def main(args):
         )
 
         limit_bar = float(iepoch + 1) / args.epochs * total_bar
-        limit_bar = min(max(0, int(limit_bar)), total_bar)
+        limit_bar = min(max(32, int(limit_bar)), total_bar)
         losses = []
         for ibatch in range(args.meta_batch):
             cur_time = random.randint(0, limit_bar)

exps/LFNA/lfna-test-hpnet.py

@@ -1,7 +1,7 @@
 #####################################################
 # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2021.04 #
 #####################################################
-# python exps/LFNA/lfna-test-hpnet.py --env_version v1 --hidden_dim 16
+# python exps/LFNA/lfna-test-hpnet.py --env_version v1 --hidden_dim 16 --layer_dim 32 --epochs 50000
 #####################################################
 import sys, time, copy, torch, random, argparse
 from tqdm import tqdm
@@ -33,17 +33,17 @@ from lfna_models import HyperNet
 def main(args):
     logger, env_info, model_kwargs = lfna_setup(args)
     dynamic_env = env_info["dynamic_env"]
-    model = get_model(dict(model_type="simple_mlp"), **model_kwargs)
+    model = get_model(**model_kwargs)
     model = model.to(args.device)
     criterion = torch.nn.MSELoss()
 
     logger.log("There are {:} weights.".format(model.get_w_container().numel()))
 
     shape_container = model.get_w_container().to_shape_container()
-    hypernet = HyperNet(shape_container, args.hidden_dim, args.task_dim)
+    hypernet = HyperNet(shape_container, args.layer_dim, args.task_dim)
     hypernet = hypernet.to(args.device)
     # task_embed = torch.nn.Parameter(torch.Tensor(env_info["total"], args.task_dim))
-    total_bar = 10
+    total_bar = 16
     task_embeds = []
     for i in range(total_bar):
         tensor = torch.Tensor(1, args.task_dim).to(args.device)
@@ -51,8 +51,12 @@ def main(args):
     for task_embed in task_embeds:
         trunc_normal_(task_embed, std=0.02)
 
+    model.train()
+    hypernet.train()
+
     parameters = list(hypernet.parameters()) + task_embeds
-    optimizer = torch.optim.Adam(parameters, lr=args.init_lr, amsgrad=True)
+    # optimizer = torch.optim.Adam(parameters, lr=args.init_lr, amsgrad=True)
+    optimizer = torch.optim.Adam(parameters, lr=args.init_lr, weight_decay=1e-5)
     lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
         optimizer,
         milestones=[
@@ -98,7 +102,7 @@ def main(args):
         lr_scheduler.step()
 
         loss_meter.update(final_loss.item())
-        if iepoch % 200 == 0:
+        if iepoch % 100 == 0:
             logger.log(
                 head_str
                 + " meta-loss: {:.4f} ({:.4f}) :: lr={:.5f}, batch={:}".format(
@@ -126,6 +130,26 @@ def main(args):
     print(model)
     print(hypernet)
 
+    w_container_per_epoch = dict()
+    for idx in range(0, total_bar):
+        future_time = env_info["{:}-timestamp".format(idx)]
+        future_x = env_info["{:}-x".format(idx)]
+        future_y = env_info["{:}-y".format(idx)]
+        future_container = hypernet(task_embeds[idx])
+        w_container_per_epoch[idx] = future_container.no_grad_clone()
+        with torch.no_grad():
+            future_y_hat = model.forward_with_container(
+                future_x, w_container_per_epoch[idx]
+            )
+            future_loss = criterion(future_y_hat, future_y)
+        logger.log("meta-test: [{:03d}] -> loss={:.4f}".format(idx, future_loss.item()))
+
+    save_checkpoint(
+        {"w_container_per_epoch": w_container_per_epoch},
+        logger.path(None) / "final-ckp.pth",
+        logger,
+    )
+
     logger.log("-" * 200 + "\n")
     logger.close()
 
@@ -150,6 +174,12 @@ if __name__ == "__main__":
         required=True,
         help="The hidden dimension.",
     )
+    parser.add_argument(
+        "--layer_dim",
+        type=int,
+        required=True,
+        help="The hidden dimension.",
+    )
     #####
     parser.add_argument(
         "--init_lr",
@@ -181,7 +211,7 @@ if __name__ == "__main__":
     if args.rand_seed is None or args.rand_seed < 0:
         args.rand_seed = random.randint(1, 100000)
     assert args.save_dir is not None, "The save dir argument can not be None"
-    args.task_dim = args.hidden_dim
+    args.task_dim = args.layer_dim
     args.save_dir = "{:}-{:}-d{:}".format(
         args.save_dir, args.env_version, args.hidden_dim
     )

exps/LFNA/lfna-ttss-hpnet.py

@@ -31,7 +31,7 @@ from lfna_models import HyperNet_VX as HyperNet
 def main(args):
     logger, env_info, model_kwargs = lfna_setup(args)
     dynamic_env = env_info["dynamic_env"]
-    model = get_model(dict(model_type="simple_mlp"), **model_kwargs)
+    model = get_model(**model_kwargs)
 
     total_time = env_info["total"]
     for i in range(total_time):

exps/LFNA/lfna_models.py

@@ -4,6 +4,8 @@
 import copy
 import torch
 
+import torch.nn.functional as F
+
 from xlayers import super_core
 from xlayers import trunc_normal_
 from models.xcore import get_model
@@ -29,13 +31,15 @@ class HyperNet(super_core.SuperModule):
         trunc_normal_(self._super_layer_embed, std=0.02)
 
         model_kwargs = dict(
+            config=dict(model_type="dual_norm_mlp"),
             input_dim=layer_embeding + task_embedding,
             output_dim=max(self._numel_per_layer),
-            hidden_dims=[layer_embeding * 4] * 4,
+            hidden_dims=[layer_embeding * 4] * 3,
             act_cls="gelu",
             norm_cls="layer_norm_1d",
+            dropout=0.1,
         )
-        self._generator = get_model(dict(model_type="norm_mlp"), **model_kwargs)
+        self._generator = get_model(**model_kwargs)
         """
         model_kwargs = dict(
             input_dim=layer_embeding + task_embedding,
@@ -50,8 +54,12 @@ class HyperNet(super_core.SuperModule):
         print("generator: {:}".format(self._generator))
 
     def forward_raw(self, task_embed):
+        # task_embed = F.normalize(task_embed, dim=-1, p=2)
+        # layer_embed = F.normalize(self._super_layer_embed, dim=-1, p=2)
+        layer_embed = self._super_layer_embed
         task_embed = task_embed.view(1, -1).expand(self._num_layers, -1)
-        joint_embed = torch.cat((task_embed, self._super_layer_embed), dim=-1)
+
+        joint_embed = torch.cat((task_embed, layer_embed), dim=-1)
         weights = self._generator(joint_embed)
         if self._return_container:
             weights = torch.split(weights, 1)