bug fix

examples/cpd_nonrigid3d_cuda.py

@@ -0,0 +1,27 @@
+import numpy as np
+use_cuda = True
+if use_cuda:
+    import cupy as cp
+    to_cpu = cp.asnumpy
+    cp.cuda.set_allocator(cp.cuda.MemoryPool().malloc)
+else:
+    cp = np
+    to_cpu = lambda x: x
+import open3d as o3
+import transformations as trans
+from probreg import cpd
+from probreg import callbacks
+import utils
+import time
+
+source, target = utils.prepare_source_and_target_nonrigid_3d('face-x.txt', 'face-y.txt', voxel_size=5.0)
+source = cp.asarray(source.points, dtype=cp.float32)
+target = cp.asarray(target.points, dtype=cp.float32)
+
+acpd = cpd.NonRigidCPD(source, use_cuda=use_cuda)
+start = time.time()
+tf_param, _, _ = acpd.registration(target)
+elapsed = time.time() - start
+print("time: ", elapsed)
+
+print("result: ", to_cpu(tf_param.w), to_cpu(tf_param.g))

probreg/cupy_utils.py

@@ -1,13 +1,13 @@
 import cupy as cp
 
-_BLOCK_SIZE = 512
+_BLOCK_SIZE = 16
 
 squard_norm_outer_kernel = cp.RawKernel(r'''
 extern "C" __global__
 void squard_norm_outer_kernel(const float* x, const float* y,
                               int dim, int nx, int ny, float* res) {
-    int row = threadIdx.y + (blockIdx.y * blockDim.y);
-    int col = threadIdx.x + (blockIdx.x * blockDim.x);
+    int row = threadIdx.x + (blockIdx.x * blockDim.x);
+    int col = threadIdx.y + (blockIdx.y * blockDim.y);
     if (row < nx && col < ny) {
         res[row * ny + col] = 0.0;
         for (int i = 0; i < dim; ++i) {
@@ -24,19 +24,19 @@ def squared_kernel_sum(x, y):
     nx = xc.shape[0]
     ny = yc.shape[0]
     dim = xc.shape[1]
-    res = cp.array((nx, ny), dtype=cp.float32, order='C')
-    grid = (nx * ny) // _BLOCK_SIZE + 1
-    squard_norm_outer_kernel((grid,), (_BLOCK_SIZE,), (xc, yc, nx, ny, dim, res))
+    res = cp.zeros((nx, ny), dtype=cp.float32, order='C')
+    grid = ((nx + _BLOCK_SIZE - 1) // _BLOCK_SIZE, (ny + _BLOCK_SIZE - 1) // _BLOCK_SIZE, 1)
+    squard_norm_outer_kernel(grid, (_BLOCK_SIZE, _BLOCK_SIZE, 1), (xc, yc, dim, nx, ny, res))
     return res.sum() / (nx * ny * dim)
 
 
-def rbf_kernel(x, y, gamma):
+def rbf_kernel(x, y, beta):
     xc = cp.asarray(x, dtype=cp.float32, order='C')
     yc = cp.asarray(y, dtype=cp.float32, order='C')
     nx = xc.shape[0]
     ny = yc.shape[0]
     dim = xc.shape[1]
-    res = cp.array((nx, ny), dtype=cp.float32, order='C')
-    grid = (nx * ny) // _BLOCK_SIZE + 1
-    squard_norm_outer_kernel((grid,), (_BLOCK_SIZE,), (xc, yc, nx, ny, dim, res))
-    return (-res / (2.0 * gamma)).exp()
+    res = cp.zeros((nx, ny), dtype=cp.float32, order='C')
+    grid = ((nx + _BLOCK_SIZE - 1) // _BLOCK_SIZE, (ny + _BLOCK_SIZE - 1) // _BLOCK_SIZE, 1)
+    squard_norm_outer_kernel(grid, (_BLOCK_SIZE, _BLOCK_SIZE, 1), (xc, yc, dim, nx, ny, res))
+    return cp.exp((-res / (2.0 * beta)))

probreg/version.py

@@ -1 +1 @@
-__version__ = '0.2.1'
+__version__ = '0.2.2'