update selective learning to support load configuration from json

Author: yisongfu

src/basicts/runners/callback/__init__.py

@@ -4,7 +4,6 @@
 from .curriculum_learrning import CurriculumLearning
 from .early_stopping import EarlyStopping
 from .grad_accumulation import GradAccumulation
-from .koopa_mask_init import KoopaMaskInitCallbackFullTrain
 from .no_bp import NoBP
 from .selective_learning import SelectiveLearning
 

src/basicts/runners/callback/selective_learning.py

@@ -42,14 +42,19 @@ def __init__(
             ckpt_path: Optional[str] = None):
 
         super().__init__()
+
+        # config
         self.r_u = r_u
         self.r_a = r_a
-        self.estimator = estimator(estimator_config)
+        self.estimator = estimator
+        self.estimator_config = estimator_config
         self.ckpt_path = ckpt_path
 
-        if self.r_a is not None and self.estimator is None:
+        self.estimation_model = self.estimator(estimator_config)
+
+        if self.r_a is not None and self.estimation_model is None:
             raise RuntimeError("Anomaly mask ratio is set but estimation model is not provided.")
-        if self.estimator is not None and self.ckpt_path is None:
+        if self.estimation_model is not None and self.ckpt_path is None:
             raise RuntimeError("Estimation model is set but checkpoint path is not provided.")
 
         self.history_residual: torch.Tensor = None
@@ -59,7 +64,7 @@ def __init__(
     def on_train_start(self, runner: "BasicTSRunner"):
         runner.logger.info(f"Use selective learning with r_u={self.r_u}, r_a={self.r_a}.")
         self._load_estimator(runner)
-        self.estimator.eval()
+        self.estimation_model.eval()
         self.num_samples = len(runner.train_data_loader.dataset)
         runner.train_data_loader = _DataLoaderWithIndex(runner.train_data_loader)
 
@@ -86,7 +91,7 @@ def on_compute_loss(self, runner: "BasicTSRunner", **kwargs):
             # Anomaly mask
             if self.r_a is not None:
                 with torch.no_grad():
-                    est_foward_return = runner._forward(self.estimator, data, step=0)
+                    est_foward_return = runner._forward(self.estimation_model, data, step=0)
                 residual_lb = torch.abs(est_foward_return["prediction"] - forward_return["targets"])
                 dist = residual - residual_lb
                 thresholds = torch.quantile(
@@ -103,24 +108,24 @@ def on_epoch_end(self, runner: "BasicTSRunner", **kwargs):
 
     def _load_estimator(self, runner: "BasicTSRunner"):
 
-        runner.logger.info(f"Building estimation model {self.estimator.__class__.__name__}.")
-        self.estimator = to_device(self.estimator)
+        runner.logger.info(f"Building estimation model {self.estimation_model.__class__.__name__}.")
+        self.estimation_model = to_device(self.estimation_model)
 
         # DDP
         if torch.distributed.is_initialized():
-            self.estimator = DDP(
-                self.estimator,
+            self.estimation_model = DDP(
+                self.estimation_model,
                 device_ids=[get_local_rank()],
                 find_unused_parameters=runner.cfg.ddp_find_unused_parameters
             )
 
         # load model weights
         try:
             checkpoint_dict = load_ckpt(None, ckpt_path=self.ckpt_path, logger=runner.logger)
-            if isinstance(self.estimator, DDP):
-                self.estimator.module.load_state_dict(checkpoint_dict["model_state_dict"])
+            if isinstance(self.estimation_model, DDP):
+                self.estimation_model.module.load_state_dict(checkpoint_dict["model_state_dict"])
             else:
-                self.estimator.load_state_dict(checkpoint_dict["model_state_dict"])
+                self.estimation_model.load_state_dict(checkpoint_dict["model_state_dict"])
         except (IndexError, OSError) as e:
             raise OSError(f"Ckpt file {self.ckpt_path} does not exist") from e
 

src/test.py

@@ -0,0 +1,31 @@
+from basicts.models.DLinear import DLinear, DLinearConfig
+from basicts.models.iTransformer import iTransformerConfig, iTransformerForForecasting
+from basicts import BasicTSLauncher
+from basicts.configs import BasicTSForecastingConfig
+from basicts.runners.callback import SelectiveLearning
+
+
+if __name__ == "__main__":
+
+    cb = SelectiveLearning(
+                r_u=0.3,
+                r_a=0.3,
+                estimator=DLinear,
+                estimator_config=DLinearConfig(input_len=336, output_len=336),
+                ckpt_path="checkpoints/DLinear/ETTh1_100_336_336/1f037d3a0fb4a6de40ce3dcb2656b136/DLinear_best_val_MSE.pt"
+            )
+    
+    BasicTSLauncher.launch_training(
+        BasicTSForecastingConfig(
+            model=iTransformerForForecasting,
+            input_len=336,
+            output_len=336,
+            use_timestamps=False,
+            model_config=iTransformerConfig(
+                input_len=336,
+                output_len=336,
+                num_features=7),
+            dataset_name="ETTh1",
+            gpus="0",
+            callbacks=[cb],
+    ))