Updates

exps/LFNA/lfna.py

@@ -93,7 +93,7 @@ def epoch_evaluate(loader, meta_model, base_model, criterion, device, logger):
     return loss_meter
 
 
-def pretrain(base_model, meta_model, criterion, xenv, args, logger):
+def pretrain_v2(base_model, meta_model, criterion, xenv, args, logger):
     optimizer = torch.optim.Adam(
         meta_model.parameters(),
         lr=args.lr,
@@ -164,6 +164,69 @@ def pretrain(base_model, meta_model, criterion, xenv, args, logger):
         start_time = time.time()
 
 
+def pretrain(base_model, meta_model, criterion, xenv, args, logger):
+    base_model.train()
+    meta_model.train()
+    optimizer = torch.optim.Adam(
+        meta_model.parameters(),
+        lr=args.lr,
+        weight_decay=args.weight_decay,
+        amsgrad=True,
+    )
+    logger.log("Pre-train the meta-model")
+    logger.log("Using the optimizer: {:}".format(optimizer))
+
+    meta_model.set_best_dir(logger.path(None) / "ckps-basic-pretrain")
+    meta_model.set_best_name("pretrain-{:}.pth".format(args.rand_seed))
+    per_epoch_time, start_time = AverageMeter(), time.time()
+    for iepoch in range(args.epochs):
+        left_time = "Time Left: {:}".format(
+            convert_secs2time(per_epoch_time.avg * (args.epochs - iepoch), True)
+        )
+        losses = []
+        optimizer.zero_grad()
+        for ibatch in range(args.meta_batch):
+            rand_index = random.randint(0, meta_model.meta_length - xenv.seq_length - 1)
+            timestamps = meta_model.meta_timestamps[
+                rand_index : rand_index + xenv.seq_length
+            ]
+
+            seq_timestamps, (seq_inputs, seq_targets) = xenv.seq_call(timestamps)
+            time_embeds = meta_model.super_meta_embed[
+                rand_index : rand_index + xenv.seq_length
+            ]
+            [seq_containers], time_embeds = meta_model(
+                None, torch.unsqueeze(time_embeds, dim=0)
+            )
+            seq_inputs, seq_targets = seq_inputs.to(args.device), seq_targets.to(
+                args.device
+            )
+            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()
+        final_loss.backward()
+        optimizer.step()
+        # success
+        success, best_score = meta_model.save_best(-final_loss.item())
+        logger.log(
+            "{:} [{:04d}/{:}] loss : {:.5f}".format(
+                time_string(),
+                iepoch,
+                args.epochs,
+                final_loss.item(),
+            )
+            + ", batch={:}".format(len(losses))
+            + ", success={:}, best_score={:.4f}".format(success, -best_score)
+            + " {:}".format(left_time)
+        )
+        per_epoch_time.update(time.time() - start_time)
+        start_time = time.time()
+
+
 def main(args):
     logger, env_info, model_kwargs = lfna_setup(args)
     train_env = get_synthetic_env(mode="train", version=args.env_version)