estimate_chunk_size

nabu/cuda/utils.py

@@ -45,6 +45,17 @@ def count_cuda_devices():
 
 
 
+
+# ~ def get_gpus_mem(n_devices):
+    # ~ mem = []
+    # ~ for i in range(n_devices):
+        # ~ ctx = get_cuda_context(device_id=i, cleanup_at_exit=False)
+        # ~ dev = ctx.get_device()
+        # ~ mem.append(dev.total_memory()/1e9)
+        # ~ ctx.pop()
+    # ~ return mem
+
+
 def get_gpu_memory(device_id):
     """
     Return the total memory (in GigaBytes) of a device.

nabu/resources/computations.py

@@ -7,13 +7,6 @@ computations.py: determine computational needs, chunking method to be used, etc.
 from silx.image.tomography import get_next_power
 
 
-def get_chunk_size(nabu_config, dataset, available_mem):
-    """
-    Return the maximum chunk size that can be used with a given amount of memory.
-    """
-    pass
-
-
 def estimate_required_memory(nabu_config, dataset, chunk_size=None):
     """
     Estimate the memory (RAM) needed for a reconstruction.
@@ -27,7 +20,7 @@ def estimate_required_memory(nabu_config, dataset, chunk_size=None):
         Dataset structure.
     """
 
-    Ny, Nx = dataset.radio_dims
+    Nx, Ny = dataset.radio_dims
     Na = dataset.n_angles
     if chunk_size is not None:
         Ny = chunk_size
@@ -83,8 +76,32 @@ def estimate_required_memory(nabu_config, dataset, chunk_size=None):
         Nz_rec = (rec_config["end_z"] - rec_config["start_z"] + 1)//binning_z
         if chunk_size:
             Nz_rec = chunk_size//binning_z
-
         reconstructed_volume_size = Nx_rec * Ny_rec * Nz_rec * 4 # float32
         total_memory_needed += reconstructed_volume_size
 
     return total_memory_needed
+
+
+
+def estimate_chunk_size(available_memory_GB, nabu_config, dataset, min_chunk=50):
+    """
+    Estimate the maximum chunk size given an avaiable amount of memory.
+
+    Parameters
+    -----------
+    available_memory_GB: float
+        available memory in Giga Bytes (GB - not GiB !).
+    nabu_config: dict
+        Dictionary describing the nabu configuration. Usually extracted from
+        the configuration file.
+    dataset: nabu.resources.dataset_analyzer.DatasetAnalyzer
+        Dataset structure.
+    """
+    # Deriving a closed-form formula for estimate_required_memory would be cumbersome
+    # so we use this inelegant way instead
+    step = min_chunk
+    chunk_size = min_chunk
+    while (estimate_required_memory(nabu_config, dataset, chunk_size=chunk_size)/1e9 < available_memory_GB) and chunk_size < dataset.radio_dims[1]:
+        chunk_size += step
+    return chunk_size
+