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) padded = xp.pad(dem, max_dist, mode='constant', constant_values=xp.nan)
pos_angles = xp.zeros((n_dirs, n_radii, rows, cols)) # Process one direction at a time to limit GPU memory.
neg_angles = xp.zeros((n_dirs, n_radii, rows, cols)) # 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): for d_idx in range(n_dirs):
ddx, ddy = dx_dir[d_idx], dy_dir[d_idx] 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: if not radii_remaining:
break break
pos_angles[d_idx] = running_pos # Store results on CPU, free GPU memory before next direction
neg_angles[d_idx] = running_neg 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 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) 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 # pos/neg are now numpy arrays (CPU) — combine on CPU
svf_combined = xp.zeros_like(dem) svf_combined = np.zeros((rows, cols), dtype=np.float32)
for r_idx in range(len(radii_m)): for r_idx in range(len(radii_m)):
horizon = xp.maximum(pos_angles[:, r_idx], neg_angles[:, r_idx]) horizon = np.maximum(pos_angles[:, r_idx], neg_angles[:, r_idx])
svf_r = xp.mean(xp.cos(horizon) ** 2, axis=0) svf_r = np.mean(np.cos(horizon) ** 2, axis=0)
svf_combined += svf_r * radius_weights[r_idx] 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 svf_np[nan_mask] = np.nan
_save_tif(output, svf_np, transform, crs) _save_tif(output, svf_np, transform, crs)
logger.info(f" ✓ SVF terminé ({time.time()-t0:.1f}s){' [GPU]' if _gpu_mod.HAS_GPU else ''}") 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: else:
angles = neg_angles angles = neg_angles
# Mean across directions and radii (equal weight) # Mean across directions and radii (equal weight) — on CPU now
openness = xp.mean(angles, axis=(0, 1)) openness = np.mean(angles, axis=(0, 1))
openness_result = to_cpu(xp.degrees(openness)).astype(np.float32) openness_result = np.degrees(openness).astype(np.float32)
openness_result[nan_mask] = np.nan openness_result[nan_mask] = np.nan
# Std normalization for cross-tile comparability # 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) weight_total = np.sum(weights)
n_radii = len(radii_m) n_radii = len(radii_m)
pos_combined = xp.zeros_like(dem) pos_combined = np.zeros((rows, cols), dtype=np.float64)
neg_combined = xp.zeros_like(dem) neg_combined = np.zeros((rows, cols), dtype=np.float64)
for r_idx in range(n_radii): for r_idx in range(n_radii):
for d_idx in range(n_dirs): 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) 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) 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 aniso_result[nan_mask] = np.nan
# Std normalization for cross-tile comparability # Std normalization for cross-tile comparability