orisys-tech/orisys-sdk
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

improve gpu computation in HHD computation

Browse Source
This commit is contained in:
unknown
2026-01-13 13:41:28 +08:00
parent f3f28c2c4c
commit a4b8149131
3 changed files with 222 additions and 163 deletions
Download Patch File Download Diff File Expand all files Collapse all files
config/compound_eye.png
BIN
View File
Binary file not shown.
Side by Side Swipe Overlay

Before

Width:  |  Height:  |  Size: 2.0 MiB

After

Width:  |  Height:  |  Size: 2.0 MiB

example1_basic_pipeline.py
+4 -4
View File
@@ -19,11 +19,11 @@ def main():
# 步骤1：创建传感器对象 (必要步骤)，config_name为配置文件名（内置配置文件名）或路径（自定义配置文件路径） # 步骤1：创建传感器对象 (必要步骤)，config_name为配置文件名（内置配置文件名）或路径（自定义配置文件路径）
# 方式1: 从视频文件读取 # 方式1: 从视频文件读取
# sensor = orisys.Sensor("s1.mp4", config_name="prototype", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True) sensor = orisys.Sensor("s1.mp4", config_name="prototype", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True)
# 方式2: 从摄像头读取（取消注释使用） # 方式2: 从摄像头读取（取消注释使用）
sensor = orisys.Sensor(0, config_name="./config/lens_test.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True) # sensor = orisys.Sensor(0, config_name="./config/lens_test.json", cal_path= "./config/cal.npy", isstitch=True, cuda=True, verbose=True)
print("\n按 'q' 退出程序，按 'r' 重置光流追踪器") print("\n按 'q' 退出程序，按 'r' 重置光流追踪器")
@@ -33,8 +33,8 @@ def main():
# 步骤 2.1: 获取并拼接图像 (必要步骤) # 步骤 2.1: 获取并拼接图像 (必要步骤)
sensor.get_img() sensor.get_img()
print("img_size_raw:", sensor.img_size_raw) print("img_size_raw:", sensor.img_size_raw)
if sensor.detect_motion(threshold=2): # if sensor.detect_motion(threshold=2):
print("motion detected") # print("motion detected")
# 步骤 2.2: 计算形变场 (必要步骤) # 步骤 2.2: 计算形变场 (必要步骤)
sensor.compute_deformation() sensor.compute_deformation()
# 步骤 2.3: 计算触碰点 # 步骤 2.3: 计算触碰点
src/orisys/HHD.py
+200 -141
View File
@@ -104,6 +104,10 @@ class nHHD(object):
self.psolver.prepare(False) self.psolver.prepare(False)
# Example pattern for stream usage:
self.stream1 = cp.cuda.Stream()
self.stream2 = cp.cuda.Stream()
# create the 3-component decomposition # create the 3-component decomposition
def decompose(self, vfield, verbose=0): def decompose(self, vfield, verbose=0):
@@ -228,20 +232,21 @@ class nHHD(object):
dudy, dudx = cp.gradient(vfield_gpu[:,:,0], dx0, dx1) dudy, dudx = cp.gradient(vfield_gpu[:,:,0], dx0, dx1)
dvdy, dvdx = cp.gradient(vfield_gpu[:,:,1], dx0, dx1) dvdy, dvdx = cp.gradient(vfield_gpu[:,:,1], dx0, dx1)
div_gpu = dudx + dvdy div_gpu = cp.add(dudx, dvdy)
curlw_gpu = dvdx - dudy curlw_gpu = cp.subtract(dvdx, dudy)
if verbose > 0: if verbose > 0:
print('div =', div_gpu.shape, float(div_gpu.min()), float(div_gpu.max())) print('div =', div_gpu.shape, float(div_gpu.min()), float(div_gpu.max()))
print('curlw =', curlw_gpu.shape, float(curlw_gpu.min()), float(curlw_gpu.max())) print('curlw =', curlw_gpu.shape, float(curlw_gpu.min()), float(curlw_gpu.max()))
# Compute potentials on GPU (solve_gpu returns numpy, so we convert back)
nD_cpu = self.psolver.solve_gpu(cp.asnumpy(div_gpu), verbose > 0)
nRu_cpu = self.psolver.solve_gpu(cp.asnumpy(curlw_gpu), verbose > 0)
# Keep potentials on GPU # Compute potentials on GPU (solve_gpu returns numpy, so we convert back)
nD_gpu = cp.asarray(nD_cpu, dtype=cp.float32) with self.stream1:
nRu_gpu = cp.asarray(nRu_cpu, dtype=cp.float32) nD_gpu = self.psolver._solve_gpu_internal(div_gpu)
with self.stream2:
nRu_gpu = self.psolver._solve_gpu_internal(curlw_gpu)
self.stream1.synchronize()
self.stream2.synchronize()
if verbose > 0: if verbose > 0:
print('nD =', nD_gpu.shape, float(nD_gpu.min()), float(nD_gpu.max())) print('nD =', nD_gpu.shape, float(nD_gpu.min()), float(nD_gpu.max()))
@@ -256,15 +261,15 @@ class nHHD(object):
# Compute rotated gradient on GPU # Compute rotated gradient on GPU
rdy, rdx = cp.gradient(nRu_gpu, dx0, dx1) rdy, rdx = cp.gradient(nRu_gpu, dx0, dx1)
rdy *= -1.0 rdy = cp.multiply(rdy, -1.0)
r_gpu = cp.stack((rdy, rdx), axis=-1) r_gpu = cp.stack((rdy, rdx), axis=-1)
if verbose > 0: if verbose > 0:
print('r =', r_gpu.shape, float(r_gpu.min()), float(r_gpu.max())) print('r =', r_gpu.shape, float(r_gpu.min()), float(r_gpu.max()))
# Compute h on GPU # Compute h on GPU
h_gpu = d_gpu + r_gpu h_gpu = cp.add(d_gpu, r_gpu)
h_gpu = vfield_gpu - h_gpu h_gpu = cp.subtract(vfield_gpu, h_gpu)
if verbose > 0: if verbose > 0:
print('h =', h_gpu.shape, float(h_gpu.min()), float(h_gpu.max())) print('h =', h_gpu.shape, float(h_gpu.min()), float(h_gpu.max()))
@@ -276,11 +281,11 @@ class nHHD(object):
self._nRu_gpu = nRu_gpu self._nRu_gpu = nRu_gpu
# Also store CPU copies for compatibility # Also store CPU copies for compatibility
self.d = cp.asnumpy(d_gpu) # self.d = cp.asnumpy(d_gpu)
self.h = cp.asnumpy(h_gpu) # self.h = cp.asnumpy(h_gpu)
self.r = cp.asnumpy(r_gpu) # self.r = cp.asnumpy(r_gpu)
self.nRu = cp.asnumpy(nRu_gpu) # self.nRu = cp.asnumpy(nRu_gpu)
self.nD = nD_cpu # self.nD = cp.asnumpy(nD_gpu)
return (d_gpu, h_gpu, r_gpu, nRu_gpu) return (d_gpu, h_gpu, r_gpu, nRu_gpu)
@@ -375,13 +380,18 @@ class nHHD(object):
return self.get_force(vfield, verbose) return self.get_force(vfield, verbose)
# Decompose the vector field # Decompose the vector field
self.decompose(vfield, verbose) self.decompose_gpu(vfield, verbose)
# Keep data on GPU - avoid transfers by computing directly # Keep data on GPU - avoid transfers by computing directly
d_gpu = cp.asarray(self.d, dtype=cp.float32) # d_gpu = cp.asarray(self.d, dtype=cp.float32)
h_gpu = cp.asarray(self.h, dtype=cp.float32) # h_gpu = cp.asarray(self.h, dtype=cp.float32)
r_gpu = cp.asarray(self.r, dtype=cp.float32) # r_gpu = cp.asarray(self.r, dtype=cp.float32)
nRu_gpu = cp.asarray(self.nRu, dtype=cp.float32) # nRu_gpu = cp.asarray(self.nRu, dtype=cp.float32)
d_gpu = self._d_gpu
h_gpu = self._h_gpu
r_gpu = self._r_gpu
nRu_gpu = self._nRu_gpu
# Use reduction kernels for better performance # Use reduction kernels for better performance
# Compute Sn using L1 norm (sum of absolute values) # Compute Sn using L1 norm (sum of absolute values)
@@ -410,10 +420,17 @@ class nHHD(object):
r_y = r_gpu[:, :, 1] r_y = r_gpu[:, :, 1]
# Compute torque with fused sum # Compute torque with fused sum
Stor = float(cp.sum( # Stor = float(cp.sum(
(mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x + # (mgrid_0 - max_0) * r_y - (mgrid_1 - max_1) * r_x +
(mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x # (mgrid_0 - min_0) * r_y - (mgrid_1 - min_1) * r_x
)) # ))
# Replace lines 413-416:
if KERNEL_AVAILABLE and _torque_kernel is not None:
torque = _torque_kernel(
r_x, r_y, mgrid_0, mgrid_1,
float(max_0), float(max_1), float(min_0), float(min_1)
)
Stor = float(cp.sum(torque))
self.Sn, self.Stx, self.Sty, self.Stor = Sn, Stx, Sty, Stor self.Sn, self.Stx, self.Sty, self.Stor = Sn, Stx, Sty, Stor
output = numpy.array([Sn, Stx, Sty, Stor], dtype=numpy.float32) output = numpy.array([Sn, Stx, Sty, Stor], dtype=numpy.float32)
@@ -430,12 +447,14 @@ class nHHD(object):
cp.get_default_memory_pool().free_all_blocks() cp.get_default_memory_pool().free_all_blocks()
def get_shear(self): def get_shear(self):
if not CUPY_AVAILABLE:
return self.h return self.h
return cp.asnumpy(self._h_gpu)
def get_normal(self): def get_normal(self):
if not CUPY_AVAILABLE:
return self.d return self.d
return cp.asnumpy(self._d_gpu)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -574,72 +593,72 @@ class StructuredGrid(object):
return (dudx, dwdy, dudz, dvdx) return (dudx, dwdy, dudz, dvdx)
def divcurl_gpu(self, vfield, verbose=False): # def divcurl_gpu(self, vfield, verbose=False):
""" # """
GPU-accelerated computation of divergence and curl of vector field. # GPU-accelerated computation of divergence and curl of vector field.
Uses CuPy for faster gradient computation on GPU. # Uses CuPy for faster gradient computation on GPU.
""" # """
if not CUPY_AVAILABLE: # if not CUPY_AVAILABLE:
if verbose: # if verbose:
print('CuPy not available, falling back to CPU') # print('CuPy not available, falling back to CPU')
return self.divcurl(vfield, verbose) # return self.divcurl(vfield, verbose)
if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims): # if (vfield.shape[-1] != self.dim) or (vfield.shape[0:self.dim] != self.dims):
raise ValueError("Dimensions of vector field should match that of the grid") # raise ValueError("Dimensions of vector field should match that of the grid")
if verbose: # if verbose:
print(' Computing div and curl on GPU...', end ='') # print(' Computing div and curl on GPU...', end ='')
sys.stdout.flush() # sys.stdout.flush()
mtimer = Timer() # mtimer = Timer()
if self.dim == 2: # if self.dim == 2:
# Transfer to GPU # # Transfer to GPU
vfield_gpu = cp.asarray(vfield, dtype=cp.float32) # vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
# Compute gradients on GPU # # Compute gradients on GPU
dudy, dudx = cp.gradient(vfield_gpu[:,:,0], self.dx[0], self.dx[1]) # dudy, dudx = cp.gradient(vfield_gpu[:,:,0], self.dx[0], self.dx[1])
dvdy, dvdx = cp.gradient(vfield_gpu[:,:,1], self.dx[0], self.dx[1]) # dvdy, dvdx = cp.gradient(vfield_gpu[:,:,1], self.dx[0], self.dx[1])
# Compute div and curl on GPU # # Compute div and curl on GPU
div_gpu = dudx + dvdy # div_gpu = dudx + dvdy
curl_gpu = dvdx - dudy # curl_gpu = dvdx - dudy
# Transfer back to CPU # # Transfer back to CPU
div = cp.asnumpy(div_gpu) # div = cp.asnumpy(div_gpu)
curl = cp.asnumpy(curl_gpu) # curl = cp.asnumpy(curl_gpu)
if verbose: # if verbose:
print(' Done!', end ='') # print(' Done!', end ='')
mtimer.end() # mtimer.end()
return (div, curl) # return (div, curl)
elif self.dim == 3: # elif self.dim == 3:
# Transfer to GPU # # Transfer to GPU
vfield_gpu = cp.asarray(vfield, dtype=cp.float32) # vfield_gpu = cp.asarray(vfield, dtype=cp.float32)
# Compute gradients on GPU # # Compute gradients on GPU
dudz, dudy, dudx = cp.gradient(vfield_gpu[:,:,:,0], self.dx[0], self.dx[1], self.dx[2]) # dudz, dudy, dudx = cp.gradient(vfield_gpu[:,:,:,0], self.dx[0], self.dx[1], self.dx[2])
dvdz, dvdy, dvdx = cp.gradient(vfield_gpu[:,:,:,1], self.dx[0], self.dx[1], self.dx[2]) # dvdz, dvdy, dvdx = cp.gradient(vfield_gpu[:,:,:,1], self.dx[0], self.dx[1], self.dx[2])
dwdz, dwdy, dwdx = cp.gradient(vfield_gpu[:,:,:,2], self.dx[0], self.dx[1], self.dx[2]) # dwdz, dwdy, dwdx = cp.gradient(vfield_gpu[:,:,:,2], self.dx[0], self.dx[1], self.dx[2])
# Compute div and curl components on GPU # # Compute div and curl components on GPU
div_gpu = dudx + dvdy + dwdz # div_gpu = dudx + dvdy + dwdz
curlu_gpu = dwdy - dvdz # curlu_gpu = dwdy - dvdz
curlv_gpu = dudz - dwdx # curlv_gpu = dudz - dwdx
curlw_gpu = dvdx - dudy # curlw_gpu = dvdx - dudy
# Transfer back to CPU # # Transfer back to CPU
div = cp.asnumpy(div_gpu) # div = cp.asnumpy(div_gpu)
curlu = cp.asnumpy(curlu_gpu) # curlu = cp.asnumpy(curlu_gpu)
curlv = cp.asnumpy(curlv_gpu) # curlv = cp.asnumpy(curlv_gpu)
curlw = cp.asnumpy(curlw_gpu) # curlw = cp.asnumpy(curlw_gpu)
if verbose: # if verbose:
print(' Done!', end ='') # print(' Done!', end ='')
mtimer.end() # mtimer.end()
return (div, curlu, curlv, curlw) # return (div, curlu, curlv, curlw)
def curl3D(self, vfield, verbose=False): def curl3D(self, vfield, verbose=False):
@@ -1187,7 +1206,7 @@ class PoissonSolver(object):
# -------------------------------------------------------------------------- # --------------------------------------------------------------------------
# GPU-accelerated solve using CuPy # GPU-accelerated solve using CuPy
# -------------------------------------------------------------------------- # --------------------------------------------------------------------------
def solve_gpu(self, f, verbose=False): def solve_gpu(self, f, verbose=False, return_gpu=False):
""" """
GPU-accelerated Poisson solver using CuPy. GPU-accelerated Poisson solver using CuPy.
Falls back to CPU if CuPy is not available. Falls back to CPU if CuPy is not available.
@@ -1241,110 +1260,150 @@ class PoissonSolver(object):
# Scale by grid spacing (on GPU) # Scale by grid spacing (on GPU)
p_gpu *= self._gpu_cache['gdx_prod'] p_gpu *= self._gpu_cache['gdx_prod']
if verbose:
print('Poisson solver (GPU) finished')
gtimer.end()
# Transfer result back to CPU # Transfer result back to CPU
p = cp.asnumpy(p_gpu) if return_gpu:
return p_gpu # Return CuPy array
return cp.asnumpy(p_gpu)
# p = cp.asnumpy(p_gpu)
elif self.stype == 'S': elif self.stype == 'S':
raise SyntaxError('Spatial solver not suggested for structured grids') raise SyntaxError('Spatial solver not suggested for structured grids')
# ---------------------------------------------------------------------- # ----------------------------------------------------------------------
# unstructured points - not optimized for GPU # unstructured points - not optimized for GPU
elif self.ptype == 'P': # elif self.ptype == 'P':
if verbose: # if verbose:
print(' - GPU solver not optimized for unstructured points, using CPU') # print(' - GPU solver not optimized for unstructured points, using CPU')
return self.solve(f, verbose) # return self.solve(f, verbose)
# ---------------------------------------------------------------------- # ----------------------------------------------------------------------
if verbose:
print('Poisson solver (GPU) finished')
gtimer.end()
return p # return p
def solve_gpu_batch(self, f_list, verbose=False): def _solve_gpu_internal(self, f_gpu):
"""Internal GPU solve that accepts and returns GPU arrays."""
""" """
Batch GPU-accelerated Poisson solver for multiple inputs. GPU-accelerated Poisson solver using CuPy.
Solves multiple Poisson equations in a single GPU call. Falls back to CPU if CuPy is not available.
Args: Args:
f_list: List of input fields (numpy arrays) f_gpu: Input field (numpy array)
verbose: Print timing information
Returns: Returns:
List of solutions (numpy arrays on CPU) p_gpu: Solution (CuPy array)
""" """
if not CUPY_AVAILABLE:
if verbose:
print("CuPy not available, falling back to CPU solver")
return [self.solve(f, verbose) for f in f_list]
if not self.ready:
self.prepare(verbose)
gtimer = Timer()
if verbose:
print(f'\nSolving {len(f_list)} Poisson Eq. on GPU (batched)')
# Verify all inputs have same shape
fshape = f_list[0].shape
for f in f_list:
if f.shape != fshape:
raise ValueError("All input fields must have the same shape for batch processing")
# ---------------------------------------------------------------------- # ----------------------------------------------------------------------
# regular grid # regular grid
if self.ptype == 'G': if self.ptype == 'G':
if self.gdims != fshape:
raise ValueError("Shape of function should match shape of grid")
# convolution in frequency domain using GPU # convolution in frequency domain using GPU
if self.stype == 'F': if self.stype == 'F':
# Initialize GPU cache if needed # Initialize GPU cache if needed
if self._gpu_cache is None: if self._gpu_cache is None:
if verbose:
print(' - Transferring Green\'s function to GPU...')
self._gpu_cache = { self._gpu_cache = {
'G_gpu': cp.asarray(self.G, dtype=cp.float32), 'G_gpu': cp.asarray(self.G, dtype=cp.float32),
'gdx_prod': numpy.prod(self.gdx) 'gdx_prod': numpy.prod(self.gdx)
} }
results = []
# Process all inputs in batch
for i, f in enumerate(f_list):
# Transfer input to GPU
f_gpu = cp.asarray(f, dtype=cp.float32)
# Perform FFT convolution on GPU # Perform FFT convolution on GPU
# CuPy's fftconvolve is GPU-accelerated
p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache['G_gpu'], mode='same') p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache['G_gpu'], mode='same')
# Scale by grid spacing (on GPU) # Scale by grid spacing (on GPU)
p_gpu *= self._gpu_cache['gdx_prod'] p_gpu *= self._gpu_cache['gdx_prod']
# Transfer result back to CPU return p_gpu # Return CuPy array
results.append(cp.asnumpy(p_gpu))
if verbose:
print(f' - Solved {i+1}/{len(f_list)}')
if verbose: # def solve_gpu_batch(self, f_list, verbose=False):
print('Poisson solver (GPU batch) finished') # """
gtimer.end() # Batch GPU-accelerated Poisson solver for multiple inputs.
# Solves multiple Poisson equations in a single GPU call.
return results # Args:
# f_list: List of input fields (numpy arrays)
# verbose: Print timing information
elif self.stype == 'S': # Returns:
raise SyntaxError('Spatial solver not suggested for structured grids') # List of solutions (numpy arrays on CPU)
# """
# if not CUPY_AVAILABLE:
# if verbose:
# print("CuPy not available, falling back to CPU solver")
# return [self.solve(f, verbose) for f in f_list]
# ---------------------------------------------------------------------- # if not self.ready:
# unstructured points - not optimized for GPU # self.prepare(verbose)
elif self.ptype == 'P':
if verbose:
print(' - GPU solver not optimized for unstructured points, using CPU')
return [self.solve(f, verbose) for f in f_list]
# If we reach here, it's an unsupported case # gtimer = Timer()
raise ValueError(f"Unsupported point type '{self.ptype}' for GPU batch solver") # if verbose:
# print(f'\nSolving {len(f_list)} Poisson Eq. on GPU (batched)')
# # Verify all inputs have same shape
# fshape = f_list[0].shape
# for f in f_list:
# if f.shape != fshape:
# raise ValueError("All input fields must have the same shape for batch processing")
# # ----------------------------------------------------------------------
# # regular grid
# if self.ptype == 'G':
# if self.gdims != fshape:
# raise ValueError("Shape of function should match shape of grid")
# # convolution in frequency domain using GPU
# if self.stype == 'F':
# # Initialize GPU cache if needed
# if self._gpu_cache is None:
# if verbose:
# print(' - Transferring Green\'s function to GPU...')
# self._gpu_cache = {
# 'G_gpu': cp.asarray(self.G, dtype=cp.float32),
# 'gdx_prod': numpy.prod(self.gdx)
# }
# results = []
# # Process all inputs in batch
# for i, f in enumerate(f_list):
# # Transfer input to GPU
# f_gpu = cp.asarray(f, dtype=cp.float32)
# # Perform FFT convolution on GPU
# p_gpu = cp_signal.fftconvolve(f_gpu, self._gpu_cache['G_gpu'], mode='same')
# # Scale by grid spacing (on GPU)
# p_gpu *= self._gpu_cache['gdx_prod']
# # Transfer result back to CPU
# results.append(cp.asnumpy(p_gpu))
# if verbose:
# print(f' - Solved {i+1}/{len(f_list)}')
# if verbose:
# print('Poisson solver (GPU batch) finished')
# gtimer.end()
# return results
# elif self.stype == 'S':
# raise SyntaxError('Spatial solver not suggested for structured grids')
# # ----------------------------------------------------------------------
# # unstructured points - not optimized for GPU
# elif self.ptype == 'P':
# if verbose:
# print(' - GPU solver not optimized for unstructured points, using CPU')
# return [self.solve(f, verbose) for f in f_list]
# # If we reach here, it's an unsupported case
# raise ValueError(f"Unsupported point type '{self.ptype}' for GPU batch solver")
def clear_gpu_cache(self): def clear_gpu_cache(self):
"""Clear GPU cache to free GPU memory""" """Clear GPU cache to free GPU memory"""