Atb operator CUDA Out Of Memory error
- Dominant language
- MATLAB
- Stars
- 803
- Forks
- 262
- Avg merge
- 9d 13h
- Merged PRs (30d)
- 8
Description
## Background
I'm simulating a local tomography setup where the region of interest must be covered with `N` scans. I implemented an iterative gradient-based reconstruction using the `Ax` and `Atb` operators to solve the following minimization problem:
$$\mathbf{x}^* = argmin_{\mathbf{x}} \quad \frac{1}{N} \sum_{i=1}^{N}
\lVert y_i - A_i(\mathbf{x})\lVert ^2 \quad \text{s.t.} \quad \mathbf{x^*} \in \mathcal{X}$$
where $\mathcal{X}$ is the feasible set, $y_i$ the observation obtained with the $A_i$ forward operator (i.e. a sample is scanned N times with the forward operator `A_i` corresponding to the local ROI `i`.
Starting from an initial estimate, we iterate a number of iterations `K` updating the estimate at every step. Each update does:
$$\mathbf{x}^{k+1} = \mathbf{x}^{k} - \eta \nabla_{\textbf{x}}\mathcal{L}(x^k)$$
where the gradient is
$$\nabla_{\textbf{x}}\mathcal{L}(x) = \frac{-1}{N} \sum_{i=1}^{N} A^T_i(y_i - A_i(\mathbf{x}))$$
and where $A^T_i$ is the adjoint (back-projection) operator.
## Expected Behavior
Memory usage is stable regardless of the number of iterations used.
## Actual Behavior
Memory grows when using the operator iteratively (heap size grows linearly with the number of iterations while the resident size remains stable according to a quick memray profiling). As the point of Tigre is making it easier to use iterative algorithms, I must be messing up somewhere but have not managed to find where so far. I suspected that defining a new geometry every time the operators are used could be the issue (hence the `ax_geos` and `atb_geos` dictionaries), but that did not change the behavior. It looks like objects are accumulating in memory despite not being used anymore. I have tried using the [`IterativeReconAlg`](https://github.com/CERN/TIGRE/blob/master/Python/tigre/algorithms/iterative_recon_alg.py#L33) class as a reference but I have not been able to find the issue in my code.
Any pointers to where I could be keeping things in memory?
## Code to reproduce the problem (If applicable)
A minimal code example naively implementing this setting that will break in the iteration 707 after 1h18' in a V100 with the error:
```
../Common/CUDA/TIGRE_common.cpp (7): Texture object creation fail
../Common/CUDA/TIGRE_common.cpp (14): CBCT:CUDA:Atb out of memory
```
```python
from functools import partial
import numpy as np
import tigre
from skimage.data import shepp_logan_phantom
from tqdm import tqdm
def pad_around_center(image: np.ndarray, center: tuple):
"""This function pads an image around a given center coordinate so that the image has equal dimensions in every direction from the center"""
squeezed = False
if image.ndim == 3:
image = image[0]
squeezed = True
distance_to_left = center[1]
distance_to_right = image.shape[1] - center[1]
left_padding = int(max(0, distance_to_right - distance_to_left))
right_padding = int(max(0, distance_to_left - distance_to_right))
distance_to_top = center[0]
distance_to_bottom = image.shape[0] - center[0]
top_padding = int(max(0, distance_to_bottom - distance_to_top))
bottom_padding = int(max(0, distance_to_top - distance_to_bottom))
# print(f"Center: {center}")
# print(f"Distances - Left: {distance_to_left}, Right: {distance_to_right}, Top: {distance_to_top}, Bottom: {distance_to_bottom}")
# print(f"Padding - Left: {left_padding}, Right: {right_padding}, Top: {top_padding}, Bottom: {bottom_padding}")
image_padded = np.pad(
image,
((top_padding, bottom_padding), (left_padding, right_padding)),
mode="constant",
constant_values=0,
)
if squeezed:
image_padded = image_padded[np.newaxis, ...]
return image_padded
def crop_padding(image, pad_width):
"""This function crops the padding from an image"""
return image[:, pad_width:-pad_width, pad_width:-pad_width]
def mse(image1, image2):
"""This function calculates the mean squared error between two images"""
return np.mean((image1 - image2) ** 2)
def psnr(image1, image2):
"""This function calculates the peak signal-to-noise ratio between two images"""
mse_value = mse(image1, image2)
max_value = np.max(image1)
return 10 * np.log10(max_value**2 / mse_value)
np.random.seed(42) # Use any fixed number for reproducibility
def _Ax(
i: int,
fov: tuple,
theta: np.ndarray,
centers_of_rotation: list,
x: np.ndarray, ## x is the entire domain, not just a roi
geos = dict
):
# Center the domain around the rotation center
rot_center = centers_of_rotation[i]
if x.ndim == 2:
x = x[np.newaxis, ...]
centered_x = pad_around_center(
x, rot_center
) # Padded to simulate local tomography with the sample moving
# Calculate projection windows considering both coordinate systems
window_half_width = fov[1] // 2
# Max because TIGRE's parallel geometry defines the number of horizontal pixels in the detector as the max of the last two dimensions of the volume
detector_center = np.max((centered_x.shape[-2] // 2, centered_x.shape[-1] // 2))
# ROI in detector coordinates
det_start = int(detector_center - window_half_width)
det_end = int(detector_center + window_half_width)
# Define geometry and calculate projection
if centered_x.shape not in geos:
geo = tigre.geometry(mode="parallel", nVoxel=np.array(centered_x.shape))
geos[centered_x.shape] = geo
else:
geo = geos[centered_x.shape]
y = tigre.Ax(centered_x, geo, theta)
# Crop the projection to the detector
det_y = y[:, :, det_start:det_end]
return det_y
def _Atb(
theta: np.ndarray,
y: np.ndarray,
geos: dict,
padding_factor: float = 1,
):
# y is a sinogram where each row is a projection, hence the first dimension is the number of projections and the second and third dims are the number of detector pixels (h,w)
fov = y.shape[-1]
# Pad sinogram
pad_width = int((fov * padding_factor) // 2)
padded_y = np.pad(y, ((0, 0), (0, 0), (pad_width, pad_width)), mode="edge")
# Backproject
local_domain = (padded_y.shape[-2], padded_y.shape[-1], padded_y.shape[-1])
if local_domain not in geos:
local_geo = tigre.geometry(mode="parallel", nVoxel=np.array(local_domain))
geos[local_domain] = local_geo
else:
local_geo = geos[local_domain]
x_tilda = crop_padding(tigre.Atb(padded_y, local_geo, theta), pad_width)
return x_tilda
def stitch(scans: dict[tuple, np.ndarray], domain: tuple):
Nscans = len(scans)
output = np.zeros((Nscans, 1, domain[0], domain[1]), dtype=np.float32)
for i, (rot_center, roi) in enumerate(scans.items()):
half_roi = int(roi.shape[-2] // 2), int(roi.shape[-1] // 2)
s = (
slice(rot_center[0] - half_roi[0], rot_center[0] + half_roi[0]),
slice(rot_center[1] - half_roi[1], rot_center[1] + half_roi[1]),
)
output[i, 0][s] = roi
nelements = np.sum(output.astype(bool), axis=0)
nelements = np.where(nelements == 0, 1, nelements)
output = np.sum(output, axis=0) / nelements # TODO: This is a naive way of averaging the local scans
return output
def main():
gt = shepp_logan_phantom().astype(np.float32)[None, ...] # TODO: Forced 3D
domain = gt.shape
hr_fov = (100, 100)
spacing = (50, 50)
PADDING = 1 # This is a percentage!
NANGLES = 1000
angles = np.linspace(0, 2 * np.pi, NANGLES)
centers = [(100, 300), (350, 100), (150, 200), (350, 350), (200, 100), (250, 150), (50, 250), (300, 50), (100, 150), (200, 350), (150, 50), (300, 300), (350, 200), (150, 300), (50, 150), (200, 200), (50, 100), (250, 250), (50, 350), (300, 150), (100, 250), (350, 50), (150, 150), (350, 300), (200, 50), (250, 100), (200, 300), (50, 200), (250, 350), (100, 100), (300, 250), (100, 350), (350, 150), (150, 250), (200, 150), (50, 50), (250, 200), (50, 300), (300, 100), (100, 200), (300, 350), (150, 100), (350, 250), (250, 50), (150, 350), (200, 250), (250, 300), (100, 50), (300, 200)]
print(centers)
N = len(centers)
ax_geos = {}
atb_geos = {}
Ax = partial(
_Ax,
fov=hr_fov,
centers_of_rotation=centers,
theta=angles,
geos=ax_geos
)
Atb = partial(
_Atb,
theta=angles,
padding_factor=PADDING,
geos=atb_geos
)
ys = np.array([Ax(i=i, x=gt) for i in range(len(centers))]).astype(np.float32)
x = np.zeros(domain, dtype=np.float32)
# Optimization parameters
learning_rate=5e-6
tolerance=1e-6
max_iterations=1000
lambda_weight=1
K = 10 # Simulates doing the optimization K times with different parameters
for _ in range(K):
pbar = tqdm(range(max_iterations))
for k in pbar:
# print(f"{len(ax_geos)=}")
# print(f"{len(atb_geos)=}")
step_ys = np.array([Ax(i=i, x=x) for i in range(N)]).astype(np.float32)
gradients = {
centers[i]: (Atb(y=(ys[i] - step_ys[i])))
for i in range(N)
}
gradient = stitch(gradients, domain[1:]).astype(np.float32)
grad_fidelity = lambda_weight * gradient
grad = -grad_fidelity
x_new = x - learning_rate * grad
x_new = np.maximum(x_new, 0)
dNorm = np.linalg.norm(x_new - x)
if dNorm < tolerance:
break
# Update progress bar
progress_desc = f"dNorm: {dNorm:.4f}"
if gt is not None:
gt_loss = mse(x_new, gt)
progress_desc += f" | GT loss: {gt_loss:.4f}"
pbar.set_description(progress_desc)
# Update
x = x_new
if __name__ == "__main__":
main()
```
## Specifications
- python version: 3.11
- OS: x86_64 GNU/Linux
- CUDA version: 12.4
- pytigre version: 2.4.0
Contributor guide
Assessment
This issue has not been assessed yet.