From b9dd336f8b77abd25c71bb1bbf35e045c6899432 Mon Sep 17 00:00:00 2001 From: Kai ZHAO <694691@qq.com> Date: Sat, 9 May 2026 08:18:06 +0800 Subject: [PATCH 01/16] initialize docs --- docs/api/distance_transforms.md | 7 ++++++ docs/api/morphological_ops.md | 5 +++++ docs/index.md | 27 ++++++++++++++++++++++ mkdocs.yml | 40 +++++++++++++++++++++++++++++++++ requirements-docs.txt | 5 +++++ 5 files changed, 84 insertions(+) create mode 100644 docs/api/distance_transforms.md create mode 100644 docs/api/morphological_ops.md create mode 100644 docs/index.md create mode 100644 mkdocs.yml create mode 100644 requirements-docs.txt diff --git a/docs/api/distance_transforms.md b/docs/api/distance_transforms.md new file mode 100644 index 0000000..87a44e8 --- /dev/null +++ b/docs/api/distance_transforms.md @@ -0,0 +1,7 @@ +# Distance Transforms + +::: torchmorph.euclidean_distance_transform + +::: torchmorph.chamfer_distance_transform + +::: torchmorph.brute_force_distance_transform diff --git a/docs/api/morphological_ops.md b/docs/api/morphological_ops.md new file mode 100644 index 0000000..c14cb5a --- /dev/null +++ b/docs/api/morphological_ops.md @@ -0,0 +1,5 @@ +# Morphological Operations + +::: torchmorph.binary_dilation + +::: torchmorph.binary_erosion diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 0000000..4a2f4cb --- /dev/null +++ b/docs/index.md @@ -0,0 +1,27 @@ +# torchmorph + +GPU-accelerated morphological image processing operations for PyTorch. + +## Installation + +```bash +pip install torchmorph +``` + +## Quick Start + +```python +import torch +import torchmorph as tm + +x = torch.zeros(1, 1, 64, 64, device='cuda') +x[0, 0, 10:20, 10:20] = 1 + +dist = tm.euclidean_distance_transform(x) +dilated = tm.binary_dilation(x) +``` + +## API Reference + +- [Distance Transforms](api/distance_transforms.md) — Euclidean, Chamfer, and brute-force distance transforms +- [Morphological Operations](api/morphological_ops.md) — Binary dilation and erosion diff --git a/mkdocs.yml b/mkdocs.yml new file mode 100644 index 0000000..af9a73e --- /dev/null +++ b/mkdocs.yml @@ -0,0 +1,40 @@ +site_name: torchmorph +site_url: https://intcomp.org/tm/docs/ +site_description: GPU-accelerated morphological operations for PyTorch +repo_url: https://github.com/intcomp/torchmorph +repo_name: intcomp/torchmorph + +theme: + name: material + palette: + scheme: default + primary: indigo + features: + - navigation.sections + - content.code.copy + +plugins: + - search + - mkdocstrings: + handlers: + python: + options: + docstring_style: google + show_source: false + show_root_heading: true + show_symbol_type_heading: true + merge_init_into_class: true + +markdown_extensions: + - pymdownx.highlight: + anchor_linenums: true + - pymdownx.superfences + - pymdownx.inlinehilite + - admonition + - pymdownx.details + +nav: + - Home: index.md + - API Reference: + - Distance Transforms: api/distance_transforms.md + - Morphological Operations: api/morphological_ops.md diff --git a/requirements-docs.txt b/requirements-docs.txt new file mode 100644 index 0000000..b8b3cf2 --- /dev/null +++ b/requirements-docs.txt @@ -0,0 +1,5 @@ +-r requirements-test.txt +mkdocs-material +mkdocstrings[python] +pymdown-extensions + From ba1441005c79f51d7d44a15726af8562f781370b Mon Sep 17 00:00:00 2001 From: yuhandeng Date: Tue, 12 May 2026 02:44:59 +0800 Subject: [PATCH 02/16] add morphology api --- docs/api/morphological_ops.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/api/morphological_ops.md b/docs/api/morphological_ops.md index c14cb5a..6e6fb74 100644 --- a/docs/api/morphological_ops.md +++ b/docs/api/morphological_ops.md @@ -1,5 +1,10 @@ # Morphological Operations +::: torchmorph.generate_binary_structure ::: torchmorph.binary_dilation ::: torchmorph.binary_erosion + +::: torchmorph.binary_opening + +::: torchmorph.binary_closing \ No newline at end of file From 3fd88508f74d106ce9fd91af9612aa1b01acbfb4 Mon Sep 17 00:00:00 2001 From: yuhandeng Date: Tue, 12 May 2026 03:03:08 +0800 Subject: [PATCH 03/16] update docstring --- torchmorph/distance_transform.py | 236 +++++++++++++++-------------- torchmorph/morphology/binary.py | 161 +++++++++++++++++++- torchmorph/morphology/structure.py | 32 +++- 3 files changed, 305 insertions(+), 124 deletions(-) diff --git a/torchmorph/distance_transform.py b/torchmorph/distance_transform.py index 3bb7a19..de3b173 100644 --- a/torchmorph/distance_transform.py +++ b/torchmorph/distance_transform.py @@ -14,64 +14,63 @@ def euclidean_distance_transform( indices: Optional[torch.Tensor] = None, algorithm: str = "exact", ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor], None]: - """Exact Euclidean Distance Transform (EDT) using Felzenszwalb algorithm. + """Compute the Euclidean distance transform of an N-dimensional binary tensor + + The exact algorithm treats ``input`` as ``(B, C, *spatial)`` and processes + each batch/channel slice independently. Zero values are background sites; + non-zero values are foreground sites whose distance to the nearest + background site is computed. Args: - input: Binary input tensor (0 = background, non-zero = foreground). - Must be in (B, C, Spatial...) format where Spatial can be 1D, 2D, or 3D. - For single images, use unsqueeze to add batch and channel dims. - sampling: Spacing of elements along each spatial dimension. If a single - number, the spacing is uniform in all spatial dimensions. If a - sequence, it must match the number of spatial dimensions. - Default is None (unit spacing for all spatial dimensions). - Note: When sampling is not unit spacing, only "exact" algorithm is used. - return_distances: Whether to calculate the distance transform. - Default is True. - return_indices: Whether to calculate the feature transform (indices - of closest background element). Default is False. - distances: Optional output tensor for distances. If provided, must have - the same shape as input. If None and return_distances is True, - a new tensor will be created and returned. - indices: Optional output tensor for indices. If provided, must have shape - (spatial_ndim, ...) where ... matches input shape. If None and - return_indices is True, a new tensor will be created and returned. - algorithm: Algorithm to use for distance transform. Options: - - "exact": Use Felzenszwalb's exact algorithm (default). - - "jfa": Use Jump Flooding Algorithm (fast but approximate). - Only available for 2D/3D with unit sampling. - - "auto": Automatically choose based on input (uses JFA when - applicable, otherwise exact). + input (torch.Tensor): CUDA tensor with shape ``(B, C, *spatial)`` and at + least one spatial dimension. Add batch and channel dimensions for a + single image or volume. + sampling (Optional[float | Sequence[float]]): Physical spacing for the + spatial dimensions. ``None`` uses unit spacing, a scalar is + broadcast to all spatial dimensions, and a sequence must have either + one value or one value per spatial dimension. + return_distances (bool): Whether to compute distance values. + return_indices (bool): Whether to compute nearest-background coordinates. + distances (Optional[torch.Tensor]): Optional preallocated distance + output with the same shape as ``input``. Provided tensors are filled + in-place and omitted from the return value. + indices (Optional[torch.Tensor]): Optional preallocated index output + with shape ``(spatial_ndim, *input.shape)``. Provided tensors are + filled in-place and omitted from the return value. + algorithm (str): Algorithm selector. ``"exact"`` is the standard + N-dimensional path. ``"jfa"`` requests the approximate Jump + Flooding Algorithm for unit-spaced 2D/3D distance-only use and + otherwise falls back to exact. ``"auto"`` may choose JFA for + eligible 2D unit-spaced distance computations. Use ``"exact"`` + when requesting indices, using multi-channel inputs, or relying on + documented ``(B, C, *spatial)`` batch/channel semantics. Returns: - Depending on return_distances, return_indices, and whether output tensors - are provided: - - Returns distance tensor only when return_distances=True and distances=None - - Returns indices tensor only when return_indices=True and indices=None - - Returns tuple of (distances, indices) when both conditions above are met - - Returns None if output tensors are provided for all requested outputs - - Example: + torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None: Requested + outputs that were not supplied through ``distances`` or ``indices``. + For ``algorithm="exact"``, distances have shape ``input.shape`` and + indices have shape ``(spatial_ndim, *input.shape)``. Returns ``None`` + when all requested outputs were preallocated. + + Examples: + ```pycon + >>> import torch >>> import torchmorph as tm - >>> # 2D image: (B, C, H, W) - >>> x = torch.zeros(1, 1, 64, 64, device='cuda') - >>> x[0, 0, 10:20, 10:20] = 1 - >>> dist = tm.euclidean_distance_transform(x) - >>> dist, indices = tm.euclidean_distance_transform(x, return_indices=True) - >>> dist = tm.euclidean_distance_transform(x, sampling=[0.5, 1.0]) - >>> # Using JFA algorithm (faster for large images) - >>> dist = tm.euclidean_distance_transform(x, algorithm="jfa") - >>> # Using pre-allocated output tensors - >>> dist_out = torch.empty_like(x) - >>> tm.euclidean_distance_transform(x, distances=dist_out) # Returns None, fills dist_out - >>> # 3D volume: (B, C, D, H, W) - >>> x_3d = torch.zeros(2, 1, 32, 64, 64, device='cuda') - >>> dist_3d = tm.euclidean_distance_transform(x_3d, sampling=[2.0, 1.0, 1.0]) + >>> x = torch.zeros(1, 1, 4, 4, 4, 4, device="cuda") + >>> x[..., 1, 1, 1, 1] = 1 + >>> dist = tm.euclidean_distance_transform(x, algorithm="exact") + >>> dist.shape + torch.Size([1, 1, 4, 4, 4, 4]) + >>> dist, idx = tm.euclidean_distance_transform(x, algorithm="exact", return_indices=True) + >>> idx.shape + torch.Size([4, 1, 1, 4, 4, 4, 4]) + ``` """ if not input.is_cuda: raise ValueError("Input tensor must be on CUDA device.") if input.ndim < 3: raise ValueError( - f"Input must be (B, C, ) format with at least 3 dimensions, got {input.shape}. " + f"Input must be (B, C, *spatial) format with at least 3 dimensions, got {input.shape}. " "For single images, use unsqueeze to add batch and channel dims." ) if input.numel() == 0: @@ -156,51 +155,45 @@ def chamfer_distance_transform( distances: Optional[torch.Tensor] = None, indices: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor], None]: - """Chamfer Distance Transform (CDT). + """Compute the chamfer distance transform of an N-dimensional binary tensor - Calculates the distance transform of the input using a chamfer metric. - The input is treated as a binary image where non-zero values are foreground - and zero values are background. Distances are computed from each foreground - pixel to the nearest background pixel. + The input is treated as ``(B, C, *spatial)`` and each batch/channel slice is + processed independently. Zero values are background sites and non-zero + values are foreground sites. The CUDA implementation supports 1D through + 16D spatial inputs. Args: - input: Binary input tensor (0 = background, non-zero = foreground). - Must be in (B, C, H, W) or (B, C, D, H, W) format for batch processing, - or (H, W) / (D, H, W) for single images. - metric: Distance metric to use: - - "chessboard": L-infinity norm (default). Also known as Chebyshev distance. - - "taxicab": L1 norm. Also known as Manhattan or city-block distance. - - "cityblock": Alias for "taxicab". - - "manhattan": Alias for "taxicab". - return_distances: Whether to calculate the distance transform. Default is True. - return_indices: Whether to calculate the feature transform (indices of closest - background element). Default is False. - distances: Optional output tensor for distances. If provided, must have - the same shape as input. If None and return_distances is True, - a new tensor will be created. - indices: Optional output tensor for indices. If provided, must have shape - (..., ndim) where ... matches input shape. If None and return_indices - is True, a new tensor will be created. + input (torch.Tensor): CUDA tensor with shape ``(B, C, *spatial)`` and at + least one spatial dimension. + metric (str): Chamfer metric. ``"chessboard"`` computes L-infinity / + Chebyshev distance. ``"taxicab"`` computes L1 / Manhattan distance. + ``"cityblock"`` and ``"manhattan"`` are aliases for ``"taxicab"``. + return_distances (bool): Whether to compute distance values. + return_indices (bool): Whether to compute nearest-background coordinates. + distances (Optional[torch.Tensor]): Optional preallocated distance + output with the same shape as ``input``. + indices (Optional[torch.Tensor]): Optional preallocated index output + with shape ``(spatial_ndim, *input.shape)``. Returns: - Depending on return_distances, return_indices, and whether output tensors - are provided: - - Returns distance tensor only when return_distances=True and distances=None - - Returns indices tensor only when return_indices=True and indices=None - - Returns tuple of (distances, indices) when both conditions above are met - - Returns None if output tensors are provided for all requested outputs - - Example: + torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None: Requested + outputs that were not supplied through ``distances`` or ``indices``. + Distances have shape ``input.shape`` and indices have shape + ``(spatial_ndim, *input.shape)``. + + Examples: + ```pycon + >>> import torch >>> import torchmorph as tm - >>> # 2D image with batch: (B, C, H, W) - >>> x = torch.zeros(1, 1, 64, 64, device='cuda') - >>> x[0, 0, 10:20, 10:20] = 1 - >>> dist = tm.chamfer_distance_transform(x) # chessboard by default - >>> dist = tm.chamfer_distance_transform(x, metric='taxicab') - >>> dist, indices = tm.chamfer_distance_transform(x, return_indices=True) - >>> # Using pre-allocated output tensors - >>> dist_out = torch.empty_like(x) - >>> tm.chamfer_distance_transform(x, distances=dist_out) # Returns None, fills dist_out + >>> x = torch.zeros(1, 1, 5, 5, 5, device="cuda") + >>> x[..., 2, 2, 2] = 1 + >>> dist = tm.chamfer_distance_transform(x, metric="chessboard") + >>> dist.shape + torch.Size([1, 1, 5, 5, 5]) + >>> dist, idx = tm.chamfer_distance_transform(x, metric="taxicab", return_indices=True) + >>> idx.shape + torch.Size([3, 1, 1, 5, 5, 5]) + ``` """ if not input.is_cuda: raise ValueError("Input tensor must be on CUDA device.") @@ -271,46 +264,55 @@ def brute_force_distance_transform( distances: Optional[torch.Tensor] = None, indices: Optional[torch.Tensor] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor], None]: - """Brute-force distance transform. + """Compute the brute-force distance transform of an N-dimensional binary tensor - Calculates the distance transform of the input using a brute-force algorithm. - The algorithm computes the distance from each foreground pixel to ALL - background pixels and finds the minimum. This is $O(N*M)$ where $N$ is - number of foreground pixels and $M$ is number of background pixels. + The input is treated as ``(B, C, *spatial)`` and each batch/channel slice is + processed independently. For every foreground site, the kernel compares all + background sites and writes the minimum distance. The CUDA template dispatch + supports 1D through 8D spatial inputs. Args: - input: Binary input tensor (0 = background, non-zero = foreground). - Must be in (B, C, Spatial...) format. - metric: Distance metric to use: - - "euclidean": L2 norm (default). - - "taxicab": L1 norm (Manhattan distance). - - "chessboard": L-infinity norm (Chebyshev distance). - sampling: Spacing of elements along each spatial dimension. If a single - number, the spacing is uniform in all spatial dimensions. If a - sequence, it must match the number of spatial dimensions. - Default is None (unit spacing for all spatial dimensions). - return_distances: Whether to calculate the distance transform. - Default is True. - return_indices: Whether to calculate the feature transform (indices - of closest background element). Default is False. - distances: Optional output tensor for distances. If provided, must have - the same shape as input. - indices: Optional output tensor for indices. If provided, must have shape - (spatial_ndim, ...) where ... matches input shape. + input (torch.Tensor): CUDA tensor with shape ``(B, C, *spatial)`` and at + least one spatial dimension. + metric (str): Distance metric. ``"euclidean"`` computes L2 distance, + ``"taxicab"`` computes L1 / Manhattan distance, and + ``"chessboard"`` computes L-infinity / Chebyshev distance. + sampling (Optional[float | Sequence[float]]): Physical spacing for the + spatial dimensions. ``None`` uses unit spacing, a scalar is + broadcast to all spatial dimensions, and a sequence must have either + one value or one value per spatial dimension. + return_distances (bool): Whether to compute distance values. + return_indices (bool): Whether to compute nearest-background coordinates. + distances (Optional[torch.Tensor]): Optional preallocated distance + output with the same shape as ``input``. + indices (Optional[torch.Tensor]): Optional preallocated index output + with shape ``(spatial_ndim, *input.shape)``. Returns: - Depending on return_distances, return_indices, and whether output tensors - are provided: - - Returns distance tensor only when return_distances=True and distances=None - - Returns indices tensor only when return_indices=True and indices=None - - Returns tuple of (distances, indices) when both conditions above are met - - Returns None if output tensors are provided for all requested outputs + torch.Tensor | tuple[torch.Tensor, torch.Tensor] | None: Requested + outputs that were not supplied through ``distances`` or ``indices``. + Distances have shape ``input.shape`` and indices have shape + ``(spatial_ndim, *input.shape)``. + + Examples: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.zeros(1, 1, 3, 3, 3, 3, device="cuda") + >>> x[..., 1, 1, 1, 1] = 1 + >>> dist = tm.brute_force_distance_transform(x, metric="euclidean") + >>> dist.shape + torch.Size([1, 1, 3, 3, 3, 3]) + >>> dist, idx = tm.brute_force_distance_transform(x, return_indices=True) + >>> idx.shape + torch.Size([4, 1, 1, 3, 3, 3, 3]) + ``` """ if not input.is_cuda: raise ValueError("Input tensor must be on CUDA device.") if input.ndim < 3: raise ValueError( - f"Input must be (B, C, ) format with at least 3 dimensions, got {input.shape}." + f"Input must be (B, C, *spatial) format with at least 3 dimensions, got {input.shape}." ) if input.numel() == 0: raise ValueError(f"Invalid input: empty tensor with shape {input.shape}.") diff --git a/torchmorph/morphology/binary.py b/torchmorph/morphology/binary.py index 49ec8b9..0249092 100644 --- a/torchmorph/morphology/binary.py +++ b/torchmorph/morphology/binary.py @@ -121,10 +121,46 @@ def binary_erosion( border_value: bool = False, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """ - N-dimensional binary erosion for `(B, C, Spatial...)` tensors. + """Erode foreground regions in an N-dimensional binary tensor + + Treats non-zero input values as foreground and applies binary erosion to + each ``(B, C)`` sample independently. A foreground value survives only when + every active element of the structuring element overlaps foreground. + + Args: + input (torch.Tensor): Tensor with shape ``(B, C, *spatial)``. Non-zero + values are treated as foreground. + structure (Optional[torch.Tensor]): Structuring element with one axis + per spatial dimension. Non-zero values are active. If ``None``, a + rank-matched connectivity-1 structure is generated. + iterations (int): Number of erosion passes. Values less than ``1`` run + until the result no longer changes. + mask (Optional[torch.Tensor]): Boolean mask that restricts which output + locations may change. Use the same ``(B, C, *spatial)`` layout as + ``input``. + output (Optional[torch.Tensor]): Optional tensor to fill in-place with + the boolean result. The same tensor is returned. + border_value (bool): Value assumed outside the spatial boundary. + origin (int | tuple[int, ...]): Structuring-element origin offset. A + scalar is applied to every spatial dimension; a tuple must match the + number of spatial dimensions. - For a single image or volume, add batch and channel dimensions first. + Returns: + torch.Tensor: Boolean tensor with the same shape as ``input``. + + Examples: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.ones(1, 1, 5, 5) + >>> x[0, 0, 0, :] = 0 + >>> tm.binary_erosion(x).to(torch.int32)[0, 0] + tensor([[0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + [0, 1, 1, 1, 0], + [0, 1, 1, 1, 0], + [0, 0, 0, 0, 0]], dtype=torch.int32) + ``` """ return _binary_morphology( input=input, @@ -147,7 +183,47 @@ def binary_dilation( border_value: bool = False, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """binary dilation for `(B, C, Spatial...)` tensors.""" + """Dilate foreground regions in an N-dimensional binary tensor + + Treats non-zero input values as foreground and applies binary dilation to + each ``(B, C)`` sample independently. A location becomes foreground when any + active element of the structuring element overlaps foreground. + + Args: + input (torch.Tensor): Tensor with shape ``(B, C, *spatial)``. Non-zero + values are treated as foreground. + structure (Optional[torch.Tensor]): Structuring element with one axis + per spatial dimension. Non-zero values are active. If ``None``, a + rank-matched connectivity-1 structure is generated. + iterations (int): Number of dilation passes. Values less than ``1`` run + until the result no longer changes. + mask (Optional[torch.Tensor]): Boolean mask that restricts which output + locations may change. Use the same ``(B, C, *spatial)`` layout as + ``input``. + output (Optional[torch.Tensor]): Optional tensor to fill in-place with + the boolean result. The same tensor is returned. + border_value (bool): Value assumed outside the spatial boundary. + origin (int | tuple[int, ...]): Structuring-element origin offset. A + scalar is applied to every spatial dimension; a tuple must match the + number of spatial dimensions. + + Returns: + torch.Tensor: Boolean tensor with the same shape as ``input``. + + Examples: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.zeros(1, 1, 5, 5) + >>> x[0, 0, 2, 2] = 1 + >>> tm.binary_dilation(x).to(torch.int32)[0, 0] + tensor([[0, 0, 0, 0, 0], + [0, 0, 1, 0, 0], + [0, 1, 1, 1, 0], + [0, 0, 1, 0, 0], + [0, 0, 0, 0, 0]], dtype=torch.int32) + ``` + """ return _binary_morphology( input, structure, @@ -169,7 +245,44 @@ def binary_opening( border_value: bool = False, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """binary opening for '(B, C, ...)' Tensors .""" + """Remove small foreground components with binary opening + + Opening is erosion followed by dilation using the same structuring element. + The operation is applied independently to each ``(B, C)`` sample of an + ``(B, C, *spatial)`` tensor. + + Args: + input (torch.Tensor): Tensor with shape ``(B, C, *spatial)``. Non-zero + values are treated as foreground. + structure (Optional[torch.Tensor]): Structuring element with one axis + per spatial dimension. If ``None``, a rank-matched connectivity-1 + structure is generated. + iterations (int): Number of erosions followed by the same number of + dilations. Values less than ``1`` run each stage until convergence. + mask (Optional[torch.Tensor]): Boolean mask in the same + ``(B, C, *spatial)`` layout restricting which locations may change. + output (Optional[torch.Tensor]): Optional tensor to fill in-place. + border_value (bool): Value assumed outside the spatial boundary. + origin (int | tuple[int, ...]): Structuring-element origin offset. + + Returns: + torch.Tensor: Boolean tensor with the same shape as ``input``. + + Examples: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.zeros(1, 1, 5, 5) + >>> x[0, 0, 1:4, 1:4] = 1 + >>> x[0, 0, 0, 0] = 1 + >>> tm.binary_opening(x).to(torch.int32)[0, 0] + tensor([[0, 0, 0, 0, 0], + [0, 0, 1, 0, 0], + [0, 1, 1, 1, 0], + [0, 0, 1, 0, 0], + [0, 0, 0, 0, 0]], dtype=torch.int32) + ``` + """ x = binary_erosion( input, structure, @@ -200,7 +313,43 @@ def binary_closing( border_value: bool = False, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """binary closing for (B, C, ...) Tensors.""" + """Fill small background gaps with binary closing + + Closing is dilation followed by erosion using the same structuring element. + The operation is applied independently to each ``(B, C)`` sample of an + ``(B, C, *spatial)`` tensor. + + Args: + input (torch.Tensor): Tensor with shape ``(B, C, *spatial)``. Non-zero + values are treated as foreground. + structure (Optional[torch.Tensor]): Structuring element with one axis + per spatial dimension. If ``None``, a rank-matched connectivity-1 + structure is generated. + iterations (int): Number of dilations followed by the same number of + erosions. Values less than ``1`` run each stage until convergence. + mask (Optional[torch.Tensor]): Boolean mask in the same + ``(B, C, *spatial)`` layout restricting which locations may change. + output (Optional[torch.Tensor]): Optional tensor to fill in-place. + border_value (bool): Value assumed outside the spatial boundary. + origin (int | tuple[int, ...]): Structuring-element origin offset. + + Returns: + torch.Tensor: Boolean tensor with the same shape as ``input``. + + Examples: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.ones(1, 1, 5, 5) + >>> x[0, 0, 2, 2] = 0 + >>> tm.binary_closing(x, border_value=True).to(torch.int32)[0, 0] + tensor([[1, 1, 1, 1, 1], + [1, 1, 1, 1, 1], + [1, 1, 1, 1, 1], + [1, 1, 1, 1, 1], + [1, 1, 1, 1, 1]], dtype=torch.int32) + ``` + """ x = binary_dilation( input, structure, diff --git a/torchmorph/morphology/structure.py b/torchmorph/morphology/structure.py index 497c42a..d2d252f 100644 --- a/torchmorph/morphology/structure.py +++ b/torchmorph/morphology/structure.py @@ -3,7 +3,37 @@ def generate_binary_structure(rank: int, connectivity: int) -> Tensor: - """N-D generate binary structure""" + """Generate an N-dimensional binary structuring element + + The returned tensor has shape ``(3,) * rank``. Elements whose offset differs + from the center along at most ``connectivity`` axes are ``True``; all other + elements are ``False``. This matches SciPy's + ``ndimage.generate_binary_structure`` connectivity convention. + + Args: + rank (int): Number of spatial dimensions in the structuring element. + Must be at least ``1``. + connectivity (int): Neighborhood connectivity from ``1`` to ``rank``. + ``1`` includes axis-adjacent neighbors; ``rank`` includes the full + ``3 ** rank`` neighborhood. + + Returns: + torch.Tensor: Boolean tensor with shape ``(3,) * rank``. + + Examples: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> tm.generate_binary_structure(2, 1).to(dtype=torch.int32) + tensor([[0, 1, 0], + [1, 1, 1], + [0, 1, 0]], dtype=torch.int32) + >>> tm.generate_binary_structure(2, 2).to(dtype=torch.int32) + tensor([[1, 1, 1], + [1, 1, 1], + [1, 1, 1]], dtype=torch.int32) + ``` + """ if connectivity < 1 or connectivity > rank: raise ValueError(f"connectivity must be in [1, rank], got {connectivity}") From 4503fe15a02403454b758510b4b4f6c04b2b8a8e Mon Sep 17 00:00:00 2001 From: Yu Han Deng Date: Mon, 20 Jul 2026 16:11:54 +0800 Subject: [PATCH 04/16] docstring --- .gitignore | 4 +- docs/api/morphological_ops.md | 37 +++- docs/api/optimal_transport.md | 14 ++ docs/index.md | 18 +- docs/stylesheets/extra.css | 72 ++++++++ mkdocs.yml | 18 +- torchmorph/distance_transform.py | 109 +++++++++++- torchmorph/morphology/binary.py | 255 +++++++++++++++++++++++++-- torchmorph/morphology/grey.py | 270 +++++++++++++++++++++++++++-- torchmorph/morphology/structure.py | 35 +++- torchmorph/optimal_transport.py | 208 ++++++++++++++++++---- 11 files changed, 963 insertions(+), 77 deletions(-) create mode 100644 docs/api/optimal_transport.md create mode 100644 docs/stylesheets/extra.css diff --git a/.gitignore b/.gitignore index bb5af7e..98add0f 100644 --- a/.gitignore +++ b/.gitignore @@ -55,6 +55,6 @@ Thumbs.db .coverage htmlcov/ -# Local development documentation +# Local documentation artifacts /CONTEXT.md -/docs/ +/site/ diff --git a/docs/api/morphological_ops.md b/docs/api/morphological_ops.md index 6e6fb74..900bea5 100644 --- a/docs/api/morphological_ops.md +++ b/docs/api/morphological_ops.md @@ -1,10 +1,45 @@ # Morphological Operations + +All image and volume operators on this page expect CUDA tensors shaped as +`(B, C, Spatial...)`. Add batch and channel dimensions before processing a +single image or volume. + +## Structuring elements + ::: torchmorph.generate_binary_structure +::: torchmorph.iterate_structure + +## Binary morphology + ::: torchmorph.binary_dilation ::: torchmorph.binary_erosion +::: torchmorph.binary_propagation + +::: torchmorph.binary_fill_holes + +::: torchmorph.binary_hit_or_miss + ::: torchmorph.binary_opening -::: torchmorph.binary_closing \ No newline at end of file +::: torchmorph.binary_closing + +## Grayscale morphology + +::: torchmorph.grey_dilation + +::: torchmorph.grey_erosion + +::: torchmorph.grey_opening + +::: torchmorph.grey_closing + +::: torchmorph.morphological_gradient + +::: torchmorph.morphological_laplace + +::: torchmorph.white_tophat + +::: torchmorph.black_tophat diff --git a/docs/api/optimal_transport.md b/docs/api/optimal_transport.md new file mode 100644 index 0000000..c5a07bd --- /dev/null +++ b/docs/api/optimal_transport.md @@ -0,0 +1,14 @@ +# Optimal Transport + +TorchMorph provides a differentiable Sinkhorn solver that runs on CPU or CUDA. +CUDA `float32` inputs automatically use the fused implementation. + +::: torchmorph.build_cost_matrix + +::: torchmorph.SinkhornSolver + options: + members: + - data_preprocess + - forward + - plan + - potentials diff --git a/docs/index.md b/docs/index.md index 4a2f4cb..8c377c1 100644 --- a/docs/index.md +++ b/docs/index.md @@ -21,7 +21,23 @@ dist = tm.euclidean_distance_transform(x) dilated = tm.binary_dilation(x) ``` +Spatial morphology and distance-transform functions accept CUDA tensors shaped +as `(B, C, Spatial...)`, with one to eight spatial dimensions. Binary operators +treat nonzero values as foreground. + +Optimal transport works on CPU and CUDA with batches of flattened histograms: + +```python +source = torch.rand(8, 64, device="cuda") +target = torch.rand(8, 64, device="cuda") + +solver = tm.SinkhornSolver(epsilon=1.0, max_iter=200) +cost = tm.build_cost_matrix((8, 8), device=source.device) +distances = solver(source, target, cost) +``` + ## API Reference - [Distance Transforms](api/distance_transforms.md) — Euclidean, Chamfer, and brute-force distance transforms -- [Morphological Operations](api/morphological_ops.md) — Binary dilation and erosion +- [Morphological Operations](api/morphological_ops.md) — Structuring elements, binary morphology, and grayscale morphology +- [Optimal Transport](api/optimal_transport.md) — Cost matrices and differentiable Sinkhorn transport diff --git a/docs/stylesheets/extra.css b/docs/stylesheets/extra.css new file mode 100644 index 0000000..dbc24f9 --- /dev/null +++ b/docs/stylesheets/extra.css @@ -0,0 +1,72 @@ +/* Keep API headings compact and let signatures read like Python definitions. */ +.md-typeset .doc-heading { + display: flex; + align-items: center; + gap: 0.55rem; + margin: 2.4rem 0 0.75rem; + padding-bottom: 0.55rem; + border-bottom: 1px solid var(--md-default-fg-color--lightest); +} + +.md-typeset .doc-heading .doc-object-name { + color: var(--md-default-fg-color); + font-family: var(--md-code-font-family); + font-size: 0.82em; + font-weight: 700; + letter-spacing: -0.015em; +} + +.md-typeset .doc-heading code.doc-symbol { + flex: 0 0 auto; + border-radius: 999px; + padding: 0.2em 0.55em; + font-size: 0.62rem; + line-height: 1.5; +} + +/* mkdocstrings emits separate signatures as highlighted pre blocks. */ +.md-typeset .doc-object > .highlight, +.md-typeset .doc-overloads + .highlight { + margin: 0 0 1.25rem; + border: 1px solid var(--md-default-fg-color--lightest); + border-radius: 0.55rem; + background: var(--md-code-bg-color); + box-shadow: 0 0.08rem 0.22rem rgb(0 0 0 / 5%); + overflow: hidden; +} + +.md-typeset .doc-object > .highlight pre, +.md-typeset .doc-overloads + .highlight pre { + margin: 0; + padding: 0.85rem 1rem; + background: transparent; + overflow-x: auto; + font-size: 0.76rem; + line-height: 1.55; +} + +.md-typeset .doc-object > .highlight code, +.md-typeset .doc-overloads + .highlight code { + background: transparent; +} + +.md-typeset .doc-contents { + padding-left: 1rem; + border-left: 0.12rem solid var(--md-default-fg-color--lightest); +} + +@media screen and (max-width: 44.9844em) { + .md-typeset .doc-heading { + margin-top: 1.8rem; + } + + .md-typeset .doc-object > .highlight pre, + .md-typeset .doc-overloads + .highlight pre { + padding: 0.7rem 0.8rem; + font-size: 0.7rem; + } + + .md-typeset .doc-contents { + padding-left: 0.7rem; + } +} diff --git a/mkdocs.yml b/mkdocs.yml index af9a73e..e537064 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -7,12 +7,23 @@ repo_name: intcomp/torchmorph theme: name: material palette: - scheme: default - primary: indigo + - scheme: default + primary: indigo + toggle: + icon: material/brightness-7 + name: Switch to dark mode + - scheme: slate + primary: indigo + toggle: + icon: material/brightness-4 + name: Switch to light mode features: - navigation.sections - content.code.copy +extra_css: + - stylesheets/extra.css + plugins: - search - mkdocstrings: @@ -24,6 +35,8 @@ plugins: show_root_heading: true show_symbol_type_heading: true merge_init_into_class: true + separate_signature: true + line_length: 88 markdown_extensions: - pymdownx.highlight: @@ -38,3 +51,4 @@ nav: - API Reference: - Distance Transforms: api/distance_transforms.md - Morphological Operations: api/morphological_ops.md + - Optimal Transport: api/optimal_transport.md diff --git a/torchmorph/distance_transform.py b/torchmorph/distance_transform.py index f5ebf3e..074c82b 100644 --- a/torchmorph/distance_transform.py +++ b/torchmorph/distance_transform.py @@ -77,7 +77,43 @@ def euclidean_distance_transform( distances: Tensor | None = None, indices: Tensor | None = None, ) -> Tensor | tuple[Tensor, Tensor] | None: - """Euclidean distance transform for (B, C, Spatial...) CUDA tensors.""" + """Compute the exact Euclidean distance to the nearest background point + + Distances are computed independently for each batch and channel. Nonzero + input values are foreground and zero values are background. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + sampling (float or sequence[float], optional): Positive spacing along + each spatial axis. A scalar or one-element sequence is broadcast. + return_distances (bool): Return a newly allocated distance tensor unless + ``distances`` is supplied. + return_indices (bool): Return a newly allocated nearest-background index + tensor unless ``indices`` is supplied. + distances (torch.Tensor, optional): Preallocated distance output with + the same shape and device as ``input``. Supplying it enables + distance computation regardless of ``return_distances``. + indices (torch.Tensor, optional): Preallocated index output with shape + ``(spatial_ndim, *input.shape)`` on the input device. Supplying it + enables index computation regardless of ``return_indices``. + + Returns: + torch.Tensor, tuple[torch.Tensor, torch.Tensor], or None: Requested + newly allocated outputs. Distances have ``float32`` dtype and input + shape; indices have ``int32`` dtype and shape + ``(spatial_ndim, *input.shape)``. Preallocated outputs are filled but + omitted from the return value, so the function returns ``None`` when + every requested output is preallocated. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.tensor([[[0, 1, 1, 0]]], device="cuda") + >>> tm.euclidean_distance_transform(x) + tensor([[[0., 1., 1., 0.]]], device='cuda:0') + ``` + """ spatial_ndim, return_distances, return_indices = _prepare_distance_transform( input, return_distances, @@ -110,7 +146,39 @@ def chamfer_distance_transform( distances: Tensor | None = None, indices: Tensor | None = None, ) -> Tensor | tuple[Tensor, Tensor] | None: - """Chamfer distance transform for (B, C, Spatial...) CUDA tensors.""" + """Compute a chamfer distance to the nearest background point + + Distances are computed independently for each batch and channel. Nonzero + input values are foreground and zero values are background. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + metric (str): ``"chessboard"`` or ``"taxicab"``. ``"cityblock"`` and + ``"manhattan"`` are accepted aliases for ``"taxicab"``. + return_distances (bool): Return a newly allocated distance tensor unless + ``distances`` is supplied. + return_indices (bool): Return a newly allocated nearest-background index + tensor unless ``indices`` is supplied. + distances (torch.Tensor, optional): Preallocated distance output with + the same shape and device as ``input``. + indices (torch.Tensor, optional): Preallocated index output with shape + ``(spatial_ndim, *input.shape)`` on the input device. + + Returns: + torch.Tensor, tuple[torch.Tensor, torch.Tensor], or None: Requested + newly allocated outputs. Distances are ``float32`` and indices are + ``int32``. Preallocated outputs are filled but omitted from the return + value. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.tensor([[[0, 1, 1, 0]]], device="cuda") + >>> tm.chamfer_distance_transform(x, metric="taxicab") + tensor([[[0., 1., 1., 0.]]], device='cuda:0') + ``` + """ _, return_distances, return_indices = _prepare_distance_transform( input, return_distances, @@ -147,7 +215,42 @@ def brute_force_distance_transform( distances: Tensor | None = None, indices: Tensor | None = None, ) -> Tensor | tuple[Tensor, Tensor] | None: - """Brute-force distance transform for (B, C, Spatial...) CUDA tensors.""" + """Compute distances to background points by exhaustive search + + This reference-style transform supports Euclidean, taxicab, and chessboard + metrics and computes each batch and channel independently. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``; + nonzero values are foreground and zeros are background. + metric (str): One of ``"euclidean"``, ``"taxicab"``, or + ``"chessboard"``. + sampling (float or sequence[float], optional): Positive spatial spacing. + A scalar or one-element sequence is broadcast to every axis. + return_distances (bool): Return a newly allocated distance tensor unless + ``distances`` is supplied. + return_indices (bool): Return a newly allocated nearest-background index + tensor unless ``indices`` is supplied. + distances (torch.Tensor, optional): Preallocated distance output with + the same shape and device as ``input``. + indices (torch.Tensor, optional): Preallocated index output with shape + ``(spatial_ndim, *input.shape)`` on the input device. + + Returns: + torch.Tensor, tuple[torch.Tensor, torch.Tensor], or None: Requested + newly allocated outputs. Distances are ``float32`` and indices are + ``int32``. Preallocated outputs are filled but omitted from the return + value. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.tensor([[[0, 1, 1, 0]]], device="cuda") + >>> tm.brute_force_distance_transform(x, metric="euclidean") + tensor([[[0., 1., 1., 0.]]], device='cuda:0') + ``` + """ spatial_ndim, return_distances, return_indices = _prepare_distance_transform( input, return_distances, diff --git a/torchmorph/morphology/binary.py b/torchmorph/morphology/binary.py index 7e566d3..fac2b8d 100644 --- a/torchmorph/morphology/binary.py +++ b/torchmorph/morphology/binary.py @@ -96,12 +96,43 @@ def binary_erosion( output: Tensor | None = None, border_value: bool = False, origin: int | tuple[int, ...] = 0, -) -> Tensor: - """ - N-dimensional binary erosion for `(B, C, Spatial...)` CUDA tensors. - - For a single image or volume, add batch and channel dimensions first. - """ +) -> Tensor: + """Erode binary objects in a batched CUDA tensor + + A value is foreground when it is nonzero. The operation supports one to + eight spatial dimensions and always treats the first two dimensions as + batch and channel. When ``iterations`` is less than one, erosion continues + until the result no longer changes. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + structure (torch.Tensor, optional): Spatial structuring element. Nonzero + entries participate in the neighborhood. If ``None``, uses + :func:`generate_binary_structure` with connectivity ``1``. + iterations (int): Number of sequential erosions. Values less than ``1`` + repeat the operation until convergence. + mask (torch.Tensor, optional): Tensor with the same shape as ``input``. + Only locations where the mask is nonzero may change. + output (torch.Tensor, optional): Preallocated tensor with the same shape + and device as ``input``. When supplied, receives and is returned as + the result. + border_value (bool): Value used outside the input boundary. + origin (int or tuple[int, ...]): Offset of the structuring-element + anchor. A scalar is applied to every spatial dimension. + + Returns: + torch.Tensor: Boolean erosion result, or ``output`` when it is supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.ones((1, 1, 5, 5), device="cuda") + >>> y = tm.binary_erosion(x, structure=torch.ones(3, 3)) + >>> y.dtype, y.sum().item() + (torch.bool, 9) + ``` + """ return _binary_morphology( input=input, structure=structure, @@ -122,8 +153,44 @@ def binary_dilation( output: Tensor | None = None, border_value: bool = False, origin: int | tuple[int, ...] = 0, -) -> Tensor: - """N-dimensional binary dilation for `(B, C, Spatial...)` CUDA tensors.""" +) -> Tensor: + """Dilate binary objects in a batched CUDA tensor + + A value is foreground when it is nonzero. The operation supports one to + eight spatial dimensions and always treats the first two dimensions as + batch and channel. When ``iterations`` is less than one, dilation continues + until the result no longer changes. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + structure (torch.Tensor, optional): Spatial structuring element. Nonzero + entries participate in the neighborhood. If ``None``, uses + :func:`generate_binary_structure` with connectivity ``1``. + iterations (int): Number of sequential dilations. Values less than ``1`` + repeat the operation until convergence. + mask (torch.Tensor, optional): Tensor with the same shape as ``input``. + Only locations where the mask is nonzero may change. + output (torch.Tensor, optional): Preallocated tensor with the same shape + and device as ``input``. When supplied, receives and is returned as + the result. + border_value (bool): Value used outside the input boundary. + origin (int or tuple[int, ...]): Offset of the structuring-element + anchor. A scalar is applied to every spatial dimension. + + Returns: + torch.Tensor: Boolean dilation result, or ``output`` when it is supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.zeros((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 2, 2] = 1 + >>> y = tm.binary_dilation(x, structure=torch.ones(3, 3)) + >>> y.dtype, y.sum().item() + (torch.bool, 9) + ``` + """ return _binary_morphology( input, structure, @@ -143,8 +210,39 @@ def binary_propagation( output: Tensor | None = None, border_value: bool = False, origin: int | tuple[int, ...] = 0, -) -> Tensor: - """N-dimensional binary propagation for `(B, C, Spatial...)` CUDA tensors.""" +) -> Tensor: + """Propagate a binary seed through a mask until convergence + + This repeatedly dilates ``input`` while allowing changes only where + ``mask`` is nonzero. It is useful for morphological reconstruction. + + Args: + input (torch.Tensor): Binary seed CUDA tensor with shape + ``(B, C, Spatial...)``; nonzero values are foreground. + structure (torch.Tensor, optional): Spatial structuring element. If + ``None``, uses connectivity ``1``. + mask (torch.Tensor, optional): Tensor with the same shape as ``input``. + Only nonzero mask locations may change. If ``None``, propagation is + unconstrained. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + border_value (bool): Value used outside the input boundary. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Boolean converged result, or ``output`` when supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> seed = torch.zeros((1, 1, 5, 5), device="cuda") + >>> seed[0, 0, 2, 2] = 1 + >>> mask = torch.ones_like(seed) + >>> tm.binary_propagation(seed, mask=mask).all().item() + True + ``` + """ return _binary_morphology( input, structure, @@ -162,8 +260,35 @@ def binary_fill_holes( structure: Tensor | None = None, output: Tensor | None = None, origin: int | tuple[int, ...] = 0, -) -> Tensor: - """Fill holes in binary objects for `(B, C, Spatial...)` CUDA tensors.""" +) -> Tensor: + """Fill holes enclosed by binary objects + + Background connected to the tensor boundary is reconstructed and removed; + enclosed background regions become foreground. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``; + nonzero values are foreground. + structure (torch.Tensor, optional): Spatial structuring element used to + determine background connectivity. If ``None``, uses connectivity + ``1``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Boolean tensor with enclosed holes filled, or ``output``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.ones((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 2, 2] = 0 + >>> tm.binary_fill_holes(x)[0, 0, 2, 2].item() + True + ``` + """ validate_bcs_input(input) validate_output(input, output) @@ -192,8 +317,39 @@ def binary_hit_or_miss( output: Tensor | None = None, origin1: int | tuple[int, ...] = 0, origin2: int | tuple[int, ...] | None = None, -) -> Tensor: - """N-dimensional binary hit-or-miss transform for `(B, C, Spatial...)` CUDA tensors.""" +) -> Tensor: + """Find binary configurations with the hit-or-miss transform + + A location matches when ``structure1`` fits the foreground and + ``structure2`` fits the background at the same anchor. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``; + nonzero values are foreground. + structure1 (torch.Tensor, optional): Foreground structuring element. If + ``None``, uses connectivity ``1``. + structure2 (torch.Tensor, optional): Background structuring element. If + ``None``, uses the logical complement of ``structure1``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + origin1 (int or tuple[int, ...]): Anchor offset for ``structure1``. + origin2 (int or tuple[int, ...], optional): Anchor offset for + ``structure2``. If ``None``, uses ``origin1``. + + Returns: + torch.Tensor: Boolean tensor marking matching locations, or ``output``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.zeros((1, 1, 3, 3), device="cuda") + >>> x[0, 0, 1, 1] = 1 + >>> hit = tm.binary_hit_or_miss(x, structure1=torch.ones(1, 1)) + >>> hit[0, 0, 1, 1].item() + True + ``` + """ spatial_ndim = validate_bcs_input(input) validate_output(input, output) origin1 = _normalize_origin(origin1, spatial_ndim) @@ -236,8 +392,40 @@ def binary_opening( output: Tensor | None = None, border_value: bool = False, origin: int | tuple[int, ...] = 0, -) -> Tensor: - """N-dimensional binary opening for `(B, C, Spatial...)` CUDA tensors.""" +) -> Tensor: + """Open binary objects by erosion followed by dilation + + Opening removes foreground features that cannot contain the structuring + element while largely preserving larger objects. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``; + nonzero values are foreground. + structure (torch.Tensor, optional): Spatial structuring element. If + ``None``, uses connectivity ``1``. + iterations (int): Number of erosions followed by the same number of + dilations. Values less than ``1`` iterate each stage to convergence. + mask (torch.Tensor, optional): Same-shaped tensor whose nonzero entries + identify locations that may change. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + border_value (bool): Value used outside the input boundary. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Boolean opened tensor, or ``output`` when supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.zeros((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 1:4, 1:4] = 1 + >>> x[0, 0, 0, 0] = 1 + >>> tm.binary_opening(x).sum().item() + 5 + ``` + """ x = binary_erosion( input, structure, @@ -267,8 +455,39 @@ def binary_closing( output: Tensor | None = None, border_value: bool = False, origin: int | tuple[int, ...] = 0, -) -> Tensor: - """N-dimensional binary closing for `(B, C, Spatial...)` CUDA tensors.""" +) -> Tensor: + """Close binary objects by dilation followed by erosion + + Closing fills small gaps and joins nearby foreground regions while largely + preserving the extent of larger objects. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``; + nonzero values are foreground. + structure (torch.Tensor, optional): Spatial structuring element. If + ``None``, uses connectivity ``1``. + iterations (int): Number of dilations followed by the same number of + erosions. Values less than ``1`` iterate each stage to convergence. + mask (torch.Tensor, optional): Same-shaped tensor whose nonzero entries + identify locations that may change. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + border_value (bool): Value used outside the input boundary. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Boolean closed tensor, or ``output`` when supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.zeros((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 2, 1] = x[0, 0, 2, 3] = 1 + >>> tm.binary_closing(x)[0, 0, 2, 2].item() + True + ``` + """ x = binary_dilation( input, structure, diff --git a/torchmorph/morphology/grey.py b/torchmorph/morphology/grey.py index 0f53b37..9f55494 100644 --- a/torchmorph/morphology/grey.py +++ b/torchmorph/morphology/grey.py @@ -115,9 +115,37 @@ def grey_erosion( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional grey erosion for ``(B, C, Spatial...)`` CUDA tensors. - - Computation uses float32. The result is float32 unless ``output`` is given. + """Erode a batched grayscale CUDA tensor + + Computes the local minimum after subtracting the non-flat ``structure``. + Computation uses ``float32`` and supports one to eight spatial dimensions. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. A scalar is broadcast to every spatial axis. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. Must match ``structure`` when both are supplied. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Eroded tensor in ``float32``, or ``output`` when supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.arange(9, device="cuda").reshape(1, 1, 3, 3) + >>> tm.grey_erosion(x, size=3).shape + torch.Size([1, 1, 3, 3]) + ``` """ return _grey_morphology( input, @@ -142,9 +170,37 @@ def grey_dilation( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional grey dilation for ``(B, C, Spatial...)`` CUDA tensors. - - Computation uses float32. The result is float32 unless ``output`` is given. + """Dilate a batched grayscale CUDA tensor + + Computes the local maximum after adding the non-flat ``structure``. + Computation uses ``float32`` and supports one to eight spatial dimensions. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. A scalar is broadcast to every spatial axis. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. Must match ``structure`` when both are supplied. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Dilated tensor in ``float32``, or ``output`` when supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.arange(9, device="cuda").reshape(1, 1, 3, 3) + >>> tm.grey_dilation(x, size=3).dtype + torch.float32 + ``` """ return _grey_morphology( input, @@ -169,9 +225,38 @@ def grey_opening( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional grey opening for ``(B, C, Spatial...)`` CUDA tensors. - - Computation uses float32. The result is float32 unless ``output`` is given. + """Open a grayscale tensor by erosion followed by dilation + + Opening suppresses bright features smaller than the selected structuring + element. Computation uses ``float32``. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Opened tensor in ``float32``, or ``output`` when supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.zeros((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 2, 2] = 1 + >>> tm.grey_opening(x, size=3).max().item() + 0.0 + ``` """ kwargs = _element_kwargs(size, footprint, structure, mode, cval, origin) eroded = grey_erosion(input, **kwargs) @@ -188,9 +273,38 @@ def grey_closing( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional grey closing for ``(B, C, Spatial...)`` CUDA tensors. - - Computation uses float32. The result is float32 unless ``output`` is given. + """Close a grayscale tensor by dilation followed by erosion + + Closing suppresses dark features smaller than the selected structuring + element. Computation uses ``float32``. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Closed tensor in ``float32``, or ``output`` when supplied. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.ones((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 2, 2] = 0 + >>> tm.grey_closing(x, size=3).min().item() + 1.0 + ``` """ kwargs = _element_kwargs(size, footprint, structure, mode, cval, origin) dilated = grey_dilation(input, **kwargs) @@ -207,7 +321,38 @@ def morphological_gradient( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional morphological gradient for ``(B, C, Spatial...)`` CUDA tensors.""" + """Compute the morphological gradient of a grayscale tensor + + The gradient is ``dilation(input) - erosion(input)`` and emphasizes local + intensity transitions. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Morphological gradient in ``float32``, or ``output``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.arange(25, device="cuda").reshape(1, 1, 5, 5) + >>> tm.morphological_gradient(x, size=3).shape + torch.Size([1, 1, 5, 5]) + ``` + """ validate_bcs_input(input) validate_output(input, output) @@ -231,7 +376,38 @@ def white_tophat( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional white top-hat filter for ``(B, C, Spatial...)`` CUDA tensors.""" + """Extract small bright features with a white top-hat filter + + Computes ``input - grey_opening(input)`` using ``float32`` arithmetic. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: White top-hat response in ``float32``, or ``output``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.zeros((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 2, 2] = 2 + >>> tm.white_tophat(x, size=3).max().item() + 2.0 + ``` + """ validate_bcs_input(input) validate_output(input, output) @@ -253,7 +429,38 @@ def black_tophat( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional black top-hat filter for ``(B, C, Spatial...)`` CUDA tensors.""" + """Extract small dark features with a black top-hat filter + + Computes ``grey_closing(input) - input`` using ``float32`` arithmetic. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Black top-hat response in ``float32``, or ``output``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.ones((1, 1, 5, 5), device="cuda") + >>> x[0, 0, 2, 2] = 0 + >>> tm.black_tophat(x, size=3).max().item() + 1.0 + ``` + """ validate_bcs_input(input) validate_output(input, output) @@ -275,7 +482,38 @@ def morphological_laplace( cval: float = 0.0, origin: int | tuple[int, ...] = 0, ) -> Tensor: - """N-dimensional morphological Laplace for ``(B, C, Spatial...)`` CUDA tensors.""" + """Compute the morphological Laplacian of a grayscale tensor + + Computes ``dilation(input) + erosion(input) - 2 * input`` using ``float32`` + arithmetic. + + Args: + input (torch.Tensor): CUDA tensor with shape ``(B, C, Spatial...)``. + size (int or tuple[int, ...], optional): Shape of a flat, full + structuring element. + footprint (torch.Tensor, optional): Boolean mask selecting participating + positions. + structure (torch.Tensor, optional): Non-flat additive structuring + element. It takes precedence over ``footprint`` and ``size``. + output (torch.Tensor, optional): Preallocated result tensor with the + same shape and device as ``input``. + mode (str): Boundary mode: ``"reflect"``, ``"constant"``, ``"nearest"``, + ``"mirror"``, or ``"wrap"``. + cval (float): Boundary value used when ``mode="constant"``. + origin (int or tuple[int, ...]): Structuring-element anchor offset. + + Returns: + torch.Tensor: Morphological Laplacian in ``float32``, or ``output``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> x = torch.arange(25, device="cuda").reshape(1, 1, 5, 5) + >>> tm.morphological_laplace(x, size=3).dtype + torch.float32 + ``` + """ validate_bcs_input(input) validate_output(input, output) diff --git a/torchmorph/morphology/structure.py b/torchmorph/morphology/structure.py index 00fbde7..2f500e3 100644 --- a/torchmorph/morphology/structure.py +++ b/torchmorph/morphology/structure.py @@ -20,7 +20,7 @@ def generate_binary_structure(rank: int, connectivity: int) -> Tensor: Returns: torch.Tensor: Boolean tensor with shape ``(3,) * rank``. - Examples: + Example: ```pycon >>> import torch >>> import torchmorph as tm @@ -92,7 +92,38 @@ def iterate_structure( iterations: int, origin: int | tuple[int, ...] | None = None, ) -> Tensor | tuple[Tensor, list[int]]: - """Iterate a binary structure by dilating it with itself.""" + """Dilate a binary structuring element with itself repeatedly + + The result is equivalent to applying the original structuring element + ``iterations`` times in a morphology operation. Nonzero input values are + treated as part of the structure. + + Args: + structure (torch.Tensor): N-dimensional binary structuring element. + iterations (int): Number of copies to combine. Values below ``2`` + return a boolean clone of ``structure``. + origin (int or tuple[int, ...], optional): Original anchor offset. A + scalar is applied to every dimension. When supplied, the adjusted + origin is returned with the iterated structure. + + Returns: + torch.Tensor or tuple[torch.Tensor, list[int]]: Boolean iterated + structure. If ``origin`` is supplied, returns ``(structure, origin)`` + with each origin component multiplied by ``iterations``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> structure = tm.generate_binary_structure(2, 1) + >>> tm.iterate_structure(structure, 2).to(dtype=torch.int32) + tensor([[0, 0, 1, 0, 0], + [0, 1, 1, 1, 0], + [1, 1, 1, 1, 1], + [0, 1, 1, 1, 0], + [0, 0, 1, 0, 0]], dtype=torch.int32) + ``` + """ structure = structure != 0 if iterations < 2: diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py index 3652845..86315af 100644 --- a/torchmorph/optimal_transport.py +++ b/torchmorph/optimal_transport.py @@ -5,11 +5,29 @@ def build_cost_matrix(shape, p=2, device=None) -> Tensor: - """Pairwise L^p distance matrix between the points of a spatial grid. - - Returns a (d, d) cost matrix for the d flattened grid points of `shape`. - Flatten grid-shaped distributions to rows of an (n, d) tensor and pass the - result as `cost_matrix` to :class:`SinkhornSolver`. + """Build pairwise distances between points on a spatial grid + + The grid contains ``d = prod(shape)`` points in row-major order. Flatten + grid-shaped distributions to ``(n, d)`` before passing this matrix to + :class:`SinkhornSolver`. + + Args: + shape (sequence[int]): Spatial grid shape. + p (float): Norm degree passed to :func:`torch.cdist`. + device (torch.device or str, optional): Device for the returned matrix. + + Returns: + torch.Tensor: Symmetric ``float32`` cost matrix with shape ``(d, d)``. + + Example: + ```pycon + >>> import torchmorph as tm + >>> cost = tm.build_cost_matrix((2, 3)) + >>> cost.shape + torch.Size([6, 6]) + >>> cost[0, 1].item() + 1.0 + ``` """ coords = torch.stack( torch.meshgrid([torch.arange(s, device=device) for s in shape], indexing="ij"), dim=-1 @@ -151,19 +169,32 @@ def backward(ctx, grad_output): class SinkhornSolver(nn.Module): - """Entropy-regularized balanced optimal transport as a differentiable module. - - forward(source, target, cost_matrix=None) takes two (n, d) batches of - histograms sharing one (d, d) cost matrix and returns the (n,) transport - distances , with K = exp(-C / epsilon). Gradients w.r.t. source and - target are the centered dual potentials (envelope theorem), so the module - can be used directly as a loss. - - log_space=True runs the iterations entirely in the log domain, which is - numerically stable for small epsilon. The implementation is picked from - the input automatically: CUDA float32 tensors use the fused kernels, - everything else uses pure torch ops (where `threshold` > 0 enables early - stopping; the fused kernels always run `max_iter` iterations). + """Solve entropy-regularized balanced optimal transport + + The module accepts batches of flattened histograms and returns transport + costs that are differentiable with respect to both marginals. CUDA + ``float32`` inputs use fused kernels; other device and dtype combinations + use PyTorch operations. Set ``log_space=True`` for better stability at small + regularization values. + + Args: + epsilon (float): Positive entropy-regularization strength. + max_iter (int): Positive maximum number of Sinkhorn iterations. + threshold (float): Nonnegative early-stopping tolerance checked by the + PyTorch implementation. Fused CUDA kernels always run ``max_iter``. + p (float): Norm used by the default one-dimensional cost matrix. + log_space (bool): Run Sinkhorn iterations in the log domain. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> source = torch.tensor([[1.0, 0.0, 0.0]]) + >>> target = torch.tensor([[0.0, 0.0, 1.0]]) + >>> solver = tm.SinkhornSolver(epsilon=1.0, max_iter=200) + >>> solver(source, target).shape + torch.Size([1]) + ``` """ def __init__(self, epsilon=1.0, max_iter=100, threshold=0.0, p=2, log_space=False): @@ -186,12 +217,42 @@ def extra_repr(self): f"threshold={self.threshold}, log_space={self.log_space}" ) - def data_preprocess(self, source, target, cost_matrix: Optional[Tensor] = None): - """Validate (n, d) inputs, clamp negatives, and normalize each row to unit mass. - - All ops are differentiable, so forward gradients flow through the - normalization. When `cost_matrix` is None the d bins are treated as - points on a line. + def data_preprocess( + self, + source, + target, + cost_matrix: Optional[Tensor] = None, + ) -> tuple[Tensor, Tensor, Tensor]: + """Validate and normalize batched transport marginals + + Negative values are clamped to zero and every row is normalized to unit + mass. These operations remain in the autograd graph. If no cost matrix + is supplied, the ``d`` bins are treated as points on a line. + + Args: + source (torch.Tensor): Floating-point source marginals with shape + ``(n, d)``. + target (torch.Tensor): Floating-point target marginals with the same + shape as ``source``. + cost_matrix (torch.Tensor, optional): Shared pairwise cost matrix + with shape ``(d, d)``. + + Returns: + tuple[torch.Tensor, torch.Tensor, torch.Tensor]: Normalized source, + normalized target, and contiguous cost matrix on the source device + and with the source dtype. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> solver = tm.SinkhornSolver() + >>> a, b, cost = solver.data_preprocess( + ... torch.tensor([[1.0, 1.0]]), torch.tensor([[3.0, 1.0]]) + ... ) + >>> a.sum().item(), b.sum().item(), cost.shape + (1.0, 1.0, torch.Size([2, 2])) + ``` """ if source.ndim != 2 or source.shape != target.shape: raise ValueError("source and target must be (n, d) tensors of the same shape.") @@ -212,8 +273,34 @@ def data_preprocess(self, source, target, cost_matrix: Optional[Tensor] = None): target = target / target.sum(dim=-1, keepdim=True).clamp(min=1e-12) return source, target, cost_matrix - def forward(self, source: Tensor, target: Tensor, cost_matrix: Optional[Tensor] = None): - """(n,) transport distances between the rows of source and target.""" + def forward( + self, + source: Tensor, + target: Tensor, + cost_matrix: Optional[Tensor] = None, + ) -> Tensor: + """Compute transport costs between corresponding marginal rows + + Args: + source (torch.Tensor): Floating-point source marginals with shape + ``(n, d)``. + target (torch.Tensor): Floating-point target marginals with shape + ``(n, d)``. + cost_matrix (torch.Tensor, optional): Shared cost matrix with shape + ``(d, d)``. If ``None``, uses a one-dimensional grid cost. + + Returns: + torch.Tensor: Differentiable transport costs with shape ``(n,)``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> solver = tm.SinkhornSolver(max_iter=200) + >>> solver(torch.ones(2, 4), torch.ones(2, 4)).shape + torch.Size([2]) + ``` + """ source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix) return _SinkhornDistance.apply( source, @@ -226,17 +313,74 @@ def forward(self, source: Tensor, target: Tensor, cost_matrix: Optional[Tensor] ) @torch.no_grad() - def plan(self, source: Tensor, target: Tensor, cost_matrix: Optional[Tensor] = None): - """(n, d, d) transport plans, reconstructed in log space.""" + def plan( + self, + source: Tensor, + target: Tensor, + cost_matrix: Optional[Tensor] = None, + ) -> Tensor: + """Reconstruct transport plans in the log domain + + This method runs without gradient tracking. + + Args: + source (torch.Tensor): Floating-point source marginals with shape + ``(n, d)``. + target (torch.Tensor): Floating-point target marginals with shape + ``(n, d)``. + cost_matrix (torch.Tensor, optional): Shared cost matrix with shape + ``(d, d)``. + + Returns: + torch.Tensor: Transport plans with shape ``(n, d, d)``. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> solver = tm.SinkhornSolver(max_iter=200) + >>> solver.plan(torch.ones(1, 3), torch.ones(1, 3)).shape + torch.Size([1, 3, 3]) + ``` + """ source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix) log_u, log_v = self._iterate(source, target, cost_matrix) return torch.exp(log_u.unsqueeze(-1) - cost_matrix / self.epsilon + log_v.unsqueeze(-2)) @torch.no_grad() - def potentials(self, source: Tensor, target: Tensor, cost_matrix: Optional[Tensor] = None): - """Centered dual potentials (f, g), each (n, d). - - These are the gradients of the entropic OT cost w.r.t. the marginals. + def potentials( + self, + source: Tensor, + target: Tensor, + cost_matrix: Optional[Tensor] = None, + ) -> tuple[Tensor, Tensor]: + """Compute centered dual potentials for both marginals + + The potentials are the envelope-theorem gradients of the entropic + transport cost with respect to normalized marginals. This method runs + without gradient tracking. + + Args: + source (torch.Tensor): Floating-point source marginals with shape + ``(n, d)``. + target (torch.Tensor): Floating-point target marginals with shape + ``(n, d)``. + cost_matrix (torch.Tensor, optional): Shared cost matrix with shape + ``(d, d)``. + + Returns: + tuple[torch.Tensor, torch.Tensor]: Source and target potentials, + each with shape ``(n, d)`` and zero mean along the last axis. + + Example: + ```pycon + >>> import torch + >>> import torchmorph as tm + >>> solver = tm.SinkhornSolver(max_iter=200) + >>> f, g = solver.potentials(torch.ones(1, 3), torch.ones(1, 3)) + >>> f.shape, g.shape + (torch.Size([1, 3]), torch.Size([1, 3])) + ``` """ source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix) log_u, log_v = self._iterate(source, target, cost_matrix) From 08bec220c073ce7d4252e9fee4c263f78a687993 Mon Sep 17 00:00:00 2001 From: Yu Han Deng Date: Mon, 20 Jul 2026 16:32:52 +0800 Subject: [PATCH 05/16] fix HTML --- mkdocs.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/mkdocs.yml b/mkdocs.yml index e537064..cef34fc 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -33,6 +33,7 @@ plugins: docstring_style: google show_source: false show_root_heading: true + heading_level: 3 show_symbol_type_heading: true merge_init_into_class: true separate_signature: true From 75df059ca8671cd550046781c4eb48b9dc33e9bb Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Tue, 21 Jul 2026 18:47:53 +0800 Subject: [PATCH 06/16] Set documentation site URL --- mkdocs.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mkdocs.yml b/mkdocs.yml index cef34fc..cba0ea2 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -1,5 +1,5 @@ site_name: torchmorph -site_url: https://intcomp.org/tm/docs/ +site_url: https://intcomp.github.io/tm/ site_description: GPU-accelerated morphological operations for PyTorch repo_url: https://github.com/intcomp/torchmorph repo_name: intcomp/torchmorph From 182078d2dd23fc309faf35277b0faa6985a93079 Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Tue, 21 Jul 2026 19:01:47 +0800 Subject: [PATCH 07/16] Revert "Set documentation site URL" This reverts commit 75df059ca8671cd550046781c4eb48b9dc33e9bb. --- mkdocs.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mkdocs.yml b/mkdocs.yml index cba0ea2..cef34fc 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -1,5 +1,5 @@ site_name: torchmorph -site_url: https://intcomp.github.io/tm/ +site_url: https://intcomp.org/tm/docs/ site_description: GPU-accelerated morphological operations for PyTorch repo_url: https://github.com/intcomp/torchmorph repo_name: intcomp/torchmorph From 0836af3dc7a72bf329b15013f6b6eb07d41f26ca Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Tue, 21 Jul 2026 19:06:59 +0800 Subject: [PATCH 08/16] Reapply "Set documentation site URL" This reverts commit 182078d2dd23fc309faf35277b0faa6985a93079. --- mkdocs.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/mkdocs.yml b/mkdocs.yml index cef34fc..cba0ea2 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -1,5 +1,5 @@ site_name: torchmorph -site_url: https://intcomp.org/tm/docs/ +site_url: https://intcomp.github.io/tm/ site_description: GPU-accelerated morphological operations for PyTorch repo_url: https://github.com/intcomp/torchmorph repo_name: intcomp/torchmorph From e7d130aafc228dc3b59f58c764f68d49384c0bf8 Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Tue, 21 Jul 2026 19:16:37 +0800 Subject: [PATCH 09/16] Split morphology API documentation --- docs/api/binary_morphology.md | 18 +++++++++++++ docs/api/grayscale_morphology.md | 20 ++++++++++++++ docs/api/morphological_ops.md | 45 -------------------------------- docs/api/structuring_elements.md | 8 ++++++ docs/index.md | 4 ++- mkdocs.yml | 4 ++- 6 files changed, 52 insertions(+), 47 deletions(-) create mode 100644 docs/api/binary_morphology.md create mode 100644 docs/api/grayscale_morphology.md delete mode 100644 docs/api/morphological_ops.md create mode 100644 docs/api/structuring_elements.md diff --git a/docs/api/binary_morphology.md b/docs/api/binary_morphology.md new file mode 100644 index 0000000..c015eb5 --- /dev/null +++ b/docs/api/binary_morphology.md @@ -0,0 +1,18 @@ +# Binary Morphology + +Binary operators treat nonzero values as foreground and expect tensors shaped +as `(B, C, Spatial...)`. + +::: torchmorph.binary_dilation + +::: torchmorph.binary_erosion + +::: torchmorph.binary_propagation + +::: torchmorph.binary_fill_holes + +::: torchmorph.binary_hit_or_miss + +::: torchmorph.binary_opening + +::: torchmorph.binary_closing diff --git a/docs/api/grayscale_morphology.md b/docs/api/grayscale_morphology.md new file mode 100644 index 0000000..5ff72e7 --- /dev/null +++ b/docs/api/grayscale_morphology.md @@ -0,0 +1,20 @@ +# Grayscale Morphology + +Grayscale operators process floating-point or integer-valued tensors shaped as +`(B, C, Spatial...)`. + +::: torchmorph.grey_dilation + +::: torchmorph.grey_erosion + +::: torchmorph.grey_opening + +::: torchmorph.grey_closing + +::: torchmorph.morphological_gradient + +::: torchmorph.morphological_laplace + +::: torchmorph.white_tophat + +::: torchmorph.black_tophat diff --git a/docs/api/morphological_ops.md b/docs/api/morphological_ops.md deleted file mode 100644 index 900bea5..0000000 --- a/docs/api/morphological_ops.md +++ /dev/null @@ -1,45 +0,0 @@ -# Morphological Operations - -All image and volume operators on this page expect CUDA tensors shaped as -`(B, C, Spatial...)`. Add batch and channel dimensions before processing a -single image or volume. - -## Structuring elements - -::: torchmorph.generate_binary_structure - -::: torchmorph.iterate_structure - -## Binary morphology - -::: torchmorph.binary_dilation - -::: torchmorph.binary_erosion - -::: torchmorph.binary_propagation - -::: torchmorph.binary_fill_holes - -::: torchmorph.binary_hit_or_miss - -::: torchmorph.binary_opening - -::: torchmorph.binary_closing - -## Grayscale morphology - -::: torchmorph.grey_dilation - -::: torchmorph.grey_erosion - -::: torchmorph.grey_opening - -::: torchmorph.grey_closing - -::: torchmorph.morphological_gradient - -::: torchmorph.morphological_laplace - -::: torchmorph.white_tophat - -::: torchmorph.black_tophat diff --git a/docs/api/structuring_elements.md b/docs/api/structuring_elements.md new file mode 100644 index 0000000..36e0493 --- /dev/null +++ b/docs/api/structuring_elements.md @@ -0,0 +1,8 @@ +# Structuring Elements + +Create and expand multidimensional connectivity structures for morphological +operations. + +::: torchmorph.generate_binary_structure + +::: torchmorph.iterate_structure diff --git a/docs/index.md b/docs/index.md index 8c377c1..f58996c 100644 --- a/docs/index.md +++ b/docs/index.md @@ -39,5 +39,7 @@ distances = solver(source, target, cost) ## API Reference - [Distance Transforms](api/distance_transforms.md) — Euclidean, Chamfer, and brute-force distance transforms -- [Morphological Operations](api/morphological_ops.md) — Structuring elements, binary morphology, and grayscale morphology +- [Structuring Elements](api/structuring_elements.md) — Multidimensional connectivity structures +- [Binary Morphology](api/binary_morphology.md) — Binary dilation, erosion, filling, and related operators +- [Grayscale Morphology](api/grayscale_morphology.md) — Grayscale dilation, erosion, gradients, and top-hats - [Optimal Transport](api/optimal_transport.md) — Cost matrices and differentiable Sinkhorn transport diff --git a/mkdocs.yml b/mkdocs.yml index cba0ea2..a7fc5ef 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -51,5 +51,7 @@ nav: - Home: index.md - API Reference: - Distance Transforms: api/distance_transforms.md - - Morphological Operations: api/morphological_ops.md + - Structuring Elements: api/structuring_elements.md + - Binary Morphology: api/binary_morphology.md + - Grayscale Morphology: api/grayscale_morphology.md - Optimal Transport: api/optimal_transport.md From e5c8cbc06e7f33ca94573879b958a8a86dfb2f51 Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Tue, 21 Jul 2026 19:24:11 +0800 Subject: [PATCH 10/16] Limit API table of contents depth --- mkdocs.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/mkdocs.yml b/mkdocs.yml index a7fc5ef..c66b14e 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -40,6 +40,8 @@ plugins: line_length: 88 markdown_extensions: + - toc: + toc_depth: 3 - pymdownx.highlight: anchor_linenums: true - pymdownx.superfences From c993feaa64a53f9b0664bc5cd9005c1d9d0477ce Mon Sep 17 00:00:00 2001 From: Yuhandeng Date: Mon, 27 Jul 2026 18:03:38 +0800 Subject: [PATCH 11/16] Add short API anchors and URLs --- docs/{api => }/binary_morphology.md | 0 docs/{api => }/distance_transforms.md | 0 docs/{api => }/grayscale_morphology.md | 0 docs/index.md | 10 +++++----- docs/{api => }/optimal_transport.md | 0 docs/{api => }/structuring_elements.md | 0 docs/stylesheets/extra.css | 6 ++++++ mkdocs.yml | 13 ++++++++----- mkdocs_hooks.py | 16 ++++++++++++++++ 9 files changed, 35 insertions(+), 10 deletions(-) rename docs/{api => }/binary_morphology.md (100%) rename docs/{api => }/distance_transforms.md (100%) rename docs/{api => }/grayscale_morphology.md (100%) rename docs/{api => }/optimal_transport.md (100%) rename docs/{api => }/structuring_elements.md (100%) create mode 100644 mkdocs_hooks.py diff --git a/docs/api/binary_morphology.md b/docs/binary_morphology.md similarity index 100% rename from docs/api/binary_morphology.md rename to docs/binary_morphology.md diff --git a/docs/api/distance_transforms.md b/docs/distance_transforms.md similarity index 100% rename from docs/api/distance_transforms.md rename to docs/distance_transforms.md diff --git a/docs/api/grayscale_morphology.md b/docs/grayscale_morphology.md similarity index 100% rename from docs/api/grayscale_morphology.md rename to docs/grayscale_morphology.md diff --git a/docs/index.md b/docs/index.md index f58996c..bc7e9b6 100644 --- a/docs/index.md +++ b/docs/index.md @@ -38,8 +38,8 @@ distances = solver(source, target, cost) ## API Reference -- [Distance Transforms](api/distance_transforms.md) — Euclidean, Chamfer, and brute-force distance transforms -- [Structuring Elements](api/structuring_elements.md) — Multidimensional connectivity structures -- [Binary Morphology](api/binary_morphology.md) — Binary dilation, erosion, filling, and related operators -- [Grayscale Morphology](api/grayscale_morphology.md) — Grayscale dilation, erosion, gradients, and top-hats -- [Optimal Transport](api/optimal_transport.md) — Cost matrices and differentiable Sinkhorn transport +- [Distance Transforms](distance_transforms.md) — Euclidean, Chamfer, and brute-force distance transforms +- [Structuring Elements](structuring_elements.md) — Multidimensional connectivity structures +- [Binary Morphology](binary_morphology.md) — Binary dilation, erosion, filling, and related operators +- [Grayscale Morphology](grayscale_morphology.md) — Grayscale dilation, erosion, gradients, and top-hats +- [Optimal Transport](optimal_transport.md) — Cost matrices and differentiable Sinkhorn transport diff --git a/docs/api/optimal_transport.md b/docs/optimal_transport.md similarity index 100% rename from docs/api/optimal_transport.md rename to docs/optimal_transport.md diff --git a/docs/api/structuring_elements.md b/docs/structuring_elements.md similarity index 100% rename from docs/api/structuring_elements.md rename to docs/structuring_elements.md diff --git a/docs/stylesheets/extra.css b/docs/stylesheets/extra.css index dbc24f9..812fb09 100644 --- a/docs/stylesheets/extra.css +++ b/docs/stylesheets/extra.css @@ -55,6 +55,12 @@ border-left: 0.12rem solid var(--md-default-fg-color--lightest); } +.md-typeset .doc-short-anchor { + display: block; + height: 0; + scroll-margin-top: 4.8rem; +} + @media screen and (max-width: 44.9844em) { .md-typeset .doc-heading { margin-top: 1.8rem; diff --git a/mkdocs.yml b/mkdocs.yml index c66b14e..3a581e9 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -24,6 +24,9 @@ theme: extra_css: - stylesheets/extra.css +hooks: + - mkdocs_hooks.py + plugins: - search - mkdocstrings: @@ -52,8 +55,8 @@ markdown_extensions: nav: - Home: index.md - API Reference: - - Distance Transforms: api/distance_transforms.md - - Structuring Elements: api/structuring_elements.md - - Binary Morphology: api/binary_morphology.md - - Grayscale Morphology: api/grayscale_morphology.md - - Optimal Transport: api/optimal_transport.md + - Distance Transforms: distance_transforms.md + - Structuring Elements: structuring_elements.md + - Binary Morphology: binary_morphology.md + - Grayscale Morphology: grayscale_morphology.md + - Optimal Transport: optimal_transport.md diff --git a/mkdocs_hooks.py b/mkdocs_hooks.py new file mode 100644 index 0000000..09ebeab --- /dev/null +++ b/mkdocs_hooks.py @@ -0,0 +1,16 @@ +"""Small build-time enhancements for the generated API reference.""" + +import re + +_API_HEADING = re.compile(r'(]*>)') + + +def on_page_content(html: str, **kwargs) -> str: + """Add a short anchor alias before every mkdocstrings API heading.""" + + def add_short_anchor(match: re.Match[str]) -> str: + short_name = match.group(2).rsplit(".", 1)[-1] + anchor = f'' + return f"{anchor}\n{match.group(1)}" + + return _API_HEADING.sub(add_short_anchor, html) From d6cd11caf44b62b2c09bb6244288714d7bb876b9 Mon Sep 17 00:00:00 2001 From: Kai ZHAO <694691@qq.com> Date: Wed, 26 Aug 2026 20:02:19 +0800 Subject: [PATCH 12/16] Add arXiv badge to documentation index Co-Authored-By: Claude Opus 5 (1M context) --- docs/index.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/index.md b/docs/index.md index bc7e9b6..fae519e 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,5 +1,7 @@ # torchmorph +[![arXiv](https://img.shields.io/badge/arXiv-2608.24738-b31b1b.svg)](https://arxiv.org/abs/2608.24738) + GPU-accelerated morphological image processing operations for PyTorch. ## Installation From 831153310f28bb157128878938433431f25dcd2b Mon Sep 17 00:00:00 2001 From: Kai ZHAO <694691@qq.com> Date: Wed, 26 Aug 2026 20:03:38 +0800 Subject: [PATCH 13/16] Add arXiv badge to README Co-Authored-By: Claude Opus 5 (1M context) --- README.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/README.md b/README.md index ede098a..a8da672 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,7 @@ # TorchMorph +[![arXiv](https://img.shields.io/badge/arXiv-2608.24738-b31b1b.svg)](https://arxiv.org/abs/2608.24738) + > ⚡ CUDA-Accelerated, Batch-Parallel Morphological Transformations for PyTorch TorchMorph is a lightweight, extensible library that brings **GPU-accelerated morphological operations** into the PyTorch ecosystem. @@ -77,3 +79,10 @@ if torch.cuda.is_available(): else: print("CUDA not available; install verified for import only.") ``` + +--- + +## 📄 License + +TorchMorph is released under the [MIT License](LICENSE) — free for academic and +commercial use, provided "as is" without warranty of any kind. From b2b4893a0dd588bb3c06d68b083d1988ebe4f45a Mon Sep 17 00:00:00 2001 From: Kai ZHAO <694691@qq.com> Date: Wed, 26 Aug 2026 20:21:56 +0800 Subject: [PATCH 14/16] LICENSE --- LICENSE | 21 +++++++++++++++++++++ README.md | 1 + docs/index.md | 1 + setup.py | 2 ++ 4 files changed, 25 insertions(+) create mode 100644 LICENSE diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..6eaaaae --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Kai Zhao + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md index a8da672..e077617 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,7 @@ # TorchMorph [![arXiv](https://img.shields.io/badge/arXiv-2608.24738-b31b1b.svg)](https://arxiv.org/abs/2608.24738) +[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Paper-FFD21E.svg)](https://huggingface.co/papers/2608.24738) > ⚡ CUDA-Accelerated, Batch-Parallel Morphological Transformations for PyTorch diff --git a/docs/index.md b/docs/index.md index fae519e..2d21f34 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,6 +1,7 @@ # torchmorph [![arXiv](https://img.shields.io/badge/arXiv-2608.24738-b31b1b.svg)](https://arxiv.org/abs/2608.24738) +[![Hugging Face](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Paper-FFD21E.svg)](https://huggingface.co/papers/2608.24738) GPU-accelerated morphological image processing operations for PyTorch. diff --git a/setup.py b/setup.py index ad0712d..1a52d5e 100644 --- a/setup.py +++ b/setup.py @@ -26,6 +26,8 @@ def get_extensions(): version="0.1.0", author="Kai Zhao", description="CUDA-accelerated morphological transformations for PyTorch", + license="MIT", + license_files=["LICENSE"], url="https://github.com/kaizhao-shu/torchmorph", packages=find_packages(), python_requires=">=3.12", From a37b03c455b77639b47c0747707ce919553bccd9 Mon Sep 17 00:00:00 2001 From: Yiping Mao <1498743918@qq.com> Date: Fri, 16 Oct 2026 19:12:08 +0800 Subject: [PATCH 15/16] perf: reuse cached CUDA graphs for Sinkhorn iterations --- torchmorph/optimal_transport.py | 89 +++++++++++++++++++++------------ 1 file changed, 57 insertions(+), 32 deletions(-) diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py index 86315af..2591983 100644 --- a/torchmorph/optimal_transport.py +++ b/torchmorph/optimal_transport.py @@ -1,3 +1,4 @@ +from threading import local from typing import Optional import torch @@ -49,41 +50,61 @@ def _transpose_or_self(matrix: Tensor) -> Tensor: return matrix if torch.equal(matrix, matrix.mT) else matrix.mT.contiguous() -_GRAPH_MIN_ITER = 100 # below this, plain launches beat the graph-capture setup cost +_GRAPH_MIN_ITER = 100 _GRAPH_CHUNK = 25 +_GRAPH_CACHE = local() # one workspace per thread; changing configuration replaces it -def _run_fused(launch, max_iter, device): - """Run `launch(k)`, which enqueues k in-place Sinkhorn iterations, max_iter times in total. +def _run_fused(launch, max_iter, device, inputs, outputs, signature): + """Replay a cached chunk on fixed buffers, then copy out this call's result. - Long runs are launch-latency-bound (two tiny kernels per iteration), so a - chunk of iterations is captured into a CUDA graph once and replayed. The - iterations are fixed-point steps on ping-pong buffers, so a fallback that - runs extra iterations is always safe. + `launch(inputs, outputs, k)` enqueues k in-place iterations. Only the most + recent configuration is cached per thread, with streams isolated by the + cache key. Warmup uses scratch state and never counts as solver iterations. + Errors propagate instead of silently rerunning an already modified state. """ - if max_iter < _GRAPH_MIN_ITER or torch.cuda.is_current_stream_capturing(): - launch(max_iter) - return - try: - with torch.cuda.device(device): - # Warm up on a side stream (required before capture), doing the - # first chunk of real iterations in the process. + with torch.cuda.device(device): + if max_iter < _GRAPH_MIN_ITER or torch.cuda.is_current_stream_capturing(): + launch(inputs, outputs, max_iter) + return + + stream = torch.cuda.current_stream() + key = ( + device, stream.cuda_stream, signature, + tuple((tuple(t.shape), t.dtype, t.stride()) for t in (*inputs, *outputs)), + ) + entry = getattr(_GRAPH_CACHE, "entry", None) + if entry is None or entry[0] != key: + # Release the old workspace before allocating its replacement. + _GRAPH_CACHE.entry = None + entry = None + static_inputs = tuple(t.detach().clone() for t in inputs) + static_outputs = tuple(t.clone() for t in outputs) side_stream = torch.cuda.Stream() - side_stream.wait_stream(torch.cuda.current_stream()) + side_stream.wait_stream(stream) with torch.cuda.stream(side_stream): - launch(_GRAPH_CHUNK) - torch.cuda.current_stream().wait_stream(side_stream) + launch(static_inputs, static_outputs, _GRAPH_CHUNK) + stream.wait_stream(side_stream) graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - launch(_GRAPH_CHUNK) - for _ in range(max_iter // _GRAPH_CHUNK - 1): - graph.replay() - remainder = max_iter % _GRAPH_CHUNK - if remainder: - launch(remainder) - except RuntimeError: - launch(max_iter) # graphs unavailable; extra iterations are harmless + with torch.cuda.graph(graph, stream=side_stream): + launch(static_inputs, static_outputs, _GRAPH_CHUNK) + _GRAPH_CACHE.entry = (key, graph, static_inputs, static_outputs) + else: + _, graph, static_inputs, static_outputs = entry + for dst, src in zip(static_inputs, inputs): + dst.copy_(src) + + # Discard warmup or the previous call's solution before any replay. + for dst, src in zip(static_outputs, outputs): + dst.copy_(src) + for _ in range(max_iter // _GRAPH_CHUNK): + graph.replay() + remainder = max_iter % _GRAPH_CHUNK + if remainder: + launch(static_inputs, static_outputs, remainder) + for dst, src in zip(outputs, static_outputs): + dst.copy_(src) def _scaling_iterations(a, b, cost_matrix, epsilon, max_iter, threshold): @@ -96,7 +117,10 @@ def _scaling_iterations(a, b, cost_matrix, epsilon, max_iter, threshold): K_T = _transpose_or_self(K) u = torch.ones_like(a) v = torch.ones_like(b) - _run_fused(lambda k: _C.sinkhorn_fastiter(a, b, K, K_T, u, v, k), max_iter, a.device) + _run_fused( + lambda ins, outs, k: _C.sinkhorn_fastiter(*ins, *outs, k), + max_iter, a.device, (a, b, K, K_T), (u, v), "scaling", + ) return u.log(), v.log() u = torch.ones_like(a) @@ -122,11 +146,9 @@ def _log_iterations(a, b, cost_matrix, epsilon, max_iter, threshold): log_u = torch.zeros_like(a) log_v = torch.zeros_like(b) _run_fused( - lambda k: _C.sinkhorn_logiter( - log_a, log_b, cost_matrix, cost_T, log_u, log_v, k, epsilon - ), - max_iter, - a.device, + lambda ins, outs, k: _C.sinkhorn_logiter(*ins, *outs, k, epsilon), + max_iter, a.device, (log_a, log_b, cost_matrix, cost_T), + (log_u, log_v), ("log", epsilon), ) return log_u, log_v @@ -177,6 +199,9 @@ class SinkhornSolver(nn.Module): use PyTorch operations. Set ``log_space=True`` for better stability at small regularization values. + Long fused runs reuse a CUDA graph with fixed scratch buffers. Each thread + caches only its most recent device, stream, shape, and kernel configuration. + Args: epsilon (float): Positive entropy-regularization strength. max_iter (int): Positive maximum number of Sinkhorn iterations. From 983409206a2595836a4c9abf4c5e7482288336b3 Mon Sep 17 00:00:00 2001 From: Yiping Mao <1498743918@qq.com> Date: Wed, 21 Oct 2026 20:00:24 +0800 Subject: [PATCH 16/16] Revert "perf: reuse cached CUDA graphs for Sinkhorn iterations" This reverts commit a37b03c455b77639b47c0747707ce919553bccd9. --- torchmorph/optimal_transport.py | 89 ++++++++++++--------------------- 1 file changed, 32 insertions(+), 57 deletions(-) diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py index 2591983..86315af 100644 --- a/torchmorph/optimal_transport.py +++ b/torchmorph/optimal_transport.py @@ -1,4 +1,3 @@ -from threading import local from typing import Optional import torch @@ -50,61 +49,41 @@ def _transpose_or_self(matrix: Tensor) -> Tensor: return matrix if torch.equal(matrix, matrix.mT) else matrix.mT.contiguous() -_GRAPH_MIN_ITER = 100 +_GRAPH_MIN_ITER = 100 # below this, plain launches beat the graph-capture setup cost _GRAPH_CHUNK = 25 -_GRAPH_CACHE = local() # one workspace per thread; changing configuration replaces it -def _run_fused(launch, max_iter, device, inputs, outputs, signature): - """Replay a cached chunk on fixed buffers, then copy out this call's result. +def _run_fused(launch, max_iter, device): + """Run `launch(k)`, which enqueues k in-place Sinkhorn iterations, max_iter times in total. - `launch(inputs, outputs, k)` enqueues k in-place iterations. Only the most - recent configuration is cached per thread, with streams isolated by the - cache key. Warmup uses scratch state and never counts as solver iterations. - Errors propagate instead of silently rerunning an already modified state. + Long runs are launch-latency-bound (two tiny kernels per iteration), so a + chunk of iterations is captured into a CUDA graph once and replayed. The + iterations are fixed-point steps on ping-pong buffers, so a fallback that + runs extra iterations is always safe. """ - with torch.cuda.device(device): - if max_iter < _GRAPH_MIN_ITER or torch.cuda.is_current_stream_capturing(): - launch(inputs, outputs, max_iter) - return - - stream = torch.cuda.current_stream() - key = ( - device, stream.cuda_stream, signature, - tuple((tuple(t.shape), t.dtype, t.stride()) for t in (*inputs, *outputs)), - ) - entry = getattr(_GRAPH_CACHE, "entry", None) - if entry is None or entry[0] != key: - # Release the old workspace before allocating its replacement. - _GRAPH_CACHE.entry = None - entry = None - static_inputs = tuple(t.detach().clone() for t in inputs) - static_outputs = tuple(t.clone() for t in outputs) + if max_iter < _GRAPH_MIN_ITER or torch.cuda.is_current_stream_capturing(): + launch(max_iter) + return + try: + with torch.cuda.device(device): + # Warm up on a side stream (required before capture), doing the + # first chunk of real iterations in the process. side_stream = torch.cuda.Stream() - side_stream.wait_stream(stream) + side_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(side_stream): - launch(static_inputs, static_outputs, _GRAPH_CHUNK) - stream.wait_stream(side_stream) + launch(_GRAPH_CHUNK) + torch.cuda.current_stream().wait_stream(side_stream) graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph, stream=side_stream): - launch(static_inputs, static_outputs, _GRAPH_CHUNK) - _GRAPH_CACHE.entry = (key, graph, static_inputs, static_outputs) - else: - _, graph, static_inputs, static_outputs = entry - for dst, src in zip(static_inputs, inputs): - dst.copy_(src) - - # Discard warmup or the previous call's solution before any replay. - for dst, src in zip(static_outputs, outputs): - dst.copy_(src) - for _ in range(max_iter // _GRAPH_CHUNK): - graph.replay() - remainder = max_iter % _GRAPH_CHUNK - if remainder: - launch(static_inputs, static_outputs, remainder) - for dst, src in zip(outputs, static_outputs): - dst.copy_(src) + with torch.cuda.graph(graph): + launch(_GRAPH_CHUNK) + for _ in range(max_iter // _GRAPH_CHUNK - 1): + graph.replay() + remainder = max_iter % _GRAPH_CHUNK + if remainder: + launch(remainder) + except RuntimeError: + launch(max_iter) # graphs unavailable; extra iterations are harmless def _scaling_iterations(a, b, cost_matrix, epsilon, max_iter, threshold): @@ -117,10 +96,7 @@ def _scaling_iterations(a, b, cost_matrix, epsilon, max_iter, threshold): K_T = _transpose_or_self(K) u = torch.ones_like(a) v = torch.ones_like(b) - _run_fused( - lambda ins, outs, k: _C.sinkhorn_fastiter(*ins, *outs, k), - max_iter, a.device, (a, b, K, K_T), (u, v), "scaling", - ) + _run_fused(lambda k: _C.sinkhorn_fastiter(a, b, K, K_T, u, v, k), max_iter, a.device) return u.log(), v.log() u = torch.ones_like(a) @@ -146,9 +122,11 @@ def _log_iterations(a, b, cost_matrix, epsilon, max_iter, threshold): log_u = torch.zeros_like(a) log_v = torch.zeros_like(b) _run_fused( - lambda ins, outs, k: _C.sinkhorn_logiter(*ins, *outs, k, epsilon), - max_iter, a.device, (log_a, log_b, cost_matrix, cost_T), - (log_u, log_v), ("log", epsilon), + lambda k: _C.sinkhorn_logiter( + log_a, log_b, cost_matrix, cost_T, log_u, log_v, k, epsilon + ), + max_iter, + a.device, ) return log_u, log_v @@ -199,9 +177,6 @@ class SinkhornSolver(nn.Module): use PyTorch operations. Set ``log_space=True`` for better stability at small regularization values. - Long fused runs reuse a CUDA graph with fixed scratch buffers. Each thread - caches only its most recent device, stream, shape, and kernel configuration. - Args: epsilon (float): Positive entropy-regularization strength. max_iter (int): Positive maximum number of Sinkhorn iterations.