diff --git a/exps/LFNA/lfna.py b/exps/LFNA/lfna.py index f13c853..89be5a9 100644 --- a/exps/LFNA/lfna.py +++ b/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) - dynamic_env.reset_max_seq_length(args.seq_length) - """ - env_loader = torch.utils.data.DataLoader( - dynamic_env, + batch_sampler = EnvSampler(train_env, args.meta_batch, args.sampler_enlarge) + train_env.reset_max_seq_length(args.seq_length) + valid_env.reset_max_seq_length(args.seq_length) + valid_env_loader = torch.utils.data.DataLoader( + valid_env, batch_size=args.meta_batch, shuffle=True, num_workers=args.workers, pin_memory=True, ) - """ - env_loader = torch.utils.data.DataLoader( - dynamic_env, + train_env_loader = torch.utils.data.DataLoader( + train_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, diff --git a/exps/LFNA/lfna_meta_model.py b/exps/LFNA/lfna_meta_model.py index 79f6bd0..a3871f1 100644 --- a/exps/LFNA/lfna_meta_model.py +++ b/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 diff --git a/lib/datasets/synthetic_env.py b/lib/datasets/synthetic_env.py index e8e519c..7f5e33c 100644 --- a/lib/datasets/synthetic_env.py +++ b/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, ) diff --git a/lib/xlayers/super_core.py b/lib/xlayers/super_core.py index 03aa6c0..0055f95 100644 --- a/lib/xlayers/super_core.py +++ b/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, diff --git a/lib/xlayers/super_dropout.py b/lib/xlayers/super_dropout.py index 124f2db..d5ed994 100644 --- a/lib/xlayers/super_dropout.py +++ b/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) + )