Update baselines

exps/GeMOSA/baselines/slbm-nof.py

@@ -72,10 +72,11 @@ def main(args):
             "This task ({:}) is not supported.".format(all_env.meta_info["task"])
         )
 
-    seq_length = 10
-    seq_times = env.get_seq_times(0, seq_length)
+    seq_times = env.get_seq_times(0, args.seq_length)
     _, (allxs, allys) = env.seq_call(seq_times)
-    allxs, allys = allxs.view(-1, 1), allys.view(-1, 1)
+    allxs, allys = allxs.view(-1, allxs.shape[-1]), allys.view(-1, 1)
+    if env.meta_info["task"] == "classification":
+        allys = allys.view(-1)
 
     historical_x, historical_y = allxs.to(args.device), allys.to(args.device)
     model = get_model(**model_kwargs)
@@ -83,28 +84,28 @@ def main(args):
 
     optimizer = torch.optim.Adam(model.parameters(), lr=args.init_lr, amsgrad=True)
     lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
-            optimizer,
-            milestones=[
-                int(args.epochs * 0.25),
-                int(args.epochs * 0.5),
-                int(args.epochs * 0.75),
-            ],
-            gamma=0.3,
-        )
+        optimizer,
+        milestones=[
+            int(args.epochs * 0.25),
+            int(args.epochs * 0.5),
+            int(args.epochs * 0.75),
+        ],
+        gamma=0.3,
+    )
 
     train_metric = metric_cls(True)
     best_loss, best_param = None, None
     for _iepoch in range(args.epochs):
-            preds = model(historical_x)
-            optimizer.zero_grad()
-            loss = criterion(preds, historical_y)
-            loss.backward()
-            optimizer.step()
-            lr_scheduler.step()
-            # save best
-            if best_loss is None or best_loss > loss.item():
-                best_loss = loss.item()
-                best_param = copy.deepcopy(model.state_dict())
+        preds = model(historical_x)
+        optimizer.zero_grad()
+        loss = criterion(preds, historical_y)
+        loss.backward()
+        optimizer.step()
+        lr_scheduler.step()
+        # save best
+        if best_loss is None or best_loss > loss.item():
+            best_loss = loss.item()
+            best_param = copy.deepcopy(model.state_dict())
     model.load_state_dict(best_param)
     model.analyze_weights()
     with torch.no_grad():
@@ -126,7 +127,7 @@ def main(args):
             + need_time
         )
         # train the same data
-    
+
         # build optimizer
         xmetric = ComposeMetric(metric_cls(True), SaveMetric())
         future_x.to(args.device), future_y.to(args.device)
@@ -176,6 +177,9 @@ if __name__ == "__main__":
         required=True,
         help="The hidden dimension.",
     )
+    parser.add_argument(
+        "--seq_length", type=int, default=10, help="The sequence length."
+    )
     parser.add_argument(
         "--init_lr",
         type=float,
@@ -213,12 +217,11 @@ if __name__ == "__main__":
         args.save_dir, args.hidden_dim, args.epochs, args.init_lr, args.env_version
     )
     if args.rand_seed is None or args.rand_seed < 0:
-        args.rand_seed = random.randint(1, 100000)
-        main(args)
-    else:
         results = []
         for iseed in range(3):
-          args.rand_seed = random.randint(1, 100000)
-          result = main(args)
-          results.append(result)
-        show_mean_var(result)
+            args.rand_seed = random.randint(1, 100000)
+            result = main(args)
+            results.append(result)
+        show_mean_var(results)
+    else:
+        main(args)