Move to A800

zqiao11 · Dec 25, 2024 · ee41984 · ee41984
1 parent 09d4e94
commit ee41984
Show file tree

Hide file tree

Showing 6 changed files with 30 additions and 15 deletions.
diff --git a/project/pf/multi_scale/finetune/base/turkey_power.sh b/project/pf/multi_scale/finetune/base/turkey_power.sh
@@ -29,4 +29,4 @@ val_data=${data} \
 val_data.patch_size=${ps} \
 val_data.context_length=$cl \
 val_data.prediction_length=$pl \
-trainer.callbacks.'2'.patience=10
+model.lr=5e-6
diff --git a/project/pf/multi_scale/finetune/small/turkey_power.sh b/project/pf/multi_scale/finetune/small/turkey_power.sh
@@ -28,4 +28,5 @@ data.prediction_length=$pl \
 val_data=${data} \
 val_data.patch_size=${ps} \
 val_data.context_length=$cl \
-val_data.prediction_length=$pl
+val_data.prediction_length=$pl \
+model.lr=5e-6
diff --git a/project/pf/single_scale/finetune/base/turkey_power.sh b/project/pf/single_scale/finetune/base/turkey_power.sh
@@ -29,4 +29,4 @@ val_data=${data} \
 val_data.patch_size=${ps} \
 val_data.context_length=$cl \
 val_data.prediction_length=$pl \
-trainer.callbacks.'2'.patience=10
+model.lr=5e-6
diff --git a/project/pf/single_scale/finetune/small/turkey_power.sh b/project/pf/single_scale/finetune/small/turkey_power.sh
@@ -29,4 +29,4 @@ val_data=${data} \
 val_data.patch_size=${ps} \
 val_data.context_length=$cl \
 val_data.prediction_length=$pl \
-trainer.callbacks.'2'.patience=10
+model.lr=5e-6
diff --git a/src/uni2ts/model/multi_scale_moirai/finetune.py b/src/uni2ts/model/multi_scale_moirai/finetune.py
@@ -42,6 +42,7 @@
 from uni2ts.module.ts_embed import MultiInSizeLinear, MultiOutSizeLinear
 from uni2ts.optim import SchedulerType, get_scheduler
 from uni2ts.transform import (
+    AddNewScaleSeries,
     AddNewScaleContextSeries,
     AddObservedMask,
     AddSampleIndex,
@@ -485,6 +486,12 @@ def default_train_transform(
                     optional_fields=("past_feat_dynamic_real",),
                 )
                 #  QZ: Apply downsample to target. Create a new field 'target{i}' for each scale.
+                # + AddNewScaleSeries(
+                #     target_field="target",
+                #     ds_factor=self.ds_factor,
+                #     new_scales_target_fields=self.new_scales_target_fields,
+                #     expected_ndim=2,
+                # )
                 + AddNewScaleContextSeries(
                     target_field="target",
                     ds_factor=self.ds_factor,
@@ -624,6 +631,12 @@ def default_val_transform(
                     fields=("target",),
                     optional_fields=("past_feat_dynamic_real",),
                 )
+                # + AddNewScaleSeries(
+                #     target_field="target",
+                #     ds_factor=self.ds_factor,
+                #     new_scales_target_fields=self.new_scales_target_fields,
+                #     expected_ndim=2,
+                # )
                 + AddNewScaleContextSeries(
                     target_field="target",
                     ds_factor=self.ds_factor,

diff --git a/src/uni2ts/transform/multi_scale.py b/src/uni2ts/transform/multi_scale.py
@@ -101,7 +101,8 @@ class AddNewScaleSeries(CheckArrNDimMixin, Transformation):
     """
 
     target_field: str
-    num_new_scales: int
+    ds_factor: int
+    new_scales_target_fields: tuple[str, ...]
     expected_ndim: int = 2
 
     def __post_init__(self):
@@ -114,19 +115,19 @@ def __post_init__(self):
 
     def __call__(self, data_entry: dict[str, Any]) -> dict[str, Any]:
         self.__post_init__()
-        for i in range(self.num_new_scales):
-            data_entry[f"target{i+1}"] = self._downsample(
+        for field in self.new_scales_target_fields:
+            data_entry[field] = self._downsample(
                 data_entry,
                 self.target_field,
             )
 
-        for i in range(self.num_new_scales):
-            self.context_length_new_scales[f"target{i+1}"] = (
-                self.new_context_length_list[i]
-            )
-            self.prediction_length_new_scales[f"target{i+1}"] = (
-                self.new_prediction_length_list[i]
-            )
+        for field in self.new_scales_target_fields:
+            self.context_length_new_scales[field] = self.new_context_length_list[
+                int(field[-1])
+            ]
+            self.prediction_length_new_scales[field] = self.new_prediction_length_list[
+                int(field[-1])
+            ]
         data_entry["context_length_new_scales"] = self.context_length_new_scales
         data_entry["prediction_length_new_scales"] = self.prediction_length_new_scales
 
@@ -142,7 +143,7 @@ def _downsample(self, data_entry: dict[str, Any], field: str) -> np.ndarray:
 
         self.check_ndim(field, arr, self.expected_ndim)
         dim, time = arr.shape[:2]
-        ds_factor = 2
+        ds_factor = self.ds_factor
 
         if len(self.new_context_length_list) == 0:
             context_length = data_entry["context_length"]