From 208c5c48e23a5aa3ead9c0375382683eb0b54f6d Mon Sep 17 00:00:00 2001 From: dkazanc Date: Wed, 3 Jun 2026 15:40:17 +0100 Subject: [PATCH] adding memory estimator and tests for osem --- .../backends/httomolibgpu/httomolibgpu.yaml | 9 ++ .../supporting_funcs/recon/algorithm.py | 55 ++++++++++- tests/test_httomolibgpu.py | 98 +++++++++++++++++++ 3 files changed, 161 insertions(+), 1 deletion(-) diff --git a/httomo_backends/methods_database/packages/backends/httomolibgpu/httomolibgpu.yaml b/httomo_backends/methods_database/packages/backends/httomolibgpu/httomolibgpu.yaml index 2d8a7125..1dbe824d 100644 --- a/httomo_backends/methods_database/packages/backends/httomolibgpu/httomolibgpu.yaml +++ b/httomo_backends/methods_database/packages/backends/httomolibgpu/httomolibgpu.yaml @@ -236,6 +236,15 @@ recon: memory_gpu: multiplier: None method: module + OSEM3d_tomobar: + pattern: sinogram + output_dims_change: True + implementation: gpu_cupy + save_result_default: True + padding: True + memory_gpu: + multiplier: None + method: module rotation: find_center_vo: pattern: projection diff --git a/httomo_backends/methods_database/packages/backends/httomolibgpu/supporting_funcs/recon/algorithm.py b/httomo_backends/methods_database/packages/backends/httomolibgpu/supporting_funcs/recon/algorithm.py index 8879ed97..6d7afb48 100644 --- a/httomo_backends/methods_database/packages/backends/httomolibgpu/supporting_funcs/recon/algorithm.py +++ b/httomo_backends/methods_database/packages/backends/httomolibgpu/supporting_funcs/recon/algorithm.py @@ -24,7 +24,7 @@ from typing import Tuple import numpy as np from httomo_backends.cufft import CufftType, cufft_estimate_1d, cufft_estimate_2d -from httomolibgpu.recon.algorithm import ADMM3d_tomobar, LPRec3d_tomobar +from httomolibgpu.recon.algorithm import LPRec3d_tomobar from tomobar.supp.memory_estimator_helpers import DeviceMemStack __all__ = [ @@ -35,6 +35,7 @@ "_calc_memory_bytes_CGLS3d_tomobar", "_calc_memory_bytes_FISTA3d_tomobar", "_calc_memory_bytes_ADMM3d_tomobar", + "_calc_memory_bytes_OSEM3d_tomobar", "_calc_output_dim_FBP2d_astra", "_calc_output_dim_FBP3d_tomobar", "_calc_output_dim_LPRec3d_tomobar", @@ -42,8 +43,10 @@ "_calc_output_dim_CGLS3d_tomobar", "_calc_output_dim_FISTA3d_tomobar", "_calc_output_dim_ADMM3d_tomobar", + "_calc_output_dim_OSEM3d_tomobar", "_calc_padding_FISTA3d_tomobar", "_calc_padding_ADMM3d_tomobar", + "_calc_padding_OSEM3d_tomobar", ] @@ -54,6 +57,9 @@ def _calc_padding_FISTA3d_tomobar(**kwargs) -> Tuple[int, int]: def _calc_padding_ADMM3d_tomobar(**kwargs) -> Tuple[int, int]: return (5, 5) +def _calc_padding_OSEM3d_tomobar(**kwargs) -> Tuple[int, int]: + return (5, 5) + def __calc_output_dim_recon(non_slice_dims_shape, **kwargs): """Function to calculate output dimensions for all reconstructors. @@ -97,6 +103,9 @@ def _calc_output_dim_ADMM3d_tomobar(non_slice_dims_shape, **kwargs): return __calc_output_dim_recon(non_slice_dims_shape, **kwargs) +def _calc_output_dim_OSEM3d_tomobar(non_slice_dims_shape, **kwargs): + return __calc_output_dim_recon(non_slice_dims_shape, **kwargs) + def _calc_memory_bytes_FBP3d_tomobar( non_slice_dims_shape: Tuple[int, int], dtype: np.dtype, @@ -627,6 +636,50 @@ def _calc_memory_bytes_FISTA3d_tomobar( return (tot_memory_bytes, 0) +def _calc_memory_bytes_OSEM3d_tomobar( + non_slice_dims_shape: Tuple[int, int], + dtype: np.dtype, + **kwargs, +) -> Tuple[int, int]: + detector_pad = 0 + if "detector_pad" in kwargs: + detector_pad = kwargs["detector_pad"] + if detector_pad is True: + detector_pad = __estimate_detectorHoriz_padding(non_slice_dims_shape[1]) + elif detector_pad is False: + detector_pad = 0 + + anglesnum = non_slice_dims_shape[0] + DetectorsLengthH_padded = non_slice_dims_shape[1] + 2 * detector_pad + + # calculate the output shape + output_dims = _calc_output_dim_OSEM3d_tomobar(non_slice_dims_shape, **kwargs) + recon_data_size_original = ( + np.prod(output_dims) * dtype.itemsize + ) # recon user-defined size + + in_data_siz_pad = (anglesnum * DetectorsLengthH_padded) * dtype.itemsize + output_dims_larger_grid = (DetectorsLengthH_padded, DetectorsLengthH_padded) + + out_data_size = np.prod(output_dims_larger_grid) * dtype.itemsize + normalisation = out_data_size + Ax = in_data_siz_pad + ratio = in_data_siz_pad + backproj = out_data_size + + mlem_part = ( + recon_data_size_original + + normalisation + + Ax + + ratio + + backproj + + out_data_size + ) + regul_part = 8 * np.prod(output_dims_larger_grid) * dtype.itemsize + + tot_memory_bytes = int(mlem_part + regul_part) + return (tot_memory_bytes, 0) + def _calc_memory_bytes_ADMM3d_tomobar( non_slice_dims_shape: Tuple[int, int], dtype: np.dtype, diff --git a/tests/test_httomolibgpu.py b/tests/test_httomolibgpu.py index b428bd52..8e1395a9 100644 --- a/tests/test_httomolibgpu.py +++ b/tests/test_httomolibgpu.py @@ -34,6 +34,7 @@ CGLS3d_tomobar, FISTA3d_tomobar, ADMM3d_tomobar, + OSEM3d_tomobar, ) from httomolibgpu.misc.rescale import rescale_to_int @@ -1010,6 +1011,103 @@ def test_recon_FISTA3d_tomobar_nonOS_memoryhook( assert percents_relative_maxmem <= 100 +@pytest.mark.cupy +@pytest.mark.parametrize("slices", [3, 5]) +@pytest.mark.parametrize("recon_size_it", [2560]) +@pytest.mark.parametrize("padding", [0, 100, 200]) +def test_recon_MLEM3d_tomobar_nonOS_memoryhook( + slices, recon_size_it, padding, ensure_clean_memory +): + angles_total = 901 + detX_size = 2560 + data = cp.random.random_sample((angles_total, slices, detX_size), dtype=np.float32) + kwargs = {} + kwargs["recon_size"] = recon_size_it + + hook = MaxMemoryHook() + with hook: + recon_data = OSEM3d_tomobar( + cp.copy(data), + np.linspace(0.0 * np.pi / 180.0, 180.0 * np.pi / 180.0, data.shape[0]), + center=1200, + recon_size=recon_size_it, + iterations=1, + subsets_number=1, + regularisation_iterations=2, + nonnegativity=True, + detector_pad=padding, + ) + + # make sure estimator function is within range (80% min, 100% max) + max_mem = ( + hook.max_mem + ) # the amount of memory in bytes needed for the method according to memoryhook + + # now we estimate how much of the total memory required for this data + estimated_memory_bytes, subtract_bytes = _calc_memory_bytes_OSEM3d_tomobar( + (angles_total, detX_size), dtype=np.float32(), **kwargs + ) + estimated_memory_mb = round(slices * estimated_memory_bytes / (1024**2), 2) + max_mem -= subtract_bytes + max_mem_mb = round(max_mem / (1024**2), 2) + + # now we compare both memory estimations + difference_mb = abs(estimated_memory_mb - max_mem_mb) + percents_relative_maxmem = round((difference_mb / max_mem_mb) * 100) + # the estimated_memory_mb should be LARGER or EQUAL to max_mem_mb + # the resulting percent value should not deviate from max_mem on more than 20% + assert estimated_memory_mb >= max_mem_mb + assert percents_relative_maxmem <= 100 + +@pytest.mark.cupy +@pytest.mark.parametrize("slices", [3, 5]) +@pytest.mark.parametrize("recon_size_it", [2560]) +@pytest.mark.parametrize("padding", [0, 100, 200]) +def test_recon_OSEM3d_tomobar_memoryhook( + slices, recon_size_it, padding, ensure_clean_memory +): + angles_total = 901 + detX_size = 2560 + data = cp.random.random_sample((angles_total, slices, detX_size), dtype=np.float32) + kwargs = {} + kwargs["recon_size"] = recon_size_it + + hook = MaxMemoryHook() + with hook: + recon_data = OSEM3d_tomobar( + cp.copy(data), + np.linspace(0.0 * np.pi / 180.0, 180.0 * np.pi / 180.0, data.shape[0]), + center=1200, + recon_size=recon_size_it, + iterations=1, + subsets_number=12, + regularisation_iterations=2, + nonnegativity=True, + detector_pad=padding, + ) + + # make sure estimator function is within range (80% min, 100% max) + max_mem = ( + hook.max_mem + ) # the amount of memory in bytes needed for the method according to memoryhook + + # now we estimate how much of the total memory required for this data + estimated_memory_bytes, subtract_bytes = _calc_memory_bytes_OSEM3d_tomobar( + (angles_total, detX_size), dtype=np.float32(), **kwargs + ) + estimated_memory_mb = round(slices * estimated_memory_bytes / (1024**2), 2) + max_mem -= subtract_bytes + max_mem_mb = round(max_mem / (1024**2), 2) + + # now we compare both memory estimations + difference_mb = abs(estimated_memory_mb - max_mem_mb) + percents_relative_maxmem = round((difference_mb / max_mem_mb) * 100) + # the estimated_memory_mb should be LARGER or EQUAL to max_mem_mb + # the resulting percent value should not deviate from max_mem on more than 20% + assert estimated_memory_mb >= max_mem_mb + assert percents_relative_maxmem <= 100 + + @pytest.mark.cupy @pytest.mark.parametrize("slices", [3, 5]) @pytest.mark.parametrize("recon_size_it", [2560])