Commit abcd5223 authored by Pierre Paleo's avatar Pierre Paleo
Browse files

estimate_chunk_size

parent 4dc200d8
......@@ -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.
......
......@@ -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
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment