Rapidflow fix

.github/workflows/lightning.yml

@@ -16,7 +16,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        lightning: ["2.1.4", "2.2.5", "2.3.3", "2.4.0", "2.5.0"]
+        lightning: ["2.1.4", "2.2.5", "2.3.3", "2.4.0", "2.5.1"]
 
     steps:
       - uses: actions/checkout@v4

ptlflow/models/base_model/base_model.py

@@ -596,33 +596,24 @@ def _split_train_val_metrics(
             dataset_name = inputs_meta["dataset_name"][0].lower()
 
         log_metrics = {}
-        has_train = False
-        has_val = False
         for k, v in metrics.items():
             if k.startswith("val/"):
                 k = k[4:]
 
+            if dataset_name is not None:
+                log_metrics[f"val_{dataset_name}/full/{k}"] = v
+            else:
+                log_metrics[f"val/full/{k}"] = v
+
             if inputs_meta is not None and inputs_meta.get("is_val") is not None:
                 if inputs_meta["is_val"][0]:
                     split = "val"
-                    has_val = True
                 else:
                     split = "train"
-                    has_train = True
 
                 if dataset_name is not None:
                     log_metrics[f"val_{dataset_name}/{split}/{k}"] = v
                 else:
                     log_metrics[f"val/{split}/{k}"] = v
 
-        if has_train and has_val:
-            for k, v in metrics.items():
-                if k.startswith("val/"):
-                    k = k[4:]
-
-                if dataset_name is not None:
-                    log_metrics[f"val_{dataset_name}/full/{k}"] = v
-                else:
-                    log_metrics[f"val/full/{k}"] = v
-
         return log_metrics

ptlflow/models/rapidflow/next1d.py

@@ -111,7 +111,7 @@ def forward(self, input):
             x = F.conv2d(
                 input,
                 self.weight_h,
-                self.bias,
+                None,
                 self.stride,
                 (0, self.padding),
                 self.dilation,

ptlflow/models/rapidflow/rapidflow.py

@@ -426,47 +426,11 @@ def __init__(
         self,
         pyramid_ranges: tuple[int, int] = (32, 32),
         iters: int = 1,
-        corr_mode: str = "allpairs",
-        corr_levels: int = 1,
-        corr_range: int = 4,
-        enc_hidden_chs: int = 64,
-        enc_out_chs: int = 128,
-        enc_stem_stride: int = 4,
-        enc_mlp_ratio: float = 4,
-        enc_depth: int = 4,
-        dec_net_chs: int = 64,
-        dec_inp_chs: int = 64,
-        dec_motion_chs: int = 128,
-        dec_depth: int = 2,
-        dec_mlp_ratio: float = 4,
-        use_upsample_mask: bool = False,
-        fuse_next1d_weights: bool = False,
-        simple_io: bool = False,
-        gamma: float = 0.8,
-        max_flow: float = 400,
         **kwargs,
     ) -> None:
         super().__init__(
             pyramid_ranges,
             iters,
-            corr_mode,
-            corr_levels,
-            corr_range,
-            enc_hidden_chs,
-            enc_out_chs,
-            enc_stem_stride,
-            enc_mlp_ratio,
-            enc_depth,
-            dec_net_chs,
-            dec_inp_chs,
-            dec_motion_chs,
-            dec_depth,
-            dec_mlp_ratio,
-            use_upsample_mask,
-            fuse_next1d_weights,
-            simple_io,
-            gamma,
-            max_flow,
             **kwargs,
         )
 
@@ -476,47 +440,11 @@ def __init__(
         self,
         pyramid_ranges: tuple[int, int] = (32, 16),
         iters: int = 2,
-        corr_mode: str = "allpairs",
-        corr_levels: int = 1,
-        corr_range: int = 4,
-        enc_hidden_chs: int = 64,
-        enc_out_chs: int = 128,
-        enc_stem_stride: int = 4,
-        enc_mlp_ratio: float = 4,
-        enc_depth: int = 4,
-        dec_net_chs: int = 64,
-        dec_inp_chs: int = 64,
-        dec_motion_chs: int = 128,
-        dec_depth: int = 2,
-        dec_mlp_ratio: float = 4,
-        use_upsample_mask: bool = False,
-        fuse_next1d_weights: bool = False,
-        simple_io: bool = False,
-        gamma: float = 0.8,
-        max_flow: float = 400,
         **kwargs,
     ) -> None:
         super().__init__(
             pyramid_ranges,
             iters,
-            corr_mode,
-            corr_levels,
-            corr_range,
-            enc_hidden_chs,
-            enc_out_chs,
-            enc_stem_stride,
-            enc_mlp_ratio,
-            enc_depth,
-            dec_net_chs,
-            dec_inp_chs,
-            dec_motion_chs,
-            dec_depth,
-            dec_mlp_ratio,
-            use_upsample_mask,
-            fuse_next1d_weights,
-            simple_io,
-            gamma,
-            max_flow,
             **kwargs,
         )
 
@@ -526,47 +454,11 @@ def __init__(
         self,
         pyramid_ranges: tuple[int, int] = (32, 8),
         iters: int = 3,
-        corr_mode: str = "allpairs",
-        corr_levels: int = 1,
-        corr_range: int = 4,
-        enc_hidden_chs: int = 64,
-        enc_out_chs: int = 128,
-        enc_stem_stride: int = 4,
-        enc_mlp_ratio: float = 4,
-        enc_depth: int = 4,
-        dec_net_chs: int = 64,
-        dec_inp_chs: int = 64,
-        dec_motion_chs: int = 128,
-        dec_depth: int = 2,
-        dec_mlp_ratio: float = 4,
-        use_upsample_mask: bool = True,
-        fuse_next1d_weights: bool = False,
-        simple_io: bool = False,
-        gamma: float = 0.8,
-        max_flow: float = 400,
         **kwargs,
     ) -> None:
         super().__init__(
             pyramid_ranges,
             iters,
-            corr_mode,
-            corr_levels,
-            corr_range,
-            enc_hidden_chs,
-            enc_out_chs,
-            enc_stem_stride,
-            enc_mlp_ratio,
-            enc_depth,
-            dec_net_chs,
-            dec_inp_chs,
-            dec_motion_chs,
-            dec_depth,
-            dec_mlp_ratio,
-            use_upsample_mask,
-            fuse_next1d_weights,
-            simple_io,
-            gamma,
-            max_flow,
             **kwargs,
         )
 
@@ -576,47 +468,11 @@ def __init__(
         self,
         pyramid_ranges: tuple[int, int] = (32, 8),
         iters: int = 6,
-        corr_mode: str = "allpairs",
-        corr_levels: int = 1,
-        corr_range: int = 4,
-        enc_hidden_chs: int = 64,
-        enc_out_chs: int = 128,
-        enc_stem_stride: int = 4,
-        enc_mlp_ratio: float = 4,
-        enc_depth: int = 4,
-        dec_net_chs: int = 64,
-        dec_inp_chs: int = 64,
-        dec_motion_chs: int = 128,
-        dec_depth: int = 2,
-        dec_mlp_ratio: float = 4,
-        use_upsample_mask: bool = True,
-        fuse_next1d_weights: bool = False,
-        simple_io: bool = False,
-        gamma: float = 0.8,
-        max_flow: float = 400,
         **kwargs,
     ) -> None:
         super().__init__(
             pyramid_ranges,
             iters,
-            corr_mode,
-            corr_levels,
-            corr_range,
-            enc_hidden_chs,
-            enc_out_chs,
-            enc_stem_stride,
-            enc_mlp_ratio,
-            enc_depth,
-            dec_net_chs,
-            dec_inp_chs,
-            dec_motion_chs,
-            dec_depth,
-            dec_mlp_ratio,
-            use_upsample_mask,
-            fuse_next1d_weights,
-            simple_io,
-            gamma,
-            max_flow,
             **kwargs,
         )
 

ptlflow/utils/lightning/ptlflow_cli.py

@@ -27,6 +27,7 @@
 
 class PTLFlowCLI(LightningCLI):
     parser_class = LightningArgumentParser
+
     def __init__(
         self,
         model_class: Optional[

tests/ptlflow/models/test_models.py

@@ -122,11 +122,10 @@ def test_forward_fp16() -> None:
             model = model.eval()
             model = model.half()
 
-            s = make_divisible(256, model.output_stride)
             num_images = 2
             if mname in ["videoflow_bof", "videoflow_mof"]:
                 num_images = 3
-            inputs = {"images": torch.rand(1, num_images, 3, s, s)}
+            inputs = {"images": torch.rand(1, num_images, 3, 256, 256)}
 
             if torch.cuda.is_available():
                 model = model.cuda()

tests/ptlflow/test_lightning.py

@@ -0,0 +1,55 @@
+# =============================================================================
+# Copyright 2021 Henrique Morimitsu
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# =============================================================================
+
+import sys
+
+from ptlflow.utils.lightning.ptlflow_cli import PTLFlowCLI
+from ptlflow.utils.lightning.ptlflow_checkpoint_connector import (
+    _PTLFlowCheckpointConnector,
+)
+from ptlflow.utils.lightning.ptlflow_trainer import PTLFlowTrainer
+from ptlflow.utils.registry import RegisteredModel
+
+TEST_MODEL = "raft_small"
+
+
+def test_cli_no_model() -> None:
+    sys.argv = sys.argv[:1]
+    PTLFlowCLI(
+        model_class=None,
+        subclass_mode_model=False,
+        parser_kwargs={"parents": []},
+        run=False,
+        parse_only=False,
+        auto_configure_optimizers=False,
+    )
+
+
+def test_cli() -> None:
+    sys.argv = sys.argv[:1]
+    sys.argv.extend(["--model", "rapidflow"])
+    PTLFlowCLI(
+        model_class=RegisteredModel,
+        subclass_mode_model=True,
+        parser_kwargs={"parents": []},
+        run=False,
+        parse_only=False,
+        auto_configure_optimizers=False,
+    )
+
+
+def test_checkpoint_connector_and_trainer() -> None:
+    _PTLFlowCheckpointConnector(PTLFlowTrainer())