Fix ray-tracing GPU OOM: process dirs sequentially on CPU, free GPU memory between dirs

This commit is contained in:
Antoine Jacquin
2026-05-31 18:44:54 +02:00
parent e9c48f9ff8
commit 2ccedbb9e0

View File

@ -324,8 +324,10 @@ def _ray_trace_horizons(dem, rows, cols, res, n_dirs, max_dist, radii_m=None):
padded = xp.pad(dem, max_dist, mode='constant', constant_values=xp.nan)
pos_angles = xp.zeros((n_dirs, n_radii, rows, cols))
neg_angles = xp.zeros((n_dirs, n_radii, rows, cols))
# Process one direction at a time to limit GPU memory.
# Store results as flat CPU arrays — transfer back to GPU at the end.
pos_results = [None] * n_dirs
neg_results = [None] * n_dirs
for d_idx in range(n_dirs):
ddx, ddy = dx_dir[d_idx], dy_dir[d_idx]
@ -373,8 +375,19 @@ def _ray_trace_horizons(dem, rows, cols, res, n_dirs, max_dist, radii_m=None):
if not radii_remaining:
break
pos_angles[d_idx] = running_pos
neg_angles[d_idx] = running_neg
# Store results on CPU, free GPU memory before next direction
pos_results[d_idx] = to_cpu(running_pos)
neg_results[d_idx] = to_cpu(running_neg)
del running_pos, running_neg
gpu_cleanup()
# Free the large padded array
del padded
gpu_cleanup()
# Reassemble into final arrays (on CPU to avoid GPU memory pressure)
pos_angles = np.array(pos_results)
neg_angles = np.array(neg_results)
return pos_angles, neg_angles
@ -511,14 +524,14 @@ def generate_svf(dem_file, basename, vis_dir, resolution, shared=None):
pos_angles, neg_angles = _ray_trace_horizons(dem, rows, cols, res, n_dirs, max_dist, radii_m)
# SVF per radius: mean of cos²(horizon) across directions
svf_combined = xp.zeros_like(dem)
# pos/neg are now numpy arrays (CPU) — combine on CPU
svf_combined = np.zeros((rows, cols), dtype=np.float32)
for r_idx in range(len(radii_m)):
horizon = xp.maximum(pos_angles[:, r_idx], neg_angles[:, r_idx])
svf_r = xp.mean(xp.cos(horizon) ** 2, axis=0)
horizon = np.maximum(pos_angles[:, r_idx], neg_angles[:, r_idx])
svf_r = np.mean(np.cos(horizon) ** 2, axis=0)
svf_combined += svf_r * radius_weights[r_idx]
svf_np = to_cpu(svf_combined).astype(np.float32)
svf_np = svf_combined
svf_np[nan_mask] = np.nan
_save_tif(output, svf_np, transform, crs)
logger.info(f" ✓ SVF terminé ({time.time()-t0:.1f}s){' [GPU]' if _gpu_mod.HAS_GPU else ''}")
@ -557,9 +570,9 @@ def generate_openness(dem_file, basename, vis_dir, resolution, positive=True, sh
else:
angles = neg_angles
# Mean across directions and radii (equal weight)
openness = xp.mean(angles, axis=(0, 1))
openness_result = to_cpu(xp.degrees(openness)).astype(np.float32)
# Mean across directions and radii (equal weight) — on CPU now
openness = np.mean(angles, axis=(0, 1))
openness_result = np.degrees(openness).astype(np.float32)
openness_result[nan_mask] = np.nan
# Std normalization for cross-tile comparability
@ -1088,8 +1101,8 @@ def generate_aniso_open(dem_file, basename, vis_dir, resolution, shared=None):
weight_total = np.sum(weights)
n_radii = len(radii_m)
pos_combined = xp.zeros_like(dem)
neg_combined = xp.zeros_like(dem)
pos_combined = np.zeros((rows, cols), dtype=np.float64)
neg_combined = np.zeros((rows, cols), dtype=np.float64)
for r_idx in range(n_radii):
for d_idx in range(n_dirs):
@ -1097,7 +1110,7 @@ def generate_aniso_open(dem_file, basename, vis_dir, resolution, shared=None):
pos_combined += pos_angles[d_idx, r_idx] * w / (n_radii * weight_total)
neg_combined += neg_angles[d_idx, r_idx] * w / (n_radii * weight_total)
aniso_result = to_cpu(xp.degrees(pos_combined - neg_combined)).astype(np.float32)
aniso_result = np.degrees(pos_combined - neg_combined).astype(np.float32)
aniso_result[nan_mask] = np.nan
# Std normalization for cross-tile comparability