Update default stride to be derived from n_step_outputs but can be overridden

Author: sgreenbury

configs/processor.yaml

@@ -14,7 +14,7 @@ output:
 training:
   n_steps_input: 1
   n_steps_output: 4
-  stride: 4
+  stride: null
   autoencoder_checkpoint: null
   freeze_autoencoder: false
 

src/autocast/eval/processor.py

@@ -107,6 +107,12 @@ def parse_args() -> argparse.Namespace:
         default=None,
         help="Override training.n_steps_output (number of target time steps).",
     )
+    parser.add_argument(
+        "--stride",
+        type=int,
+        default=None,
+        help="Override training stride used for rollouts (defaults to n_steps_output).",
+    )
     parser.add_argument(
         "--checkpoint",
         type=Path,

src/autocast/train/configuration.py

@@ -18,6 +18,7 @@ class TrainingParams:
 
     n_steps_input: int
     n_steps_output: int
+    stride: int
     autoencoder_checkpoint: Path | None
     freeze_autoencoder: bool
 
@@ -135,6 +136,7 @@ def resolve_training_params(cfg: DictConfig, args) -> TrainingParams:
         if training_cfg is not None
         else False
     )
+    stride_cfg = training_cfg.get("stride") if training_cfg is not None else None
 
     n_steps_input = args.n_steps_input or n_steps_input_cfg
     n_steps_output = args.n_steps_output or n_steps_output_cfg
@@ -147,13 +149,26 @@ def resolve_training_params(cfg: DictConfig, args) -> TrainingParams:
         args.freeze_autoencoder if args.freeze_autoencoder is not None else freeze_cfg
     )
 
+    if stride_cfg in (None, "auto"):
+        stride_cfg = n_steps_output
+    stride_override = getattr(args, "stride", None)
+    stride = stride_override or stride_cfg or n_steps_output
+    if stride < 1:
+        msg = "stride must be >= 1."
+        raise ValueError(msg)
+
+    if training_cfg is not None:
+        with open_dict(training_cfg):
+            training_cfg["stride"] = stride
+
     if n_steps_output < 1:
         msg = "n_steps_output must be >= 1 for processor training."
         raise ValueError(msg)
 
     return TrainingParams(
         n_steps_input=n_steps_input,
         n_steps_output=n_steps_output,
+        stride=stride,
         autoencoder_checkpoint=checkpoint,
         freeze_autoencoder=freeze_autoencoder,
     )

src/autocast/train/processor.py

@@ -86,6 +86,12 @@ def parse_args() -> argparse.Namespace:
         default=None,
         help="Override training.n_steps_output (number of target time steps).",
     )
+    parser.add_argument(
+        "--stride",
+        type=int,
+        default=None,
+        help="Override training stride (rollout interval between predictions).",
+    )
     parser.add_argument(
         "--work-dir",
         type=Path,
@@ -254,8 +260,7 @@ def main() -> None:  # noqa: PLR0915
     max_rollout_steps = epd_cfg.get("max_rollout_steps", 10)
     loss_cfg = epd_cfg.get("loss_func")
     loss_func = instantiate(loss_cfg) if loss_cfg is not None else nn.MSELoss()
-    training_cfg = cfg.get("training") or {}
-    stride = training_cfg.get("stride", 1)
+    stride = training_params.stride
 
     model = EncoderProcessorDecoder(
         encoder_decoder=encoder_decoder,