Rename BackwardMode.FULL --> UNROLL and simplify backward mode config (…

…#305) * Added code so that backward mode can be resolved from string. * Replaced BackwardMode.FULL for BackwardModel.UNROLL.
facebookresearch · Sep 27, 2022 · a18eec7 · a18eec7
1 parent 26a616b
commit a18eec7
Show file tree

Hide file tree

Showing 11 changed files with 51 additions and 46 deletions.
diff --git a/README.md b/README.md
@@ -141,7 +141,7 @@ outer_optimizer = torch.optim.RMSprop([phi], lr=0.001)
 for epoch in range(10):
     solution, info = layer.forward(
         input_tensors={"x": phi.clone(), "v": torch.ones(1, 1)},
-        optimizer_kwargs={"backward_mode": th.BackwardMode.IMPLICIT})
+        optimizer_kwargs={"backward_mode": "implicit"})
     outer_loss = torch.nn.functional.mse_loss(solution["v"], v_true)
     outer_loss.backward()
     outer_optimizer.step()

diff --git a/examples/backward_modes.py b/examples/backward_modes.py
@@ -82,7 +82,7 @@ def quad_error_fn(optim_vars, aux_vars):
     optimizer_kwargs={
         "track_best_solution": True,
         "verbose": False,
-        "backward_mode": th.BackwardMode.FULL,
+        "backward_mode": "unroll",
     },
 )
 
@@ -103,7 +103,7 @@ def quad_error_fn(optim_vars, aux_vars):
     optimizer_kwargs={
         "track_best_solution": True,
         "verbose": False,
-        "backward_mode": th.BackwardMode.IMPLICIT,
+        "backward_mode": "implicit",
     },
 )
 
@@ -117,7 +117,7 @@ def quad_error_fn(optim_vars, aux_vars):
     optimizer_kwargs={
         "track_best_solution": True,
         "verbose": False,
-        "backward_mode": th.BackwardMode.TRUNCATED,
+        "backward_mode": "truncated",
         "backward_num_iterations": 5,
     },
 )
@@ -134,7 +134,7 @@ def quad_error_fn(optim_vars, aux_vars):
     optimizer_kwargs={
         "track_best_solution": True,
         "verbose": False,
-        "backward_mode": th.BackwardMode.DLM,
+        "backward_mode": "dlm",
         "dlm_epsilon": 1e-3,
     },
 )
@@ -175,7 +175,7 @@ def fit_x(data_x_np):
         optimizer_kwargs={
             "track_best_solution": True,
             "verbose": False,
-            "backward_mode": th.BackwardMode.FULL,
+            "backward_mode": "unroll",
         },
     )
     times["fwd"].append(time.time() - start)
@@ -191,7 +191,7 @@ def fit_x(data_x_np):
         optimizer_kwargs={
             "track_best_solution": True,
             "verbose": False,
-            "backward_mode": th.BackwardMode.IMPLICIT,
+            "backward_mode": "implicit",
         },
     )
     start = time.time()
@@ -205,7 +205,7 @@ def fit_x(data_x_np):
         optimizer_kwargs={
             "track_best_solution": True,
             "verbose": False,
-            "backward_mode": th.BackwardMode.TRUNCATED,
+            "backward_mode": "truncated",
             "backward_num_iterations": 5,
         },
     )
@@ -220,7 +220,7 @@ def fit_x(data_x_np):
         optimizer_kwargs={
             "track_best_solution": True,
             "verbose": False,
-            "backward_mode": th.BackwardMode.DLM,
+            "backward_mode": "dlm",
             "dlm_epsilon": 1e-3,
         },
     )

diff --git a/examples/bundle_adjustment.py b/examples/bundle_adjustment.py
@@ -17,13 +17,6 @@
 import theseus as th
 import theseus.utils.examples as theg
 
-BACKWARD_MODE = {
-    "implicit": th.BackwardMode.IMPLICIT,
-    "full": th.BackwardMode.FULL,
-    "truncated": th.BackwardMode.TRUNCATED,
-}
-
-
 # Logger
 log = logging.getLogger(__name__)
 
@@ -211,7 +204,7 @@ def run(cfg: omegaconf.OmegaConf, results_path: pathlib.Path):
             optimizer_kwargs={
                 "verbose": cfg.inner_optim.verbose,
                 "track_err_history": cfg.inner_optim.track_err_history,
-                "backward_mode": BACKWARD_MODE[cfg.inner_optim.backward_mode],
+                "backward_mode": cfg.inner_optim.backward_mode,
                 "__keep_final_step_size__": cfg.inner_optim.keep_step_size,
             },
         )

diff --git a/examples/homography_estimation.py b/examples/homography_estimation.py
@@ -32,11 +32,6 @@
 FONT_SZ = 0.5
 FONT_PT = (5, 15)
 
-BACKWARD_MODE = {
-    "implicit": th.BackwardMode.IMPLICIT,
-    "full": th.BackwardMode.FULL,
-    "truncated": th.BackwardMode.TRUNCATED,
-}
 
 # Logger
 logger = logging.getLogger(__name__)
@@ -406,7 +401,7 @@ def run(
                     "verbose": verbose,
                     "track_err_history": True,
                     "track_state_history": True,
-                    "backward_mode": BACKWARD_MODE["implicit"],
+                    "backward_mode": "implicit",
                 },
             )
             end_event.record()

diff --git a/examples/pose_graph/pose_graph_synthetic.py b/examples/pose_graph/pose_graph_synthetic.py
@@ -24,12 +24,6 @@
 from theseus.optimizer.linear import LinearSolver
 from theseus.optimizer.linearization import Linearization
 
-BACKWARD_MODE = {
-    "implicit": th.BackwardMode.IMPLICIT,
-    "full": th.BackwardMode.FULL,
-    "truncated": th.BackwardMode.TRUNCATED,
-}
-
 LINEARIZATION_MODE: Dict[str, Type[Linearization]] = {
     "sparse": th.SparseLinearization,
     "dense": th.DenseLinearization,
@@ -98,12 +92,6 @@ def run(
     dtype = torch.float64
     pr = cProfile.Profile()
 
-    BACKWARD_MODE = {
-        "implicit": th.BackwardMode.IMPLICIT,
-        "full": th.BackwardMode.FULL,
-        "truncated": th.BackwardMode.TRUNCATED,
-    }
-
     LINEARIZATION_MODE: Dict[str, Type[Linearization]] = {
         "sparse": th.SparseLinearization,
         "dense": th.DenseLinearization,
@@ -232,7 +220,7 @@ def run_batch(batch_idx: int):
             optimizer_kwargs={
                 "verbose": cfg.inner_optim.verbose,
                 "track_err_history": cfg.inner_optim.track_err_history,
-                "backward_mode": BACKWARD_MODE[cfg.inner_optim.backward_mode],
+                "backward_mode": cfg.inner_optim.backward_mode,
                 "__keep_final_step_size__": cfg.inner_optim.keep_step_size,
             },
         )

diff --git a/examples/simple_example.py b/examples/simple_example.py
@@ -49,7 +49,7 @@ def error_fn(optim_vars, aux_vars):  # returns y - v * exp(x)
 for epoch in range(20):
     solution, info = layer.forward(
         input_tensors={"x": phi.clone(), "v": torch.ones(1, 1)},
-        optimizer_kwargs={"backward_mode": th.BackwardMode.IMPLICIT},
+        optimizer_kwargs={"backward_mode": "implicit"},
     )
     outer_loss = torch.nn.functional.mse_loss(solution["v"], v_true)
     outer_loss.backward()

diff --git a/theseus/optimizer/nonlinear/nonlinear_optimizer.py b/theseus/optimizer/nonlinear/nonlinear_optimizer.py
@@ -8,7 +8,7 @@
 import warnings
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Callable, Dict, NoReturn, Optional, Type
+from typing import Any, Callable, Dict, NoReturn, Optional, Type, Union
 
 import numpy as np
 import torch
@@ -53,10 +53,35 @@ class NonlinearOptimizerInfo(OptimizerInfo):
 
 
 class BackwardMode(Enum):
-    FULL = 0
+    UNROLL = 0
     IMPLICIT = 1
     TRUNCATED = 2
     DLM = 3
+    FULL = -1
+
+    @staticmethod
+    def resolve(key: Union[str, "BackwardMode"]) -> "BackwardMode":
+        if isinstance(key, BackwardMode):
+            if key == BackwardMode.FULL:
+                warnings.warn(
+                    "BackwardMode.FULL is deprecated and will be "
+                    "replaced by BackwardMode.UNROLL in future versions.",
+                    DeprecationWarning,
+                )
+                return BackwardMode.UNROLL
+            return key
+
+        if not isinstance(key, str):
+            raise ValueError("Backward mode must be th.BackwardMode or string.")
+
+        try:
+            backward_mode = BackwardMode[key.upper()]
+        except KeyError:
+            raise ValueError(
+                f"Unrecognized backward mode f{key}."
+                f"Valid choices are full, implicit, truncated, dlm."
+            )
+        return backward_mode
 
 
 EndIterCallbackType = Callable[
@@ -351,10 +376,11 @@ def _optimize_impl(
         track_err_history: bool = False,
         track_state_history: bool = False,
         verbose: bool = False,
-        backward_mode: BackwardMode = BackwardMode.FULL,
+        backward_mode: Union[str, BackwardMode] = BackwardMode.UNROLL,
         end_iter_callback: Optional[EndIterCallbackType] = None,
         **kwargs,
     ) -> OptimizerInfo:
+        backward_mode = BackwardMode.resolve(backward_mode)
         with torch.no_grad():
             info = self._init_info(
                 track_best_solution, track_err_history, track_state_history
@@ -366,7 +392,7 @@ def _optimize_impl(
                 f"Error: {info.last_err.mean().item()}"
             )
 
-        if backward_mode in [BackwardMode.FULL, BackwardMode.DLM]:
+        if backward_mode in [BackwardMode.UNROLL, BackwardMode.DLM]:
             self._optimize_loop(
                 start_iter=0,
                 num_iter=self.params.max_iterations,

diff --git a/theseus/optimizer/nonlinear/tests/test_backwards.py b/theseus/optimizer/nonlinear/tests/test_backwards.py
@@ -85,7 +85,7 @@ def fit_x(data_x_np):
         optimizer_kwargs={
             "track_best_solution": True,
             "verbose": False,
-            "backward_mode": th.BackwardMode.FULL,
+            "backward_mode": "unroll",
         },
     )
     da_dx_full = torch.autograd.grad(updated_inputs["a"], data_x, retain_graph=True)[
@@ -111,7 +111,7 @@ def fit_x(data_x_np):
         optimizer_kwargs={
             "track_best_solution": True,
             "verbose": False,
-            "backward_mode": th.BackwardMode.TRUNCATED,
+            "backward_mode": "TRUNCATED",
             "backward_num_iterations": 5,
         },
     )

diff --git a/theseus/tests/test_dlm_perturbation.py b/theseus/tests/test_dlm_perturbation.py
@@ -97,7 +97,7 @@ def test_backward_pass_se3_runs():
             out, _ = layer.forward(
                 {"target": target_data},
                 optimizer_kwargs={
-                    "backward_mode": th.BackwardMode.DLM,
+                    "backward_mode": "dlm",
                     "verbose": False,
                 },
             )

diff --git a/theseus/theseus_layer.py b/theseus/theseus_layer.py
@@ -51,7 +51,10 @@ def forward(
                 "currently not supported."
             )
         optimizer_kwargs = optimizer_kwargs or {}
-        backward_mode = optimizer_kwargs.get("backward_mode", None)
+        # Defaults to "unroll" to avoid error, we only care to see if it's not dlm.
+        backward_mode = BackwardMode.resolve(
+            optimizer_kwargs.get("backward_mode", "unroll")
+        )
         if backward_mode == BackwardMode.DLM:
             dlm_epsilon = optimizer_kwargs.get(
                 TheseusLayerDLMForward._DLM_EPSILON_STR, 1e-2

diff --git a/theseus/utils/examples/tactile_pose_estimation/trainer.py b/theseus/utils/examples/tactile_pose_estimation/trainer.py
@@ -192,7 +192,7 @@ def _resolve_backward_mode(self, epoch: int) -> th.BackwardMode:
             logger.info("Forcing IMPLICIT backward mode.")
             return th.BackwardMode.IMPLICIT
         else:
-            return getattr(th.BackwardMode, self.cfg.inner_optim.backward_mode)
+            return self.cfg.inner_optim.backward_mode
 
     def compute_loss(
         self, epoch: int, update: bool = True