Added option to clear cuda cache when vectorization cache is cleared. (…

…#249)

* Added option to clear cuda cache when vectorization cache is cleared.

* Changed PGO synthetic script to have the option to empty cuda cache.

examples/configs/pose_graph/pose_graph_synthetic.yaml

@@ -23,6 +23,7 @@ inner_optim:
   ratio_known_poses: 0.1
   reg_w: 1e-3
   vectorize: true
+  empty_cuda_cache: false
 
 outer_optim:
   lr: 0.1

examples/pose_graph/pose_graph_synthetic.py

@@ -197,10 +197,11 @@ def run(
         linearization_cls=LINEARIZATION_MODE[cast(str, cfg.inner_optim.solver)],
         linear_solver_cls=LINEAR_SOLVER_MODE[cast(str, cfg.inner_optim.solver)],
         vectorize=cfg.inner_optim.vectorize,
+        empty_cuda_cache=cfg.inner_optim.empty_cuda_cache,
     )
 
     # Set up Theseus layer
-    theseus_optim = th.TheseusLayer(optimizer, vectorize=cfg.inner_optim.vectorize)
+    theseus_optim = th.TheseusLayer(optimizer)
     theseus_optim.to(device=device)
 
     # Outer optimization loop

theseus/core/vectorizer.py

@@ -84,7 +84,7 @@ def __repr__(self) -> str:
 class Vectorize:
     _SHARED_TOKEN = "__shared__"
 
-    def __init__(self, objective: Objective):
+    def __init__(self, objective: Objective, empty_cuda_cache: bool = False):
         # Each cost function is assigned a wrapper. The wrapper will hold the error
         # and jacobian after vectorization.
         self._cost_fn_wrappers: List[_CostFunctionWrapper] = []
@@ -127,6 +127,7 @@ def __init__(self, objective: Objective):
         )
 
         self._objective = objective
+        self._empty_cuda_cache = empty_cuda_cache
 
     # Returns a dictionary which maps every schema to information about its shared
     # variables.
@@ -288,6 +289,8 @@ def _clear_wrapper_caches(self):
             for cf in cost_fn_wrappers:
                 cf._cached_error = None
                 cf._cached_jacobians = None
+        if self._empty_cuda_cache:
+            torch.cuda.empty_cache()
 
     # This could be a static method, but writing like this makes some unit tests
     # easier

theseus/optimizer/linear/linear_optimizer.py

@@ -34,7 +34,7 @@ def __init__(
         linear_solver_kwargs: Optional[Dict[str, Any]] = None,
         **kwargs,
     ):
-        super().__init__(objective, vectorize=vectorize)
+        super().__init__(objective, vectorize=vectorize, **kwargs)
         linearization_kwargs = linearization_kwargs or {}
         linear_solver_kwargs = linear_solver_kwargs or {}
         self.linear_solver = linear_solver_cls(

theseus/optimizer/nonlinear/gauss_newton.py

@@ -27,6 +27,7 @@ def __init__(
         rel_err_tolerance: float = 1e-8,
         max_iterations: int = 20,
         step_size: float = 1.0,
+        **kwargs
     ):
         super().__init__(
             objective,
@@ -39,6 +40,7 @@ def __init__(
             rel_err_tolerance=rel_err_tolerance,
             max_iterations=max_iterations,
             step_size=step_size,
+            **kwargs
         )
 
     def compute_delta(self, **kwargs) -> torch.Tensor:

theseus/optimizer/nonlinear/levenberg_marquardt.py

@@ -39,6 +39,7 @@ def __init__(
         rel_err_tolerance: float = 1e-8,
         max_iterations: int = 20,
         step_size: float = 1.0,
+        **kwargs,
     ):
         super().__init__(
             objective,
@@ -51,6 +52,7 @@ def __init__(
             rel_err_tolerance=rel_err_tolerance,
             max_iterations=max_iterations,
             step_size=step_size,
+            **kwargs,
         )
         self._allows_ellipsoidal = _check_ellipsoidal_damping_cls(self.linear_solver)
 

theseus/optimizer/nonlinear/nonlinear_least_squares.py

@@ -44,6 +44,7 @@ def __init__(
             rel_err_tolerance=rel_err_tolerance,
             max_iterations=max_iterations,
             step_size=step_size,
+            **kwargs,
         )
 
     @abc.abstractmethod

theseus/optimizer/nonlinear/nonlinear_optimizer.py

@@ -75,7 +75,7 @@ def __init__(
         step_size: float = 1.0,
         **kwargs,
     ):
-        super().__init__(objective, vectorize=vectorize)
+        super().__init__(objective, vectorize=vectorize, **kwargs)
         linear_solver_kwargs = linear_solver_kwargs or {}
         self.linear_solver = linear_solver_cls(
             objective,

theseus/optimizer/optimizer.py

@@ -25,7 +25,9 @@ class Optimizer(abc.ABC):
     def __init__(self, objective: Objective, *args, vectorize: bool = True, **kwargs):
         self.objective = objective
         if vectorize:
-            Vectorize(self.objective)
+            Vectorize(
+                self.objective, empty_cuda_cache=kwargs.get("empty_cuda_cache", False)
+            )
         self._objectives_version = objective.current_version
 
     @abc.abstractmethod

theseus/theseus_layer.py

@@ -25,11 +25,16 @@
 
 
 class TheseusLayer(nn.Module):
-    def __init__(self, optimizer: Optimizer, vectorize: bool = True):
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        vectorize: bool = True,
+        empty_cuda_cache: bool = False,
+    ):
         super().__init__()
         self.objective = optimizer.objective
         if vectorize and not self.objective.vectorized:
-            Vectorize(self.objective)
+            Vectorize(self.objective, empty_cuda_cache=empty_cuda_cache)
         self.optimizer = optimizer
         self._objectives_version = optimizer.objective.current_version
         self._dlm_bwd_objective = None