From 6e60ec07ddfc533222b604c2f14e5befefec406d Mon Sep 17 00:00:00 2001 From: Kai ZHAO <694691@qq.com> Date: Mon, 27 Oct 2025 11:34:19 +0800 Subject: [PATCH 01/10] init distance transform --- pyproject.toml | 3 +-- test/test_distance_transform.py | 4 ++++ torchmorph/csrc/distance_transform_kernel.cu | 3 +++ 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 1dce09a..6027eac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,6 +14,5 @@ max-line-length = 100 extend-ignore = ["E203", "W503"] [tool.pytest.ini_options] -addopts = "-v" +addopts = "-v --import-mode=importlib" testpaths = ["test"] - diff --git a/test/test_distance_transform.py b/test/test_distance_transform.py index 166968c..e002200 100644 --- a/test/test_distance_transform.py +++ b/test/test_distance_transform.py @@ -14,6 +14,10 @@ def test_distance_transform(): y = tm.distance_transform(x) expected = x * 2 + # here we compare the output, i.e. results of our distance transform, + # with the expected output, e.g. the results of scipy.ndimage.distance_transform_edt + # currently, our implementation simply multiplies the input by 2, + # but eventually we have to implement the full algorithm. torch.testing.assert_close(y, expected) assert y.device.type == "cuda" assert y.shape == x.shape diff --git a/torchmorph/csrc/distance_transform_kernel.cu b/torchmorph/csrc/distance_transform_kernel.cu index 6a57f49..33d3abd 100644 --- a/torchmorph/csrc/distance_transform_kernel.cu +++ b/torchmorph/csrc/distance_transform_kernel.cu @@ -1,5 +1,8 @@ #include +// distance transform: https://en.wikipedia.org/wiki/Distance_transform +// https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.distance_transform_edt.html + __global__ void distance_transform_kernel(const float* in, float* out, int64_t N) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < N) { From 875483d93115b997e008a0b9f4fefb9bc8abd7ea Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Wed, 29 Oct 2025 02:08:05 +0800 Subject: [PATCH 02/10] =?UTF-8?q?docs:=20=E8=AF=A6=E7=BB=86=E8=AE=B0?= =?UTF-8?q?=E5=BD=95=E5=B9=B6=E4=BF=AE=E5=A4=8D=E9=A1=B9=E7=9B=AE=E7=8E=AF?= =?UTF-8?q?=E5=A2=83=E6=90=AD=E5=BB=BA=E4=B8=8E=E6=9E=84=E5=BB=BA=E6=B5=81?= =?UTF-8?q?=E7=A8=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 安装了miniconda来进行虚拟环境包管理 在执行`pip install -e. "import torch" 失败并抛出 `ModuleNotFoundError 所以在 `pyproject.toml` 的 `[build-system].requires` 列表中明确添加 `"torch"`。 这会强制 pip 在构建开始前，先将 torch 安装到其临时环境中，从而确保构建过程顺利完成。 --- pyproject.toml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 6027eac..10af0cd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,3 +16,7 @@ extend-ignore = ["E203", "W503"] [tool.pytest.ini_options] addopts = "-v --import-mode=importlib" testpaths = ["test"] + +[build-system] +requires = ["setuptools>=61.0", "wheel", "torch"] +build-backend = "setuptools.build_meta" \ No newline at end of file From 79c5acdb26faeba1f282a435c33ff36d04fc865c Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Sat, 1 Nov 2025 21:18:22 +0800 Subject: [PATCH 03/10] =?UTF-8?q?=E5=AE=9E=E7=8E=B0=E4=BA=8C=E7=BB=B4?= =?UTF-8?q?=E6=AC=A7=E5=BC=8F=E8=B7=9D=E7=A6=BB=E5=8F=98=E6=8D=A2=EF=BC=88?= =?UTF-8?q?EDT=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 CUDA 内核，分别用于行与列方向的距离变换 - 支持在 GPU 上处理二维张量 - 已通过基础单元测试验证正确性 - 注意：当前实现仅适用于二维情况，尚未推广到 N 维张量 --- test/test_distance_transform.py | 113 ++++++++++++++++--- torchmorph/csrc/distance_transform_kernel.cu | 109 ++++++++++++++---- torchmorph/csrc/torchmorph.cpp | 11 -- torchmorph/csrc/torchmorph.cu | 51 +++++++++ torchmorph/distance_transform.py | 2 +- 5 files changed, 239 insertions(+), 47 deletions(-) delete mode 100644 torchmorph/csrc/torchmorph.cpp create mode 100644 torchmorph/csrc/torchmorph.cu diff --git a/test/test_distance_transform.py b/test/test_distance_transform.py index e002200..493758a 100644 --- a/test/test_distance_transform.py +++ b/test/test_distance_transform.py @@ -1,24 +1,105 @@ import torch import pytest -import torchmorph as tm from scipy.ndimage import distance_transform_edt as dte +import torchmorph as tm +import numpy as np + +# --- 我们在这里定义所有的测试用例 --- + +# 用例 1: 我们之前成功的那个标准例子 +case_standard = np.array([ + [0, 1, 1, 1, 1], + [0, 0, 1, 1, 1], + [0, 1, 1, 1, 1], + [0, 1, 1, 1, 0], + [0, 1, 1, 0, 0] +], dtype=np.float32) + +# 用例 2: 全是背景 (0)，输出应该全是 0 +case_all_background = np.zeros((5, 5), dtype=np.float32) + +# 用例 3: 全是前景 (1)，输出应该也全是 0 (因为前景点到背景的距离未定义，SciPy默认输出0) +case_all_foreground = np.ones((5, 5), dtype=np.float32) + +# 用例 4: 只有一个背景点 (0) 在中间 +case_single_background = np.ones((5, 5), dtype=np.float32) +case_single_background[2, 2] = 0 +# 用例 5: 只有一个前景点 (1) 在中间 +case_single_foreground = np.zeros((5, 5), dtype=np.float32) +case_single_foreground[2, 2] = 1 -@pytest.mark.cuda -def test_distance_transform(): - """Test that tm.foo doubles all tensor elements.""" +# 用例 6: 非正方形的矩阵 (高 > 宽) +case_tall_matrix = np.array([ + [1, 0, 1], + [1, 1, 1], + [1, 1, 1], + [0, 1, 0], + [1, 1, 1], +], dtype=np.float32) + +# 用例 7: 非正方形的矩阵 (宽 > 高) +case_wide_matrix = np.array([ + [1, 1, 0, 1, 1], + [1, 1, 1, 1, 0], + [0, 1, 1, 1, 1], +], dtype=np.float32) + +# 用例 8: 棋盘格，考验对角线距离的计算 +case_checkerboard = np.array([ + [0, 1, 0, 1], + [1, 0, 1, 0], + [0, 1, 0, 1], + [1, 0, 1, 0], +], dtype=np.float32) + +# --- 使用 pytest.mark.parametrize 来自动运行所有测试用例 --- + +@pytest.mark.parametrize( + "input_numpy", + [ + pytest.param(case_standard, id="Standard Case"), + pytest.param(case_all_background, id="All Background"), + pytest.param(case_all_foreground, id="All Foreground"), + pytest.param(case_single_background, id="Single Background Pixel"), + pytest.param(case_single_foreground, id="Single Foreground Pixel"), + pytest.param(case_tall_matrix, id="Tall Matrix (H>W)"), + pytest.param(case_wide_matrix, id="Wide Matrix (W>H)"), + pytest.param(case_checkerboard, id="Checkerboard"), + ] +) +def test_distance_transform_comprehensive(input_numpy, request): + """ + 一个统一的测试函数，用来验证所有不同的输入情况。 + """ if not torch.cuda.is_available(): pytest.skip("CUDA not available") - x = torch.arange(6, dtype=torch.float32, device="cuda").reshape(2, 3) - y = tm.distance_transform(x) - - expected = x * 2 - # here we compare the output, i.e. results of our distance transform, - # with the expected output, e.g. the results of scipy.ndimage.distance_transform_edt - # currently, our implementation simply multiplies the input by 2, - # but eventually we have to implement the full algorithm. - torch.testing.assert_close(y, expected) - assert y.device.type == "cuda" - assert y.shape == x.shape - print("tm.foo test passed ✅") + # 准备输入数据 + x = torch.from_numpy(input_numpy).cuda() + + # 1. 运行你的 CUDA 实现 + y_cuda = tm.distance_transform(x) + + # 2. 运行 SciPy 官方实现 + y_ref_numpy = dte(input_numpy) + y_ref = torch.from_numpy(y_ref_numpy).to(torch.float32).cuda() + + # 打印结果用于直观对比 + print(f"\n\n--- Running Test: {request.node.callspec.id} ---") + print("Input Array:\n", input_numpy) + print("\nYour CUDA Implementation Output:\n", y_cuda.cpu().numpy()) + print("\nSciPy Reference Output:\n", y_ref.cpu().numpy()) + if request.node.callspec.id == "All Foreground": + # 对于这个特殊情况，我们不与 SciPy 比较。 + # 我们验证我们自己的逻辑：输出值是否都非常大 (代表无穷远)。 + print("\nSciPy has different behavior for this edge case. Verifying CUDA output is ~inf.") + # 断言所有元素都大于一个很大的阈值 + assert torch.all(y_cuda > 1e4) + else: + # 对于所有其他正常情况，我们与 SciPy 的黄金标准进行比较。 + y_ref_numpy = dte(input_numpy) + y_ref = torch.from_numpy(y_ref_numpy).to(torch.float32).cuda() + print("\nSciPy Reference Output:\n", y_ref.cpu().numpy()) + torch.testing.assert_close(y_cuda, y_ref, atol=1e-3, rtol=1e-3) + print("--- Test End ---") diff --git a/torchmorph/csrc/distance_transform_kernel.cu b/torchmorph/csrc/distance_transform_kernel.cu index 33d3abd..70527db 100644 --- a/torchmorph/csrc/distance_transform_kernel.cu +++ b/torchmorph/csrc/distance_transform_kernel.cu @@ -1,27 +1,98 @@ -#include +#include +#include -// distance transform: https://en.wikipedia.org/wiki/Distance_transform -// https://docs.scipy.org/doc/scipy/reference/generated/scipy.ndimage.distance_transform_edt.html +__global__ void edt_pass1_rows(const float* input, float* temp, int H, int W) { + int y = blockIdx.x; + if (y >= H) return; -__global__ void distance_transform_kernel(const float* in, float* out, int64_t N) { - int idx = blockIdx.x * blockDim.x + threadIdx.x; - if (idx < N) { - out[idx] = 2.0f * in[idx]; + extern __shared__ float sdata[]; + float* f = sdata; + int* v = (int*)(sdata + W); + float* z = (float*)(v + W + 1); + + for (int x = threadIdx.x; x < W; x += blockDim.x) { + float val = input[y * W + x]; + // 【关键任务修改】 + // 如果像素是 0 (背景)，则为源点 (距离0)；否则为无穷远。 + f[x] = (val < 0.5f) ? 0.0f : 1e10f; } -} + __syncthreads(); -torch::Tensor distance_transform_cuda(torch::Tensor input) { - auto output = torch::empty_like(input); - int64_t N = input.numel(); - int threads = 256; - int blocks = (N + threads - 1) / threads; + if (threadIdx.x == 0) { + int k = 0; + v[0] = 0; + z[0] = -1e10f; + z[1] = 1e10f; - distance_transform_kernel<<>>( - input.data_ptr(), - output.data_ptr(), - N - ); + for (int q = 1; q < W; q++) { + float s; + while (true) { + int p = v[k]; + s = ((f[q] + q * q) - (f[p] + p * p)) / (2.0f * (q - p)); + if (s > z[k]) { break; } + if (k == 0) { break; } + k--; + } + k++; + v[k] = q; + z[k] = s; + z[k + 1] = 1e10f; + } - return output; + k = 0; + for (int q = 0; q < W; q++) { + while (z[k + 1] < q) k++; + int p = v[k]; + temp[y * W + q] = (q - p) * (q - p) + f[p]; + } + } } +// PASS 2: 对每一列进行操作 +__global__ void edt_pass2_cols(const float* temp, float* output, int H, int W) { + int x = blockIdx.x; + if (x >= W) return; + + extern __shared__ float sdata[]; + float* f = sdata; + int* v = (int*)(sdata + H); + float* z = (float*)(v + H + 1); + + for (int y = threadIdx.x; y < H; y += blockDim.x) { + f[y] = temp[y * W + x]; + } + __syncthreads(); + + if (threadIdx.x == 0) { + int k = 0; + v[0] = 0; + z[0] = -1e10f; + z[1] = 1e10f; + + for (int q = 1; q < H; q++) { + float s; + while (true) { + int p = v[k]; + s = ((f[q] + q * q) - (f[p] + p * p)) / (2.0f * (q - p)); + if (s > z[k]) { + break; + } + if (k == 0) { + break; + } + k--; + } + k++; + v[k] = q; + z[k] = s; + z[k + 1] = 1e10f; + } + + k = 0; + for (int q = 0; q < H; q++) { + while (z[k + 1] < q) k++; + int p = v[k]; + output[q * W + x] = sqrtf((q - p) * (q - p) + f[p]); + } + } +} \ No newline at end of file diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp deleted file mode 100644 index 5d1dae8..0000000 --- a/torchmorph/csrc/torchmorph.cpp +++ /dev/null @@ -1,11 +0,0 @@ -#include - -// Declare CUDA implementations -torch::Tensor add_cuda(torch::Tensor input, float scalar); -torch::Tensor distance_transform_cuda(torch::Tensor input); - -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("add_cuda", &add_cuda, "Add tensor with scalar"); - m.def("distance_transform_cuda", &distance_transform_cuda, "Distance transform"); -} - diff --git a/torchmorph/csrc/torchmorph.cu b/torchmorph/csrc/torchmorph.cu new file mode 100644 index 0000000..dc2e4ce --- /dev/null +++ b/torchmorph/csrc/torchmorph.cu @@ -0,0 +1,51 @@ +// ========================================================================= +// 内容保存到: torchmorph/csrc/torchmorph.cpp +// ========================================================================= + +#include + +// 函数声明：告诉 C++ 编译器，这两个 CUDA 内核函数是在别的文件里定义的 +// 这样 C++ 代码才能成功调用 .cu 文件里的内核 +__global__ void edt_pass1_rows(const float* input, float* temp, int H, int W); +__global__ void edt_pass2_cols(const float* temp, float* output, int H, int W); + + + +// 主调函数 (运行在 CPU 上) +torch::Tensor distance_transform_cuda(torch::Tensor input) { + // 检查输入张量是否在 CUDA 上，以及是否为二维 + TORCH_CHECK(input.is_cuda(), "Input must be on CUDA"); + TORCH_CHECK(input.dim() == 2, "Only 2D tensors supported"); + + int H = input.size(0); + int W = input.size(1); + + // 创建临时的和最终的输出张量 + auto temp = torch::empty_like(input); + auto output = torch::empty_like(input); + + // 计算动态共享内存的大小 + size_t shared_mem_pass1 = W * sizeof(float) + (W + 1) * sizeof(int) + (W + 2) * sizeof(float); + size_t shared_mem_pass2 = H * sizeof(float) + (H + 1) * sizeof(int) + (H + 2) * sizeof(float); + + // 设置每个块的线程数 + int threads_per_block = 32; + + // <<<...>>> 语法：启动 CUDA 内核 + // 参数：Grid大小, Block大小, 共享内存大小, (可选的流) + + // Pass 1: 每行启动一个 block + edt_pass1_rows<<>>( + input.data_ptr(), temp.data_ptr(), H, W); + + // Pass 2: 每列启动一个 block + edt_pass2_cols<<>>( + temp.data_ptr(), output.data_ptr(), H, W); + + return output; +} + +// 使用 PYBIND11 将 C++ 函数绑定到 Python 模块 +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("distance_transform", &distance_transform_cuda, "CUDA-accelerated Exact Euclidean Distance Transform"); +} \ No newline at end of file diff --git a/torchmorph/distance_transform.py b/torchmorph/distance_transform.py index e4b54db..b4cd458 100644 --- a/torchmorph/distance_transform.py +++ b/torchmorph/distance_transform.py @@ -6,4 +6,4 @@ def distance_transform(input: torch.Tensor) -> torch.Tensor: """Distance Transform in CUDA.""" if not input.is_cuda: raise ValueError("Input tensor must be on CUDA device.") - return _C.distance_transform_cuda(input) + return _C.distance_transform(input) From 72943e44a723e794deb4a31af56b345a8e5834a9 Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Sun, 2 Nov 2025 04:12:02 +0800 Subject: [PATCH 04/10] =?UTF-8?q?N=E7=BB=B4=E6=89=B9=E5=A4=84=E7=90=86?= =?UTF-8?q?=E7=9A=84=E6=AC=A7=E6=B0=8F=E8=B7=9D=E7=A6=BB=E5=8F=98=E6=8D=A2?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- test/test_distance_transform.py | 138 ++++------ torchmorph/csrc/distance_transform_kernel.cu | 271 ++++++++++++++----- torchmorph/csrc/torchmorph.cpp | 10 + torchmorph/csrc/torchmorph.cu | 51 ---- torchmorph/distance_transform.py | 2 +- 5 files changed, 269 insertions(+), 203 deletions(-) create mode 100644 torchmorph/csrc/torchmorph.cpp delete mode 100644 torchmorph/csrc/torchmorph.cu diff --git a/test/test_distance_transform.py b/test/test_distance_transform.py index 493758a..3feffb4 100644 --- a/test/test_distance_transform.py +++ b/test/test_distance_transform.py @@ -4,102 +4,82 @@ import torchmorph as tm import numpy as np -# --- 我们在这里定义所有的测试用例 --- -# 用例 1: 我们之前成功的那个标准例子 -case_standard = np.array([ - [0, 1, 1, 1, 1], - [0, 0, 1, 1, 1], - [0, 1, 1, 1, 1], - [0, 1, 1, 1, 0], - [0, 1, 1, 0, 0] -], dtype=np.float32) +def batch_distance_transform_edt(batch_numpy): -# 用例 2: 全是背景 (0)，输出应该全是 0 -case_all_background = np.zeros((5, 5), dtype=np.float32) + is_single_sample = batch_numpy.ndim <= 2 + # (H, W) -> (1, H, W) + if is_single_sample: + batch_numpy = batch_numpy[np.newaxis, ...] + + results = [dte(sample) for sample in batch_numpy] + output = np.stack(results, axis=0) + # (1, H, W) -> (H, W) + if is_single_sample: + output = output.squeeze(0) + + return output -# 用例 3: 全是前景 (1)，输出应该也全是 0 (因为前景点到背景的距离未定义，SciPy默认输出0) -case_all_foreground = np.ones((5, 5), dtype=np.float32) +# 用例 1: 批处理的 2D 图像 +case_batch_2d = np.array([ + # 第 1 张图 + [[0, 1, 1, 1], [0, 0, 1, 1], [0, 1, 1, 0]], + # 第 2 张图 + [[0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 1]] +], dtype=np.float32) -# 用例 4: 只有一个背景点 (0) 在中间 -case_single_background = np.ones((5, 5), dtype=np.float32) -case_single_background[2, 2] = 0 -# 用例 5: 只有一个前景点 (1) 在中间 -case_single_foreground = np.zeros((5, 5), dtype=np.float32) -case_single_foreground[2, 2] = 1 +# 用例 2: 批处理的 3D 图像 +case_3d_sample1 = np.ones((4, 5, 6), dtype=np.float32); case_3d_sample1[1, 1, 1] = 0.0; case_3d_sample1[2, 3, 4] = 0.0 +case_3d_sample2 = np.ones((4, 5, 6), dtype=np.float32); case_3d_sample2[0, 0, 0] = 0.0 +case_batch_3d = np.stack([case_3d_sample1, case_3d_sample2], axis=0) -# 用例 6: 非正方形的矩阵 (高 > 宽) -case_tall_matrix = np.array([ - [1, 0, 1], - [1, 1, 1], - [1, 1, 1], - [0, 1, 0], - [1, 1, 1], +# 用例 3: 单张 2D 图像 (隐式批处理) +case_single_2d = np.array([ + [0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 1], [1, 0, 1, 0], ], dtype=np.float32) -# 用例 7: 非正方形的矩阵 (宽 > 高) -case_wide_matrix = np.array([ - [1, 1, 0, 1, 1], - [1, 1, 1, 1, 0], - [0, 1, 1, 1, 1], -], dtype=np.float32) -# 用例 8: 棋盘格，考验对角线距离的计算 -case_checkerboard = np.array([ - [0, 1, 0, 1], - [1, 0, 1, 0], - [0, 1, 0, 1], - [1, 0, 1, 0], -], dtype=np.float32) +# 用例 4: 单张 2D 图像 (显式批处理) +case_explicit_batch_one = case_single_2d[np.newaxis, ...] -# --- 使用 pytest.mark.parametrize 来自动运行所有测试用例 --- +# 用例 5: 含幺元维度的批处理 +case_dim_one = np.ones((2, 5, 1), dtype=np.float32) # 两张 5x1 的图片 +case_dim_one[0, 2, 0] = 0.0 +case_dim_one[1, 4, 0] = 0.0 + +# 用例 6: 1D 张量的批处理 +case_batch_1d = np.array([ + [1, 1, 0, 1, 0, 1], + [0, 1, 1, 1, 1, 0] +], dtype=np.float32) @pytest.mark.parametrize( "input_numpy", [ - pytest.param(case_standard, id="Standard Case"), - pytest.param(case_all_background, id="All Background"), - pytest.param(case_all_foreground, id="All Foreground"), - pytest.param(case_single_background, id="Single Background Pixel"), - pytest.param(case_single_foreground, id="Single Foreground Pixel"), - pytest.param(case_tall_matrix, id="Tall Matrix (H>W)"), - pytest.param(case_wide_matrix, id="Wide Matrix (W>H)"), - pytest.param(case_checkerboard, id="Checkerboard"), + pytest.param(case_batch_2d, id="批处理2D图像"), + pytest.param(case_batch_3d, id="批处理3D图像"), + pytest.param(case_single_2d, id="单张2D图像(隐式批处理)"), + pytest.param(case_explicit_batch_one, id="单张2D图像(显式批处理)"), + pytest.param(case_dim_one, id="含幺元维度的批处理"), + pytest.param(case_batch_1d, id="批处理1D数据"), ] ) -def test_distance_transform_comprehensive(input_numpy, request): - """ - 一个统一的测试函数，用来验证所有不同的输入情况。 - """ +def test_batch_processing(input_numpy, request): if not torch.cuda.is_available(): pytest.skip("CUDA not available") + x_numpy_contiguous = np.ascontiguousarray(input_numpy) + x = torch.from_numpy(x_numpy_contiguous).cuda() - # 准备输入数据 - x = torch.from_numpy(input_numpy).cuda() - - # 1. 运行你的 CUDA 实现 - y_cuda = tm.distance_transform(x) - - # 2. 运行 SciPy 官方实现 - y_ref_numpy = dte(input_numpy) + print(f"\n\n--- 正在运行测试: {request.node.callspec.id} ---") + print(f"输入张量形状: {x.shape}") + y_cuda = tm.distance_transform(x.clone()) + + y_ref_numpy = batch_distance_transform_edt(x_numpy_contiguous) y_ref = torch.from_numpy(y_ref_numpy).to(torch.float32).cuda() - - # 打印结果用于直观对比 - print(f"\n\n--- Running Test: {request.node.callspec.id} ---") - print("Input Array:\n", input_numpy) - print("\nYour CUDA Implementation Output:\n", y_cuda.cpu().numpy()) - print("\nSciPy Reference Output:\n", y_ref.cpu().numpy()) - if request.node.callspec.id == "All Foreground": - # 对于这个特殊情况，我们不与 SciPy 比较。 - # 我们验证我们自己的逻辑：输出值是否都非常大 (代表无穷远)。 - print("\nSciPy has different behavior for this edge case. Verifying CUDA output is ~inf.") - # 断言所有元素都大于一个很大的阈值 - assert torch.all(y_cuda > 1e4) - else: - # 对于所有其他正常情况，我们与 SciPy 的黄金标准进行比较。 - y_ref_numpy = dte(input_numpy) - y_ref = torch.from_numpy(y_ref_numpy).to(torch.float32).cuda() - print("\nSciPy Reference Output:\n", y_ref.cpu().numpy()) - torch.testing.assert_close(y_cuda, y_ref, atol=1e-3, rtol=1e-3) - print("--- Test End ---") + + assert y_cuda.shape == y_ref.shape, f"形状不匹配! CUDA输出: {y_cuda.shape}, SciPy应为: {y_ref.shape}" + print("CUDA 和 SciPy 输出形状匹配。") + + torch.testing.assert_close(y_cuda, y_ref, atol=1e-3, rtol=1e-3) + print("--- 断言通过 (数值接近) ---") \ No newline at end of file diff --git a/torchmorph/csrc/distance_transform_kernel.cu b/torchmorph/csrc/distance_transform_kernel.cu index 70527db..0bc5c26 100644 --- a/torchmorph/csrc/distance_transform_kernel.cu +++ b/torchmorph/csrc/distance_transform_kernel.cu @@ -1,98 +1,225 @@ -#include -#include +#include +#include + +// --- Kernel 1: 二值化内核 --- +/** + * @brief 对输入张量进行逐元素二值化。 + * @details 这是一个简单的并行操作。它将输入张量中的背景像素（值<0.5）设置为0， + * 并将前景像素（值>=0.5）设置为一个极大的值（1e20f），这在距离变换的上下文中 + * 可以被认为是无穷大。 + * @param in 输入张量的数据指针。 + * @param out 输出张量的数据指针。 + * @param N 张量中的元素总数。 + */ +__global__ void binarize_kernel(const float* in, float* out, int64_t N) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < N) { + out[idx] = (in[idx] < 0.5f) ? 0.0f : 1e20f; + } +} + +// --- Kernel 2: 1D Pass 距离平方计算内核 --- +/** + * @brief 沿着一个指定的空间维度，对N维张量执行一维抛物线下包络算法。 + * @details 这是Felzenszwalb和Huttenlocher EDT算法的核心。它通过将N维问题分解为N个 + * 一维问题来解决。此内核负责处理其中一个维度。它只计算距离的平方，以避免 + * 昂贵的开方运算并保持数值精度。 + * 每个CUDA线程块（block）负责处理一条完整的一维扫描线（slice）。 + * @param in_data 输入张量数据指针。 + * @param out_data 输出张量数据指针。 + * @param shape 描述输入张量形状的数组指针 (在GPU上)。 + * @param strides 描述输入张量步幅的数组指针 (在GPU上)。 + * @param ndim 张量的总维度数 (包括批处理维度)。 + * @param process_dim_sample 当前正在处理的空间维度索引 (0代表第一个空间维度，依此类推)。 + * @param total_slices 需要处理的一维扫描线总数 (batch_size * num_slices_per_sample)。 + * @param num_slices_per_sample 每个样本中，垂直于当前处理维度的扫描线数量。 + */ +__global__ void edt_1d_pass_sq_kernel( + const float* in_data, float* out_data, + const int64_t* shape, const int64_t* strides, + int32_t ndim, int32_t process_dim_sample, + int64_t total_slices, int64_t num_slices_per_sample +) { + // 每个线程块处理一条一维扫描线 + int64_t slice_idx = blockIdx.x; + if (slice_idx >= total_slices) return; + + + int64_t batch_idx = slice_idx / num_slices_per_sample; + int64_t slice_idx_in_sample = slice_idx % num_slices_per_sample; + int64_t batch_offset = batch_idx * strides[0]; // 获取批处理的基地址 + int64_t sample_base_offset = 0; + int64_t temp_idx = slice_idx_in_sample; + int sample_ndim = ndim - 1; + + // 从非处理维度中计算出样本内的基地址偏移 + for (int32_t d = sample_ndim - 1; d >= 0; --d) { + if (d == process_dim_sample) continue; // 跳过当前正在处理的维度 + int64_t size_of_dim = shape[d + 1]; + if (size_of_dim == 0) continue; + int64_t coord_in_dim = temp_idx % size_of_dim; + temp_idx /= size_of_dim; + sample_base_offset += coord_in_dim * strides[d + 1]; + } + + const int64_t process_dim_actual = process_dim_sample + 1; // 加上批处理维度的实际索引 + const int64_t N = shape[process_dim_actual]; // 当前处理维度的长度 + const int64_t stride = strides[process_dim_actual]; // 沿当前维度移动一个元素所需的步幅 + const int64_t base_offset = batch_offset + sample_base_offset; // 最终的起始地址 -__global__ void edt_pass1_rows(const float* input, float* temp, int H, int W) { - int y = blockIdx.x; - if (y >= H) return; extern __shared__ float sdata[]; - float* f = sdata; - int* v = (int*)(sdata + W); - float* z = (float*)(v + W + 1); - - for (int x = threadIdx.x; x < W; x += blockDim.x) { - float val = input[y * W + x]; - // 【关键任务修改】 - // 如果像素是 0 (背景)，则为源点 (距离0)；否则为无穷远。 - f[x] = (val < 0.5f) ? 0.0f : 1e10f; + float* f = sdata; // 存储函数值 g(p) = f(p) + p^2 + int* v = (int*)(sdata + N); // 存储抛物线顶点的索引 + float* z = (float*)(v + N + 1); // 存储相邻抛物线的交点 + + // 块内的所有线程协同将数据从全局内存加载到共享内存 + for (int64_t i = threadIdx.x; i < N; i += blockDim.x) { + f[i] = in_data[base_offset + i * stride]; } - __syncthreads(); + __syncthreads(); // 等待所有线程完成加载 - if (threadIdx.x == 0) { - int k = 0; - v[0] = 0; - z[0] = -1e10f; - z[1] = 1e10f; + //计算抛物线的下包络 + if (threadIdx.x == 0 && N > 0) { + int k = 0; // 下包络中的抛物线数量 + v[0] = 0; // 第一个抛物线的顶点索引为0 + z[0] = -1e20f; z[1] = 1e20f; // 初始化交点为负无穷和正无穷 - for (int q = 1; q < W; q++) { + // 遍历所有点，构建下包络 + for (int q = 1; q < N; q++) { float s; + // 寻找新的抛物线q应该插入的位置 while (true) { - int p = v[k]; + int p = v[k]; if (q == p) break; + // s 是抛物线 p 和 q 的交点的横坐标 s = ((f[q] + q * q) - (f[p] + p * p)) / (2.0f * (q - p)); + // 如果交点在当前区间的右侧，则找到了插入点 if (s > z[k]) { break; } - if (k == 0) { break; } + // 否则，抛物线p被q完全覆盖，需要移除p + if (k == 0) { break; } k--; } - k++; - v[k] = q; - z[k] = s; - z[k + 1] = 1e10f; + // 插入新的抛物线q + k++; + v[k] = q; + z[k] = s; + z[k + 1] = 1e20f; } - + // 计算距离平方 k = 0; - for (int q = 0; q < W; q++) { - while (z[k + 1] < q) k++; - int p = v[k]; - temp[y * W + q] = (q - p) * (q - p) + f[p]; + // 遍历所有点，找到其头顶上方的下包络线段，并计算距离 + for (int q = 0; q < N; q++) { + while (z[k + 1] < q) k++; // 找到点q所属的区间 + int p = v[k]; // 获取该区间的抛物线顶点索引 + // 计算距离平方: D(q)^2 = (q - p)^2 + g(p) + out_data[base_offset + q * stride] = (q - p) * (q - p) + f[p]; } } } -// PASS 2: 对每一列进行操作 -__global__ void edt_pass2_cols(const float* temp, float* output, int H, int W) { - int x = blockIdx.x; - if (x >= W) return; - extern __shared__ float sdata[]; - float* f = sdata; - int* v = (int*)(sdata + H); - float* z = (float*)(v + H + 1); - - for (int y = threadIdx.x; y < H; y += blockDim.x) { - f[y] = temp[y * W + x]; +// --- Kernel 3: 开平方根内核 --- +/** + * @brief 对张量中的每个元素计算平方根。 + * @details 这是一个简单的逐元素操作。由于之前的1D pass计算的是距离的平方， + * 此内核在所有维度处理完毕后被调用，以得到最终的欧氏距离。 + * @param data 需要进行开方操作的张量数据指针。 + * @param N 张量中的元素总数。 + */ +__global__ void sqrt_kernel(float* data, int64_t N) { + int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < N) { + data[idx] = sqrtf(data[idx]); } - __syncthreads(); +} - if (threadIdx.x == 0) { - int k = 0; - v[0] = 0; - z[0] = -1e10f; - z[1] = 1e10f; +// --- 主调函数 (Host) --- +/** + * @brief 执行N维欧氏距离变换。 + * @param input 一个N维的PyTorch张量，第一个维度被视为批处理（batch）维度。 + * @return 一个与输入形状相同的张量，包含每个点到最近前景点（值>=0.5）的欧氏距离。 + */ +torch::Tensor distance_transform_cuda(torch::Tensor input) { + auto original_input = input; + + // --- 预处理: 统一输入格式 --- + // 确保所有输入都至少是3D的 (B, ...)，方便后续统一处理。 + // 如果输入是 (H, W) 或 (L)，则变为 (1, H, W) 或 (1, L)。 + bool had_no_batch_dim = (input.dim() <= 2); + if (had_no_batch_dim) { input = input.unsqueeze(0); } - for (int q = 1; q < H; q++) { + // 检查输入张量是否在CUDA上并且是内存连续的 + TORCH_CHECK(input.is_cuda(), "Input must be on a CUDA device."); + TORCH_CHECK(input.is_contiguous(), "Input tensor must be contiguous."); - float s; - while (true) { - int p = v[k]; - s = ((f[q] + q * q) - (f[p] + p * p)) / (2.0f * (q - p)); - if (s > z[k]) { - break; - } - if (k == 0) { - break; - } - k--; - } - k++; - v[k] = q; - z[k] = s; - z[k + 1] = 1e10f; - } + if (input.numel() == 0) { return torch::empty_like(original_input); } + + // --- 获取张量元数据 --- + const auto ndim = input.dim(); + const auto sample_ndim = ndim - 1; // 空间维度 = 总维度 - 1 (batch) + const auto batch_size = input.size(0); + const int64_t N_total = input.numel(); + + auto shape_vec = input.sizes().vec(); + auto strides_vec = input.strides().vec(); + + // --- 内存分配: 使用Ping-Pong缓冲策略 --- + // 分配两个缓冲区，在处理每个维度时交替作为输入和输出，避免原地读写冲突。 + auto output = torch::empty_like(input); + auto buffer = (sample_ndim > 0) ? torch::empty_like(input) : output; - k = 0; - for (int q = 0; q < H; q++) { - while (z[k + 1] < q) k++; - int p = v[k]; - output[q * W + x] = sqrtf((q - p) * (q - p) + f[p]); + //二值化 + int threads = 256; // 定义每个线程块的线程数 + int blocks = (N_total + threads - 1) / threads; // 计算启动的线程块数 + binarize_kernel<<>>(input.data_ptr(), buffer.data_ptr(), N_total); + + //循环调用 edt_1d_pass_sq_kernel + // 将shape和strides信息从CPU内存拷贝到GPU内存，以便内核可以访问 + int64_t *shape_gpu, *strides_gpu; + cudaMalloc(&shape_gpu, ndim * sizeof(int64_t)); + cudaMalloc(&strides_gpu, ndim * sizeof(int64_t)); + cudaMemcpy(shape_gpu, shape_vec.data(), ndim * sizeof(int64_t), cudaMemcpyHostToDevice); + cudaMemcpy(strides_gpu, strides_vec.data(), ndim * sizeof(int64_t), cudaMemcpyHostToDevice); + + torch::Tensor current_input = buffer; + torch::Tensor current_output = output; + + // 遍历所有空间维度 + for (int32_t d_sample = 0; d_sample < sample_ndim; ++d_sample) { + // 为当前处理的维度计算启动内核所需的参数 + int64_t num_slices_per_sample = 1; + for(int i = 0; i < sample_ndim; ++i) { + if (i != d_sample) num_slices_per_sample *= shape_vec[i + 1]; } + int64_t total_slices = batch_size * num_slices_per_sample; + int64_t slice_len = shape_vec[d_sample + 1]; + + // 动态设置线程数和共享内存大小 + int threads_pass = (slice_len > 0 && slice_len < 256) ? slice_len : 256; + if (threads_pass == 0) threads_pass = 1; + size_t shared_mem_size = slice_len * sizeof(float) + (slice_len + 1) * sizeof(int) + (slice_len + 2) * sizeof(float); + + edt_1d_pass_sq_kernel<<>>( + current_input.data_ptr(), current_output.data_ptr(), + shape_gpu, strides_gpu, ndim, d_sample, total_slices, num_slices_per_sample + ); + // 交换输入和输出缓冲区，为下一个维度做准备 + std::swap(current_input, current_output); + } + + cudaFree(shape_gpu); + cudaFree(strides_gpu); + + //计算最终距离 + // 经过循环后，current_input 指向的是包含最终距离平方结果的张量 + sqrt_kernel<<>>(current_input.data_ptr(), N_total); + + // 如果最后一轮的输出不在我们期望的 output 张量里，就做一次拷贝 + if (current_input.data_ptr() != output.data_ptr()){ + output.copy_(current_input); } + + // 如果最初没有批处理维度，则移除我们添加的维度 + if (had_no_batch_dim) { output = output.squeeze(0); } + + return output; } \ No newline at end of file diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp new file mode 100644 index 0000000..b7f466a --- /dev/null +++ b/torchmorph/csrc/torchmorph.cpp @@ -0,0 +1,10 @@ +#include + +// Declare CUDA implementations +torch::Tensor add_cuda(torch::Tensor input, float scalar); +torch::Tensor distance_transform_cuda(torch::Tensor input); + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("add_cuda", &add_cuda, "Add tensor with scalar"); + m.def("distance_transform_cuda", &distance_transform_cuda, "Distance transform"); +} \ No newline at end of file diff --git a/torchmorph/csrc/torchmorph.cu b/torchmorph/csrc/torchmorph.cu deleted file mode 100644 index dc2e4ce..0000000 --- a/torchmorph/csrc/torchmorph.cu +++ /dev/null @@ -1,51 +0,0 @@ -// ========================================================================= -// 内容保存到: torchmorph/csrc/torchmorph.cpp -// ========================================================================= - -#include - -// 函数声明：告诉 C++ 编译器，这两个 CUDA 内核函数是在别的文件里定义的 -// 这样 C++ 代码才能成功调用 .cu 文件里的内核 -__global__ void edt_pass1_rows(const float* input, float* temp, int H, int W); -__global__ void edt_pass2_cols(const float* temp, float* output, int H, int W); - - - -// 主调函数 (运行在 CPU 上) -torch::Tensor distance_transform_cuda(torch::Tensor input) { - // 检查输入张量是否在 CUDA 上，以及是否为二维 - TORCH_CHECK(input.is_cuda(), "Input must be on CUDA"); - TORCH_CHECK(input.dim() == 2, "Only 2D tensors supported"); - - int H = input.size(0); - int W = input.size(1); - - // 创建临时的和最终的输出张量 - auto temp = torch::empty_like(input); - auto output = torch::empty_like(input); - - // 计算动态共享内存的大小 - size_t shared_mem_pass1 = W * sizeof(float) + (W + 1) * sizeof(int) + (W + 2) * sizeof(float); - size_t shared_mem_pass2 = H * sizeof(float) + (H + 1) * sizeof(int) + (H + 2) * sizeof(float); - - // 设置每个块的线程数 - int threads_per_block = 32; - - // <<<...>>> 语法：启动 CUDA 内核 - // 参数：Grid大小, Block大小, 共享内存大小, (可选的流) - - // Pass 1: 每行启动一个 block - edt_pass1_rows<<>>( - input.data_ptr(), temp.data_ptr(), H, W); - - // Pass 2: 每列启动一个 block - edt_pass2_cols<<>>( - temp.data_ptr(), output.data_ptr(), H, W); - - return output; -} - -// 使用 PYBIND11 将 C++ 函数绑定到 Python 模块 -PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("distance_transform", &distance_transform_cuda, "CUDA-accelerated Exact Euclidean Distance Transform"); -} \ No newline at end of file diff --git a/torchmorph/distance_transform.py b/torchmorph/distance_transform.py index b4cd458..7840158 100644 --- a/torchmorph/distance_transform.py +++ b/torchmorph/distance_transform.py @@ -6,4 +6,4 @@ def distance_transform(input: torch.Tensor) -> torch.Tensor: """Distance Transform in CUDA.""" if not input.is_cuda: raise ValueError("Input tensor must be on CUDA device.") - return _C.distance_transform(input) + return _C.distance_transform_cuda(input) \ No newline at end of file From f9420b25ffcca6115a49a7ec7c329c2db46583e7 Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Sun, 2 Nov 2025 15:07:10 +0800 Subject: [PATCH 05/10] =?UTF-8?q?=E4=BF=AE=E5=A4=8DBUG:=E5=8E=9F=E5=85=88?= =?UTF-8?q?=E7=89=88=E6=9C=AC=E4=B8=AD=E8=AF=AF=E5=B0=860=E5=BD=93?= =?UTF-8?q?=E6=88=90=E8=83=8C=E6=99=AF=EF=BC=8C1=E5=BD=93=E6=88=90?= =?UTF-8?q?=E5=89=8D=E6=99=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- torchmorph/csrc/distance_transform_kernel.cu | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/torchmorph/csrc/distance_transform_kernel.cu b/torchmorph/csrc/distance_transform_kernel.cu index 0bc5c26..d8cc8a7 100644 --- a/torchmorph/csrc/distance_transform_kernel.cu +++ b/torchmorph/csrc/distance_transform_kernel.cu @@ -1,11 +1,11 @@ #include #include -// --- Kernel 1: 二值化内核 --- -/** - * @brief 对输入张量进行逐元素二值化。 - * @details 这是一个简单的并行操作。它将输入张量中的背景像素（值<0.5）设置为0， - * 并将前景像素（值>=0.5）设置为一个极大的值（1e20f），这在距离变换的上下文中 +// --- Kernel 1: 二值化内核 --- +/* + * @brief 对输入张量进行逐元素二值化，为距离变换做准备。 + * @details 将前景点(in[idx] == 0)的初始距离设为0， + * 背景点的初始距离设为一个极大值(1e20f)，这在距离变换的上下文中 * 可以被认为是无穷大。 * @param in 输入张量的数据指针。 * @param out 输出张量的数据指针。 @@ -14,7 +14,9 @@ __global__ void binarize_kernel(const float* in, float* out, int64_t N) { int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < N) { - out[idx] = (in[idx] < 0.5f) ? 0.0f : 1e20f; + // 如果输入像素为0，则为前景点，其距离为0。 + // 如果输入像素非0，则为背景点，其初始距离为无穷大。 + out[idx] = (in[idx] == 0.0f) ? 0.0f : 1e20f; } } From 85bd1e546b2c53552868ff79b165e78c6a2176e5 Mon Sep 17 00:00:00 2001 From: Kai ZHAO <694691@qq.com> Date: Mon, 3 Nov 2025 16:53:10 +0800 Subject: [PATCH 06/10] returns both distance and index --- test/test_distance_transform.py | 9 ++--- torchmorph/csrc/distance_transform_kernel.cu | 37 +++++++++++--------- torchmorph/csrc/torchmorph.cpp | 4 +-- torchmorph/distance_transform.py | 8 ++++- 4 files changed, 35 insertions(+), 23 deletions(-) diff --git a/test/test_distance_transform.py b/test/test_distance_transform.py index 3feffb4..63b8568 100644 --- a/test/test_distance_transform.py +++ b/test/test_distance_transform.py @@ -73,13 +73,14 @@ def test_batch_processing(input_numpy, request): print(f"\n\n--- 正在运行测试: {request.node.callspec.id} ---") print(f"输入张量形状: {x.shape}") - y_cuda = tm.distance_transform(x.clone()) + dist_cuda, idx_cuda = tm.distance_transform(x.clone()) + print(f"Output index shape: {idx_cuda.shape}.") y_ref_numpy = batch_distance_transform_edt(x_numpy_contiguous) y_ref = torch.from_numpy(y_ref_numpy).to(torch.float32).cuda() - assert y_cuda.shape == y_ref.shape, f"形状不匹配! CUDA输出: {y_cuda.shape}, SciPy应为: {y_ref.shape}" + assert dist_cuda.shape == y_ref.shape, f"形状不匹配! CUDA输出: {dist_cuda.shape}, SciPy应为: {y_ref.shape}" print("CUDA 和 SciPy 输出形状匹配。") - torch.testing.assert_close(y_cuda, y_ref, atol=1e-3, rtol=1e-3) - print("--- 断言通过 (数值接近) ---") \ No newline at end of file + torch.testing.assert_close(dist_cuda, y_ref, atol=1e-3, rtol=1e-3) + print("--- 断言通过 (数值接近) ---") diff --git a/torchmorph/csrc/distance_transform_kernel.cu b/torchmorph/csrc/distance_transform_kernel.cu index d8cc8a7..534618f 100644 --- a/torchmorph/csrc/distance_transform_kernel.cu +++ b/torchmorph/csrc/distance_transform_kernel.cu @@ -11,7 +11,7 @@ * @param out 输出张量的数据指针。 * @param N 张量中的元素总数。 */ -__global__ void binarize_kernel(const float* in, float* out, int64_t N) { +__global__ void initialize_distance_kernel(const float* in, float* out, int64_t N) { int64_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < N) { // 如果输入像素为0，则为前景点，其距离为0。 @@ -140,7 +140,7 @@ __global__ void sqrt_kernel(float* data, int64_t N) { * @param input 一个N维的PyTorch张量，第一个维度被视为批处理（batch）维度。 * @return 一个与输入形状相同的张量，包含每个点到最近前景点（值>=0.5）的欧氏距离。 */ -torch::Tensor distance_transform_cuda(torch::Tensor input) { +std::tuple distance_transform_cuda(torch::Tensor input) { auto original_input = input; // --- 预处理: 统一输入格式 --- @@ -153,7 +153,6 @@ torch::Tensor distance_transform_cuda(torch::Tensor input) { TORCH_CHECK(input.is_cuda(), "Input must be on a CUDA device."); TORCH_CHECK(input.is_contiguous(), "Input tensor must be contiguous."); - if (input.numel() == 0) { return torch::empty_like(original_input); } // --- 获取张量元数据 --- const auto ndim = input.dim(); @@ -161,39 +160,45 @@ torch::Tensor distance_transform_cuda(torch::Tensor input) { const auto batch_size = input.size(0); const int64_t N_total = input.numel(); - auto shape_vec = input.sizes().vec(); + auto shape = input.sizes().vec(); + auto index_shape = shape; + index_shape.push_back(sample_ndim); + auto strides_vec = input.strides().vec(); // --- 内存分配: 使用Ping-Pong缓冲策略 --- // 分配两个缓冲区，在处理每个维度时交替作为输入和输出，避免原地读写冲突。 - auto output = torch::empty_like(input); - auto buffer = (sample_ndim > 0) ? torch::empty_like(input) : output; + auto distance = torch::zeros_like(input); + auto index = torch::zeros(index_shape); + auto buffer = (sample_ndim > 0) ? torch::empty_like(input) : distance; + + if (input.numel() == 0) { return std::make_tuple(distance, index); } //二值化 int threads = 256; // 定义每个线程块的线程数 int blocks = (N_total + threads - 1) / threads; // 计算启动的线程块数 - binarize_kernel<<>>(input.data_ptr(), buffer.data_ptr(), N_total); + initialize_distance_kernel<<>>(input.data_ptr(), buffer.data_ptr(), N_total); //循环调用 edt_1d_pass_sq_kernel // 将shape和strides信息从CPU内存拷贝到GPU内存，以便内核可以访问 int64_t *shape_gpu, *strides_gpu; cudaMalloc(&shape_gpu, ndim * sizeof(int64_t)); cudaMalloc(&strides_gpu, ndim * sizeof(int64_t)); - cudaMemcpy(shape_gpu, shape_vec.data(), ndim * sizeof(int64_t), cudaMemcpyHostToDevice); + cudaMemcpy(shape_gpu, shape.data(), ndim * sizeof(int64_t), cudaMemcpyHostToDevice); cudaMemcpy(strides_gpu, strides_vec.data(), ndim * sizeof(int64_t), cudaMemcpyHostToDevice); torch::Tensor current_input = buffer; - torch::Tensor current_output = output; + torch::Tensor current_output = distance; // 遍历所有空间维度 for (int32_t d_sample = 0; d_sample < sample_ndim; ++d_sample) { // 为当前处理的维度计算启动内核所需的参数 int64_t num_slices_per_sample = 1; for(int i = 0; i < sample_ndim; ++i) { - if (i != d_sample) num_slices_per_sample *= shape_vec[i + 1]; + if (i != d_sample) num_slices_per_sample *= shape[i + 1]; } int64_t total_slices = batch_size * num_slices_per_sample; - int64_t slice_len = shape_vec[d_sample + 1]; + int64_t slice_len = shape[d_sample + 1]; // 动态设置线程数和共享内存大小 int threads_pass = (slice_len > 0 && slice_len < 256) ? slice_len : 256; @@ -216,12 +221,12 @@ torch::Tensor distance_transform_cuda(torch::Tensor input) { sqrt_kernel<<>>(current_input.data_ptr(), N_total); // 如果最后一轮的输出不在我们期望的 output 张量里，就做一次拷贝 - if (current_input.data_ptr() != output.data_ptr()){ - output.copy_(current_input); + if (current_input.data_ptr() != distance.data_ptr()){ + distance.copy_(current_input); } // 如果最初没有批处理维度，则移除我们添加的维度 - if (had_no_batch_dim) { output = output.squeeze(0); } + if (had_no_batch_dim) { distance = distance.squeeze(0); } - return output; -} \ No newline at end of file + return std::make_tuple(distance, index); +} diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp index b7f466a..c79970c 100644 --- a/torchmorph/csrc/torchmorph.cpp +++ b/torchmorph/csrc/torchmorph.cpp @@ -2,9 +2,9 @@ // Declare CUDA implementations torch::Tensor add_cuda(torch::Tensor input, float scalar); -torch::Tensor distance_transform_cuda(torch::Tensor input); +std::tuple distance_transform_cuda(torch::Tensor input); PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("add_cuda", &add_cuda, "Add tensor with scalar"); m.def("distance_transform_cuda", &distance_transform_cuda, "Distance transform"); -} \ No newline at end of file +} diff --git a/torchmorph/distance_transform.py b/torchmorph/distance_transform.py index 7840158..0184be5 100644 --- a/torchmorph/distance_transform.py +++ b/torchmorph/distance_transform.py @@ -6,4 +6,10 @@ def distance_transform(input: torch.Tensor) -> torch.Tensor: """Distance Transform in CUDA.""" if not input.is_cuda: raise ValueError("Input tensor must be on CUDA device.") - return _C.distance_transform_cuda(input) \ No newline at end of file + if input.ndim < 2 or input.numel() == 0: + raise ValueError(f"Invalid input dimension: {input.shape}.") + + # binarize input + input[input != 0] = 1 + + return _C.distance_transform_cuda(input) From be4eb3dba55a34b80dc87089cd110d31715d7741 Mon Sep 17 00:00:00 2001 From: Kai ZHAO <694691@qq.com> Date: Mon, 3 Nov 2025 16:53:30 +0800 Subject: [PATCH 07/10] format --- test/test_distance_transform.py | 61 ++++++++++++++++++++------------- 1 file changed, 37 insertions(+), 24 deletions(-) diff --git a/test/test_distance_transform.py b/test/test_distance_transform.py index 63b8568..f65a0b4 100644 --- a/test/test_distance_transform.py +++ b/test/test_distance_transform.py @@ -11,48 +11,59 @@ def batch_distance_transform_edt(batch_numpy): # (H, W) -> (1, H, W) if is_single_sample: batch_numpy = batch_numpy[np.newaxis, ...] - + results = [dte(sample) for sample in batch_numpy] - output = np.stack(results, axis=0) + output = np.stack(results, axis=0) # (1, H, W) -> (H, W) if is_single_sample: output = output.squeeze(0) - + return output + # 用例 1: 批处理的 2D 图像 -case_batch_2d = np.array([ - # 第 1 张图 - [[0, 1, 1, 1], [0, 0, 1, 1], [0, 1, 1, 0]], - # 第 2 张图 - [[0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 1]] -], dtype=np.float32) +case_batch_2d = np.array( + [ + # 第 1 张图 + [[0, 1, 1, 1], [0, 0, 1, 1], [0, 1, 1, 0]], + # 第 2 张图 + [[0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 1]], + ], + dtype=np.float32, +) # 用例 2: 批处理的 3D 图像 -case_3d_sample1 = np.ones((4, 5, 6), dtype=np.float32); case_3d_sample1[1, 1, 1] = 0.0; case_3d_sample1[2, 3, 4] = 0.0 -case_3d_sample2 = np.ones((4, 5, 6), dtype=np.float32); case_3d_sample2[0, 0, 0] = 0.0 +case_3d_sample1 = np.ones((4, 5, 6), dtype=np.float32) +case_3d_sample1[1, 1, 1] = 0.0 +case_3d_sample1[2, 3, 4] = 0.0 +case_3d_sample2 = np.ones((4, 5, 6), dtype=np.float32) +case_3d_sample2[0, 0, 0] = 0.0 case_batch_3d = np.stack([case_3d_sample1, case_3d_sample2], axis=0) # 用例 3: 单张 2D 图像 (隐式批处理) -case_single_2d = np.array([ - [0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 1], [1, 0, 1, 0], -], dtype=np.float32) +case_single_2d = np.array( + [ + [0, 1, 0, 1], + [1, 0, 1, 0], + [0, 1, 0, 1], + [1, 0, 1, 0], + ], + dtype=np.float32, +) # 用例 4: 单张 2D 图像 (显式批处理) case_explicit_batch_one = case_single_2d[np.newaxis, ...] # 用例 5: 含幺元维度的批处理 -case_dim_one = np.ones((2, 5, 1), dtype=np.float32) # 两张 5x1 的图片 +case_dim_one = np.ones((2, 5, 1), dtype=np.float32) # 两张 5x1 的图片 case_dim_one[0, 2, 0] = 0.0 case_dim_one[1, 4, 0] = 0.0 # 用例 6: 1D 张量的批处理 -case_batch_1d = np.array([ - [1, 1, 0, 1, 0, 1], - [0, 1, 1, 1, 1, 0] -], dtype=np.float32) +case_batch_1d = np.array([[1, 1, 0, 1, 0, 1], [0, 1, 1, 1, 1, 0]], dtype=np.float32) + @pytest.mark.parametrize( "input_numpy", @@ -63,7 +74,7 @@ def batch_distance_transform_edt(batch_numpy): pytest.param(case_explicit_batch_one, id="单张2D图像(显式批处理)"), pytest.param(case_dim_one, id="含幺元维度的批处理"), pytest.param(case_batch_1d, id="批处理1D数据"), - ] + ], ) def test_batch_processing(input_numpy, request): if not torch.cuda.is_available(): @@ -75,12 +86,14 @@ def test_batch_processing(input_numpy, request): print(f"输入张量形状: {x.shape}") dist_cuda, idx_cuda = tm.distance_transform(x.clone()) print(f"Output index shape: {idx_cuda.shape}.") - + y_ref_numpy = batch_distance_transform_edt(x_numpy_contiguous) y_ref = torch.from_numpy(y_ref_numpy).to(torch.float32).cuda() - - assert dist_cuda.shape == y_ref.shape, f"形状不匹配! CUDA输出: {dist_cuda.shape}, SciPy应为: {y_ref.shape}" + + assert ( + dist_cuda.shape == y_ref.shape + ), f"形状不匹配! CUDA输出: {dist_cuda.shape}, SciPy应为: {y_ref.shape}" print("CUDA 和 SciPy 输出形状匹配。") - + torch.testing.assert_close(dist_cuda, y_ref, atol=1e-3, rtol=1e-3) print("--- 断言通过 (数值接近) ---") From 8f835e37e59e2749eb93ef9efb4bd88836a44453 Mon Sep 17 00:00:00 2001 From: Kai ZHAO <694691@qq.com> Date: Mon, 3 Nov 2025 18:00:28 +0800 Subject: [PATCH 08/10] benchmark --- benchmark/distance_transform.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 benchmark/distance_transform.py diff --git a/benchmark/distance_transform.py b/benchmark/distance_transform.py new file mode 100644 index 0000000..91219bb --- /dev/null +++ b/benchmark/distance_transform.py @@ -0,0 +1,24 @@ +import torch +import torch.utils.benchmark as benchmark +import scipy.ndimage as ndi +import torchmorph as tm + +for size in [64, 128, 256, 512, 1024, 2048]: + x = (torch.randn(1, 1, size, size, device="cuda") > 0).to(torch.float32) + + # TorchMorph CUDA + t1 = benchmark.Timer( + stmt="tm.distance_transform(x)", + setup="from __main__ import x, tm", + num_threads=torch.get_num_threads() + ) + # SciPy (CPU) + import numpy as np + x_np = x.cpu().squeeze().numpy() + t2 = benchmark.Timer( + stmt="ndi.distance_transform_edt(x_np)", + setup="from __main__ import x_np, ndi" + ) + + print(f"Size {size}:\n", t1.blocked_autorange()) + print(f"Size {size}:\n", t2.blocked_autorange()) From 7e108791dee292f7fcaa1e035e58384ce752adc8 Mon Sep 17 00:00:00 2001 From: Kai ZHAO <694691@qq.com> Date: Mon, 3 Nov 2025 18:27:09 +0800 Subject: [PATCH 09/10] benchmark outputs tables --- benchmark/distance_transform.py | 93 ++++++++++++++++++++++++++------- 1 file changed, 74 insertions(+), 19 deletions(-) diff --git a/benchmark/distance_transform.py b/benchmark/distance_transform.py index 91219bb..c659c82 100644 --- a/benchmark/distance_transform.py +++ b/benchmark/distance_transform.py @@ -1,24 +1,79 @@ import torch import torch.utils.benchmark as benchmark import scipy.ndimage as ndi +import numpy as np +from prettytable import PrettyTable import torchmorph as tm -for size in [64, 128, 256, 512, 1024, 2048]: - x = (torch.randn(1, 1, size, size, device="cuda") > 0).to(torch.float32) - - # TorchMorph CUDA - t1 = benchmark.Timer( - stmt="tm.distance_transform(x)", - setup="from __main__ import x, tm", - num_threads=torch.get_num_threads() - ) - # SciPy (CPU) - import numpy as np - x_np = x.cpu().squeeze().numpy() - t2 = benchmark.Timer( - stmt="ndi.distance_transform_edt(x_np)", - setup="from __main__ import x_np, ndi" - ) - - print(f"Size {size}:\n", t1.blocked_autorange()) - print(f"Size {size}:\n", t2.blocked_autorange()) +sizes = [64, 128, 256, 512, 1024] +batches = [1, 4, 8, 16] +dtype = torch.float32 +device = "cuda" +MIN_RUN = 1.0 # seconds per measurement + +torch.set_num_threads(torch.get_num_threads()) + +for B in batches: + table = PrettyTable() + table.field_names = [ + "Size", + "SciPy (ms/img)", + "Torch 1× (ms/img)", + "Torch batch (ms/img)", + "Speedup 1×", + "Speedup batch", + ] + for c in table.field_names: + table.align[c] = "r" + + for s in sizes: + # Inputs + x = (torch.randn(B, 1, s, s, device=device) > 0).to(dtype) + x_np_list = [x[i, 0].detach().cpu().numpy() for i in range(B)] + x_imgs = [x[i:i+1] for i in range(B)] + + # SciPy (CPU, one-by-one) + stmt_scipy = "out = [ndi.distance_transform_edt(arr) for arr in x_np_list]" + t_scipy = benchmark.Timer( + stmt=stmt_scipy, + setup="from __main__ import x_np_list, ndi", + num_threads=torch.get_num_threads(), + ).blocked_autorange(min_run_time=MIN_RUN) + scipy_per_img_ms = (t_scipy.median * 1e3) / B + + # Torch (CUDA, one-by-one) + stmt_torch1 = """ +for xi in x_imgs: + tm.distance_transform(xi) +""" + t_torch1 = benchmark.Timer( + stmt=stmt_torch1, + setup="from __main__ import x_imgs, tm", + num_threads=torch.get_num_threads(), + ).blocked_autorange(min_run_time=MIN_RUN) + torch1_per_img_ms = (t_torch1.median * 1e3) / B + + # Torch (CUDA, batched) + t_batch = benchmark.Timer( + stmt="tm.distance_transform(x)", + setup="from __main__ import x, tm", + num_threads=torch.get_num_threads(), + ).blocked_autorange(min_run_time=MIN_RUN) + torchB_per_img_ms = (t_batch.median * 1e3) / B + + # Speedups + speed1 = scipy_per_img_ms / torch1_per_img_ms + speedB = scipy_per_img_ms / torchB_per_img_ms + + table.add_row([ + s, + f"{scipy_per_img_ms:.3f}", + f"{torch1_per_img_ms:.3f}", + f"{torchB_per_img_ms:.3f}", + f"{speed1:.1f}×", + f"{speedB:.1f}×", + ]) + + print(f"\n=== Batch Size: {B} ===") + print(table) + From 7b6b8aae176f5dc8a7ac03898d84fb287fcea674 Mon Sep 17 00:00:00 2001 From: Kai ZHAO <694691@qq.com> Date: Mon, 3 Nov 2025 21:31:14 +0800 Subject: [PATCH 10/10] prettytable --- requirements-dev.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements-dev.txt b/requirements-dev.txt index fc961bb..0df6115 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -12,4 +12,4 @@ flake8>=6.0 setuptools>=65.0 wheel>=0.40 ninja>=1.11 # optional, speeds up torch extension builds - +prettytable>=3.16.0