[BUG] Improve memory management by adding garbage collection and cleanup tasks in evaluators

Author: Leguark

gempy_engine/core/backend_tensor.py

@@ -46,11 +46,14 @@ def get_backend_string(cls) -> str:
                 return "CPU"
 
     @classmethod
-    def change_backend_gempy(cls, engine_backend: AvailableBackends, use_gpu: bool = False, dtype: Optional[str] = None):
-        cls._change_backend(engine_backend, use_pykeops=PYKEOPS, use_gpu=use_gpu, dtype=dtype)
+    def change_backend_gempy(cls, engine_backend: AvailableBackends, use_gpu: bool = False,
+                             dtype: Optional[str] = None, grads:bool = False):
+        cls._change_backend(engine_backend, use_pykeops=PYKEOPS, use_gpu=use_gpu, dtype=dtype,
+                            grads=grads)
 
     @classmethod
-    def _change_backend(cls, engine_backend: AvailableBackends, use_pykeops: bool = False, use_gpu: bool = True, dtype: Optional[str] = None):
+    def _change_backend(cls, engine_backend: AvailableBackends, use_pykeops: bool = False,
+                        use_gpu: bool = True, dtype: Optional[str] = None, grads:bool = False):
         cls.dtype = DEFAULT_TENSOR_DTYPE if dtype is None else dtype
         cls.dtype_obj = cls.dtype
         match engine_backend:
@@ -99,6 +102,21 @@ def _change_backend(cls, engine_backend: AvailableBackends, use_pykeops: bool =
                 cls.dtype_obj = pytorch_copy.float32 if cls.dtype == "float32" else pytorch_copy.float64
                 cls.tensor_types = pytorch_copy.Tensor
 
+                torch.set_num_threads(torch.get_num_threads())  # Use all available threads
+                cls.COMPUTE_GRADS = grads  # Store the grads setting
+                if grads is False:
+                    cls._torch_no_grad_context = torch.no_grad()
+                    cls._torch_no_grad_context.__enter__()
+                else:
+                    # If there was a previous context, exit it first
+                    if hasattr(cls, '_torch_no_grad_context') and cls._torch_no_grad_context is not None:
+                        try:
+                            cls._torch_no_grad_context.__exit__(None, None, None)
+                        except:
+                            pass  # Context might already be exited
+                    cls._torch_no_grad_context = None
+                    torch.set_grad_enabled(True)
+                    
                 cls.use_pykeops = use_pykeops  # TODO: Make this compatible with pykeops
                 if (use_pykeops):
                     import pykeops

gempy_engine/modules/evaluator/generic_evaluator.py

@@ -1,4 +1,5 @@
 import numpy as np
+import gc
 from typing import Optional
 
 from gempy_engine.core.backend_tensor import BackendTensor
@@ -57,6 +58,10 @@ def generic_evaluator(
             if gz_field is not None:
                 gz_field[slice_array] = gz_chunk  # type: ignore
 
+    # Force garbage collection every few chunks to prevent memory buildup
+    if (i + 1) % 5 == 0 or i == n_chunks - 1:
+        gc.collect()
+        
     if n_chunks > 5:
         print(f"Chunking done: {n_chunks} chunks")
 
@@ -75,7 +80,9 @@ def _eval_on(
     try:
         scalar_field = (eval_kernel.T @ weights).reshape(-1)
     except ValueError:
-        pass
+        scalar_field = None
+    
+    del eval_kernel
 
     gx_field: Optional[np.ndarray] = None
     gy_field: Optional[np.ndarray] = None
@@ -85,17 +92,21 @@ def _eval_on(
         eval_gx = yield_evaluation_grad_kernel(
             solver_input, options.kernel_options, axis=0, slice_array=slice_array
         )
+        gx_field = (eval_gx.T @ weights).reshape(-1)  # Use BEFORE deleting
+        del eval_gx  # Clean up immediately after use
+        
         eval_gy = yield_evaluation_grad_kernel(
             solver_input, options.kernel_options, axis=1, slice_array=slice_array
         )
-        gx_field = (eval_gx.T @ weights).reshape(-1)
-        gy_field = (eval_gy.T @ weights).reshape(-1)
+        gy_field = (eval_gy.T @ weights).reshape(-1)  # Use BEFORE deleting
+        del eval_gy  # Clean up immediately after use
 
         if options.number_dimensions == 3:
             eval_gz = yield_evaluation_grad_kernel(
                 solver_input, options.kernel_options, axis=2, slice_array=slice_array
             )
-            gz_field = (eval_gz.T @ weights).reshape(-1)
+            gz_field = (eval_gz.T @ weights).reshape(-1)  # Use BEFORE deleting
+            del eval_gz  # Clean up immediately after use
         elif options.number_dimensions != 2:
             raise ValueError("`number_dimensions` must be 2 or 3")