supports id_sensitive=True

intel-analytics · Aug 26, 2022 · 91cb229 · 91cb229
1 parent fdf18bc
commit 91cb229
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 2 deletions.
diff --git a/python/chronos/src/bigdl/chronos/forecaster/base_forecaster.py b/python/chronos/src/bigdl/chronos/forecaster/base_forecaster.py
@@ -449,7 +449,8 @@ def predict(self, data, batch_size=32, quantize=False):
                                              horizon=self.data_config['future_seq_len'],
                                              feature_col=data.roll_feature,
                                              target_col=data.roll_target,
-                                             shuffle=False)
+                                             shuffle=False,
+                                             is_predict=True)
         # data transform
         is_local_data = isinstance(data, (np.ndarray, DataLoader))
         if is_local_data and self.distributed:
@@ -542,7 +543,8 @@ def predict_with_onnx(self, data, batch_size=32, quantize=False):
                                              horizon=self.data_config['future_seq_len'],
                                              feature_col=data.roll_feature,
                                              target_col=data.roll_target,
-                                             shuffle=False)
+                                             shuffle=False,
+                                             is_predict=True)
         if quantize:
             return _pytorch_fashion_inference(model=self.onnxruntime_int8,
                                               input_data=data,
@@ -1151,6 +1153,11 @@ def check_time_steps(tsdataset, past_seq_len, future_seq_len):
                           fixMsg="Do not specify past_seq_len and future seq_len "
                           "or call tsdataset.roll method again and specify time step")
 
+        if tsdataset.id_sensitive:
+            _id_list_len = len(tsdataset.id_col)
+            input_feature_num *= _id_list_len
+            output_feature_num *= _id_list_len
+
         return cls(past_seq_len=past_seq_len,
                    future_seq_len=future_seq_len,
                    input_feature_num=input_feature_num,

diff --git a/python/chronos/src/bigdl/chronos/forecaster/lstm_forecaster.py b/python/chronos/src/bigdl/chronos/forecaster/lstm_forecaster.py
@@ -198,6 +198,11 @@ def check_time_steps(tsdataset, past_seq_len):
                           fixMsg="Do not specify past_seq_len "
                           "or call tsdataset.roll method again and specify time step.")
 
+        if tsdataset.id_sensitive:
+            _id_list_len = len(tsdataset._id_list)
+            input_feature_num *= _id_list_len
+            output_feature_num *= _id_list_len
+
         return cls(past_seq_len=past_seq_len,
                    input_feature_num=input_feature_num,
                    output_feature_num=output_feature_num,

diff --git a/python/chronos/src/bigdl/chronos/forecaster/nbeats_forecaster.py b/python/chronos/src/bigdl/chronos/forecaster/nbeats_forecaster.py
@@ -224,6 +224,9 @@ def check_time_steps(tsdataset, past_seq_len, future_seq_len):
                           f"but found {past_seq_len, future_seq_len}",
                           fixMsg="Do not specify past_seq_len and future seq_len "
                           "or call tsdataset.roll method again and specify time step")
+
+        invalidInputError(tsdataset.id_sensitive and len(tsdataset._id_list) > 1,
+                          "NBeats only supports univariate forecasting.")
 
         return cls(past_seq_len=past_seq_len,
                    future_seq_len=future_seq_len,