data selection share handler & trainer

qlib/contrib/meta/data_selection/model.py

@@ -59,6 +59,66 @@ def __init__(
         self.max_epoch = max_epoch
         self.fitted = False
 
+    def run_epoch(self, phase, task_list, epoch, opt, loss_l, ignore_weight=False):
+        if phase == "train":  # phase 0 for training, 1 for inference
+            self.tn.train()
+            torch.set_grad_enabled(True)
+        else:
+            self.tn.eval()
+            torch.set_grad_enabled(False)
+        running_loss = 0.0
+        pred_y_all = []
+        for task in tqdm(task_list, desc=f"{phase} Task", leave=False):
+            meta_input = task.get_meta_input()
+            pred, weights = self.tn(
+                meta_input["X"],
+                meta_input["y"],
+                meta_input["time_perf"],
+                meta_input["time_belong"],
+                meta_input["X_test"],
+                ignore_weight=ignore_weight
+            ) # 这里可能因为如下原因导致pred为None;
+            if self.criterion == "mse":
+                criterion = nn.MSELoss()
+                loss = criterion(pred, meta_input["y_test"])
+            elif self.criterion == "ic_loss":
+                criterion = ICLoss()
+                loss = criterion(pred, meta_input["y_test"], meta_input["test_idx"], skip_size=50)
+
+            if np.isnan(loss.detach().item()): __import__('ipdb').set_trace()
+
+            if phase == "train":
+                opt.zero_grad()
+                norm_loss = nn.MSELoss()
+                loss.backward()
+                opt.step()
+            elif phase == "test":
+                pass
+
+            pred_y_all.append(
+                pd.DataFrame(
+                    {
+                        "pred": pd.Series(pred.detach().cpu().numpy(), index=meta_input["test_idx"]),
+                        "label": pd.Series(
+                            meta_input["y_test"].detach().cpu().numpy(), index=meta_input["test_idx"]
+                        ),
+                    }
+                )
+            )
+            running_loss += loss.detach().item()
+        running_loss = running_loss / len(task_list)
+        loss_l.setdefault(phase, []).append(running_loss)
+
+        pred_y_all = pd.concat(pred_y_all)
+        ic = (
+            pred_y_all.groupby("datetime")
+            .apply(lambda df: df["pred"].corr(df["label"], method="spearman"))
+            .mean()
+        )
+
+        R.log_metrics(**{f"loss/{phase}": running_loss, "step": epoch})
+        R.log_metrics(**{f"ic/{phase}": ic, "step": epoch})
+
     def fit(self, meta_dataset: MetaDatasetHDS):
         """
         The meta-learning-based data selection interacts directly with meta-dataset due to the close-form proxy measurement.
@@ -81,67 +141,18 @@ def fit(self, meta_dataset: MetaDatasetHDS):
             step=self.step, hist_step_n=self.hist_step_n, clip_weight=self.clip_weight, clip_method=self.clip_method
         )
 
-        train_step = 0
         opt = optim.Adam(self.tn.parameters(), lr=self.lr)
+
+        # run weight with no weight
+        for phase, task_list in zip(phases, meta_tasks_l):
+            self.run_epoch(f"{phase}_noweight", task_list, 0, opt, {}, ignore_weight=True)
+            self.run_epoch(f"{phase}_init", task_list, 0, opt, {})
+
+        # run training
         loss_l = {}
         for epoch in tqdm(range(self.max_epoch), desc="epoch"):
             for phase, task_list in zip(phases, meta_tasks_l):
-                if phase == "train":  # phase 0 for training, 1 for inference
-                    self.tn.train()
-                    torch.set_grad_enabled(True)
-                else:
-                    self.tn.eval()
-                    torch.set_grad_enabled(False)
-                running_loss = 0.0
-                pred_y_all = []
-                for task in tqdm(task_list, desc=f"{phase} Task", leave=False):
-                    meta_input = task.get_meta_input()
-                    pred, weights = self.tn(
-                        meta_input["X"],
-                        meta_input["y"],
-                        meta_input["time_perf"],
-                        meta_input["time_belong"],
-                        meta_input["X_test"],
-                    )
-                    if self.criterion == "mse":
-                        criterion = nn.MSELoss()
-                        loss = criterion(pred, meta_input["y_test"])
-                    elif self.criterion == "ic_loss":
-                        criterion = ICLoss()
-                        loss = criterion(pred, meta_input["y_test"], meta_input["test_idx"])
-
-                    if phase == "train":
-                        opt.zero_grad()
-                        norm_loss = nn.MSELoss()
-                        loss.backward()
-                        opt.step()
-                        train_step += 1
-                    elif phase == "test":
-                        pass
-
-                    pred_y_all.append(
-                        pd.DataFrame(
-                            {
-                                "pred": pd.Series(pred.detach().cpu().numpy(), index=meta_input["test_idx"]),
-                                "label": pd.Series(
-                                    meta_input["y_test"].detach().cpu().numpy(), index=meta_input["test_idx"]
-                                ),
-                            }
-                        )
-                    )
-                    running_loss += loss.detach().item()
-                running_loss = running_loss / len(task_list)
-                loss_l.setdefault(phase, []).append(running_loss)
-
-                pred_y_all = pd.concat(pred_y_all)
-                ic = (
-                    pred_y_all.groupby("datetime")
-                    .apply(lambda df: df["pred"].corr(df["label"], method="spearman"))
-                    .mean()
-                )
-
-                R.log_metrics(**{f"loss/{phase}": running_loss, "step": epoch})
-                R.log_metrics(**{f"ic/{phase}": ic, "step": epoch})
+                self.run_epoch(phase, task_list, epoch, opt, loss_l)
             R.save_objects(**{"model.pkl": self.tn})
         self.fitted = True
 

qlib/contrib/meta/data_selection/utils.py

@@ -25,7 +25,7 @@ def get_sim_mat_idx(i_sim_mat, outsample_period):
 
 
 class ICLoss(nn.Module):
-    def forward(self, pred, y, idx):
+    def forward(self, pred, y, idx, skip_size=50):
         """forward.
 
         :param pred:
@@ -41,15 +41,20 @@ def forward(self, pred, y, idx):
         diff_point.append(None)
 
         ic_all = 0.0
+        skip_n = 0
         for start_i, end_i in zip(diff_point, diff_point[1:]):
             pred_focus = pred[start_i:end_i]  # TODO: just for fake
+            if pred_focus.shape[0] < skip_size:
+                # skip some days which have very small amount of stock.
+                skip_n += 1
+                continue
             y_focus = y[start_i:end_i]
             ic_day = torch.dot(
                 (pred_focus - pred_focus.mean()) / np.sqrt(pred_focus.shape[0]) / pred_focus.std(),
                 (y_focus - y_focus.mean()) / np.sqrt(y_focus.shape[0]) / y_focus.std(),
             )
             ic_all += ic_day
-        ic_mean = ic_all / (len(diff_point) - 1)
+        ic_mean = ic_all / (len(diff_point) - 1 - skip_n)
         return -ic_mean  # ic loss
 
 

qlib/model/trainer.py

@@ -25,6 +25,38 @@
 from qlib.data.dataset.weight import Reweighter
 
 
+def _log_task_info(task_config: dict):
+    R.log_params(**flatten_dict(task_config))
+    R.save_objects(**{"task": task_config})  # keep the original format and datatype
+    R.set_tags(**{"hostname": socket.gethostname()})
+
+
+def _exe_task(task_config: dict):
+    rec = R.get_recorder()
+    # model & dataset initiation
+    model: Model = init_instance_by_config(task_config["model"])
+    dataset: Dataset = init_instance_by_config(task_config["dataset"])
+    reweighter: Reweighter = task_config.get("reweighter", None)
+    # model training
+    auto_filter_kwargs(model.fit)(dataset, reweighter=reweighter)
+    R.save_objects(**{"params.pkl": model})
+    # this dataset is saved for online inference. So the concrete data should not be dumped
+    dataset.config(dump_all=False, recursive=True)
+    R.save_objects(**{"dataset": dataset})
+    # generate records: prediction, backtest, and analysis
+    records = task_config.get("record", [])
+    if isinstance(records, dict):  # prevent only one dict
+        records = [records]
+    for record in records:
+        cls, kwargs = get_cls_kwargs(record, default_module="qlib.workflow.record_temp")
+        if cls is SignalRecord:
+            rconf = {"model": model, "dataset": dataset, "recorder": rec}
+        else:
+            rconf = {"recorder": rec}
+        r = cls(**kwargs, **rconf)
+        r.generate()
+
+
 def begin_task_train(task_config: dict, experiment_name: str, recorder_name: str = None) -> Recorder:
     """
     Begin task training to start a recorder and save the task config.
@@ -38,11 +70,8 @@ def begin_task_train(task_config: dict, experiment_name: str, recorder_name: str
         Recorder: the model recorder
     """
     with R.start(experiment_name=experiment_name, recorder_name=recorder_name):
-        R.log_params(**flatten_dict(task_config))
-        R.save_objects(**{"task": task_config})  # keep the original format and datatype
-        R.set_tags(**{"hostname": socket.gethostname()})
-        recorder: Recorder = R.get_recorder()
-    return recorder
+        _log_task_info(task_config)
+        return R.get_recorder()
 
 
 def end_task_train(rec: Recorder, experiment_name: str) -> Recorder:
@@ -58,29 +87,7 @@ def end_task_train(rec: Recorder, experiment_name: str) -> Recorder:
     """
     with R.start(experiment_name=experiment_name, recorder_id=rec.info["id"], resume=True):
         task_config = R.load_object("task")
-        # model & dataset initiation
-        model: Model = init_instance_by_config(task_config["model"])
-        dataset: Dataset = init_instance_by_config(task_config["dataset"])
-        reweighter: Reweighter = task_config.get("reweighter", None)
-        # model training
-        auto_filter_kwargs(model.fit)(dataset, reweighter=reweighter)
-        R.save_objects(**{"params.pkl": model})
-        # this dataset is saved for online inference. So the concrete data should not be dumped
-        dataset.config(dump_all=False, recursive=True)
-        R.save_objects(**{"dataset": dataset})
-        # generate records: prediction, backtest, and analysis
-        records = task_config.get("record", [])
-        if isinstance(records, dict):  # prevent only one dict
-            records = [records]
-        for record in records:
-            cls, kwargs = get_cls_kwargs(record, default_module="qlib.workflow.record_temp")
-            if cls is SignalRecord:
-                rconf = {"model": model, "dataset": dataset, "recorder": rec}
-            else:
-                rconf = {"recorder": rec}
-            r = cls(**kwargs, **rconf)
-            r.generate()
-
+        _exe_task(task_config)
     return rec
 
 
@@ -101,9 +108,10 @@ def task_train(task_config: dict, experiment_name: str, recorder_name: str = Non
     ----------
     Recorder: The instance of the recorder
     """
-    recorder = begin_task_train(task_config, experiment_name, recorder_name=recorder_name)
-    recorder = end_task_train(recorder, experiment_name)
-    return recorder
+    with R.start(experiment_name=experiment_name, recorder_name=recorder_name):
+        _log_task_info(task_config)
+        _exe_task(task_config)
+        return R.get_recorder()
 
 
 class Trainer:

qlib/utils/__init__.py

@@ -594,7 +594,7 @@ def lazy_sort_index(df: pd.DataFrame, axis=0) -> pd.DataFrame:
         sorted dataframe
     """
     idx = df.index if axis == 0 else df.columns
-    if idx.is_monotonic_increasing:
+    if idx.is_monotonic_increasing and (not isinstance(idx, pd.MultiIndex) or not idx.is_lexsorted()):
         return df
     else:
         return df.sort_index(axis=axis)
@@ -657,7 +657,7 @@ def _func(*args, **kwargs):
         for k, v in kwargs.items():
             # if `func` don't accept variable keyword arguments like `**kwargs` and have not according named arguments
             if spec.varkw is None and k not in spec.args:
-                log.warn(f"The parameter `{k}` with value `{v}` is ignored.")
+                log.warning(f"The parameter `{k}` with value `{v}` is ignored.")
             else:
                 new_kwargs[k] = v
         return func(*args, **new_kwargs)

qlib/workflow/__init__.py

@@ -3,7 +3,7 @@
 
 from contextlib import contextmanager
 from typing import Text, Optional
-from .expm import MLflowExpManager
+from .expm import ExpManager
 from .exp import Experiment
 from .recorder import Recorder
 from ..utils import Wrapper
@@ -15,7 +15,7 @@ class QlibRecorder:
     A global system that helps to manage the experiments.
     """
 
-    def __init__(self, exp_manager):
+    def __init__(self, exp_manager: ExpManager):
         self.exp_manager = exp_manager
 
     def __repr__(self):
@@ -334,6 +334,23 @@ def set_uri(self, uri: Optional[Text]):
         """
         self.exp_manager.set_uri(uri)
 
+    @contextmanager
+    def uri_context(self, uri: Text):
+        """
+        Temporarily set the exp_manager's uri to uri
+
+        Parameters
+        ----------
+        uri : Text
+            the temporal uri
+        """
+        prev_uri = self.exp_manager._current_uri
+        self.exp_manager.set_uri(uri)
+        try:
+            yield
+        finally:
+            self.exp_manager._current_uri = prev_uri
+
     def get_recorder(
         self, *, recorder_id=None, recorder_name=None, experiment_id=None, experiment_name=None
     ) -> Recorder:

qlib/workflow/expm.py

@@ -14,7 +14,7 @@
 from .recorder import Recorder
 from ..log import get_module_logger
 
-logger = get_module_logger("workflow", logging.INFO)
+logger = get_module_logger("workflow")
 
 
 class ExpManager:
@@ -258,7 +258,7 @@ def set_uri(self, uri: Optional[Text] = None):
 
         """
         if uri is None:
-            logger.info("No tracking URI is provided. Use the default tracking URI.")
+            logger.debug("No tracking URI is provided. Use the default tracking URI.")
             self._current_uri = self.default_uri
         else:
             # Temporarily re-set the current uri as the uri argument.
@@ -269,6 +269,7 @@ def set_uri(self, uri: Optional[Text] = None):
     def _set_uri(self):
         """
         Customized features for subclasses' set_uri function.
+        This method is designed for the underlying experiment backend storage.
         """
         raise NotImplementedError(f"Please implement the `_set_uri` method.")