Correctness: opt-in deterministic conv, pruning empty-mask guard, TensorField backprop regression test

MinkowskiEngine/MinkowskiConvolution.py

@@ -314,6 +314,21 @@ def forward(
         assert isinstance(input, SparseTensor)
         assert input.D == self.dimension
 
+        # Opt-in deterministic path (ME.set_deterministic(True)): sort the input
+        # into a canonical coordinate order so the backend accumulation order --
+        # and therefore the output -- is independent of the input row ordering.
+        # See MinkowskiEngine/utils/determinism.py and upstream issues #554/#504.
+        # Skipped for the use_mm (kernel_volume == 1) path, which is a plain
+        # matrix multiply with no coordinate-order-dependent accumulation.
+        if not self.use_mm:
+            from MinkowskiEngine.utils.determinism import (
+                is_deterministic,
+                sorted_coordinates,
+            )
+
+            if is_deterministic():
+                input = sorted_coordinates(input)
+
         if self.use_mm:
             # If the kernel_size == 1, the convolution is simply a matrix
             # multiplication

MinkowskiEngine/__init__.py

@@ -208,6 +208,12 @@
 
 import MinkowskiEngine.utils as utils
 
+from MinkowskiEngine.utils.determinism import (
+    set_deterministic,
+    is_deterministic,
+    sorted_coordinates,
+)
+
 import MinkowskiEngine.modules as modules
 
 from sparse_matrix_functions import (

MinkowskiEngine/utils/__init__.py

@@ -25,4 +25,5 @@
 from .collation import SparseCollation, batched_coordinates, sparse_collate, batch_sparse_collate
 # from .coords import get_coords_map
 from .init import kaiming_normal_
-from .summary import summary
+from .summary import summary
+from .determinism import set_deterministic, is_deterministic, sorted_coordinates

MinkowskiEngine/utils/determinism.py

@@ -0,0 +1,154 @@
+# Copyright (c) Chris Choy (chrischoy@ai.stanford.edu).
+#
+# 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.
+r"""Opt-in deterministic convolution helpers.
+
+MinkowskiEngine convolutions accumulate contributions with ``atomicAdd`` and
+dispatch the kernel map in an order that depends on how the input coordinates
+were hashed / inserted.  As a consequence, running the *same* convolution twice
+on the *same* set of points but in a different row order can produce outputs
+that differ in the low-order bits.  The effect is largest for strided
+(``stride > 1``) and transposed convolutions.  See upstream issues
+https://github.com/NVIDIA/MinkowskiEngine/issues/554 and
+https://github.com/NVIDIA/MinkowskiEngine/issues/504.
+
+The community-confirmed workaround is to sort the coordinates into a canonical
+order before the convolution so that the insertion order (and therefore the
+accumulation order) becomes a deterministic function of the point *set* rather
+than of the incoming row order.  This module exposes that as:
+
+* :func:`sorted_coordinates` -- return a new :class:`SparseTensor` whose rows
+  are in a canonical lexicographic coordinate order, and
+* :func:`set_deterministic` / :func:`is_deterministic` -- a process-wide flag
+  that, when enabled, makes :class:`MinkowskiConvolution` and
+  :class:`MinkowskiConvolutionTranspose` route their input through
+  :func:`sorted_coordinates` before calling the backend.
+
+This is *opt-in*: the default behavior is unchanged.
+"""
+import torch
+
+from MinkowskiSparseTensor import SparseTensor
+
+__all__ = [
+    "set_deterministic",
+    "is_deterministic",
+    "sorted_coordinates",
+]
+
+# Process-wide opt-in flag. Off by default; toggled via set_deterministic().
+_DETERMINISTIC = False
+
+
+def set_deterministic(mode: bool = True):
+    r"""Enable or disable opt-in deterministic convolution behavior.
+
+    When enabled, :class:`MinkowskiEngine.MinkowskiConvolution` and
+    :class:`MinkowskiEngine.MinkowskiConvolutionTranspose` sort their input into
+    a canonical coordinate order (via :func:`sorted_coordinates`) before calling
+    the backend, so that the convolution output is a deterministic function of
+    the input *point set*, independent of the input row ordering.
+
+    .. warning::
+
+        This costs one coordinate sort **per convolution** (a handful of stable
+        ``argsort`` passes plus one ``SparseTensor`` reconstruction). It is
+        intended for reproducibility / debugging, not for throughput-critical
+        training or inference. The default (``False``) preserves the original,
+        faster, order-dependent behavior.
+
+    Args:
+        :attr:`mode` (bool): ``True`` to enable deterministic convolutions,
+        ``False`` to restore the default behavior. Defaults to ``True``.
+    """
+    global _DETERMINISTIC
+    _DETERMINISTIC = bool(mode)
+
+
+def is_deterministic() -> bool:
+    r"""Return whether opt-in deterministic convolution behavior is enabled."""
+    return _DETERMINISTIC
+
+
+def _lexicographic_permutation(coordinates: torch.Tensor) -> torch.Tensor:
+    r"""Return a permutation that lexicographically sorts ``coordinates`` rows.
+
+    Column 0 (the batch index) is the most significant key and the last spatial
+    column is the least significant. The sort is implemented as a stable radix
+    sort -- a sequence of per-column stable ``argsort`` passes from the least to
+    the most significant column. This avoids packing the coordinates into a
+    single integer rank (which would overflow for large or high-dimensional
+    coordinates) and is correct for arbitrary spatial dimension ``D`` and for
+    negative coordinates (signed integer comparison).
+    """
+    n, ncol = coordinates.shape
+    perm = torch.arange(n, device=coordinates.device)
+    # Least significant column first, most significant (batch) column last, so
+    # that the final stable pass leaves the batch column dominant.
+    for col in reversed(range(ncol)):
+        keys = coordinates[perm, col]
+        order = torch.argsort(keys, stable=True)
+        perm = perm[order]
+    return perm
+
+
+def sorted_coordinates(sparse_tensor: SparseTensor) -> SparseTensor:
+    r"""Return a copy of ``sparse_tensor`` with rows in canonical coordinate order.
+
+    The returned :class:`SparseTensor` contains the same points and features as
+    the input, but its rows are sorted by a lexicographic ranking of the
+    coordinates (batch index most significant). Because the row (insertion)
+    order is now a deterministic function of the point set, feeding the result
+    into a convolution removes the order-dependent low-order-bit nondeterminism
+    described in the module docstring.
+
+    The reconstruction preserves the input ``tensor_stride`` and
+    ``quantization_mode`` and builds a fresh coordinate manager on the same
+    device. Gradients flow through the reordered features.
+
+    Args:
+        :attr:`sparse_tensor` (:class:`MinkowskiEngine.SparseTensor`): the tensor
+        to reorder.
+
+    Returns:
+        A new :class:`MinkowskiEngine.SparseTensor` with canonically ordered rows.
+    """
+    assert isinstance(
+        sparse_tensor, SparseTensor
+    ), "sorted_coordinates expects a SparseTensor"
+
+    coords = sparse_tensor.C
+    feats = sparse_tensor.F
+
+    # Nothing to reorder for empty / single-row tensors; return as-is.
+    if coords.shape[0] <= 1:
+        return sparse_tensor
+
+    with torch.no_grad():
+        perm = _lexicographic_permutation(coords)
+
+    sorted_coords = coords[perm]
+    sorted_feats = feats[perm]
+
+    return SparseTensor(
+        sorted_feats,
+        coordinates=sorted_coords,
+        tensor_stride=sparse_tensor.tensor_stride,
+        quantization_mode=sparse_tensor.quantization_mode,
+    )

tests/python/determinism.py

@@ -0,0 +1,173 @@
+# Copyright (c) Chris Choy (chrischoy@ai.stanford.edu).
+#
+# 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.
+#
+# Regression tests for the opt-in deterministic convolution path
+# (ME.set_deterministic / ME.utils.sorted_coordinates). See upstream issues
+# https://github.com/NVIDIA/MinkowskiEngine/issues/554 and
+# https://github.com/NVIDIA/MinkowskiEngine/issues/504.
+import unittest
+
+import numpy as np
+import torch
+
+import MinkowskiEngine as ME
+from MinkowskiEngine import (
+    SparseTensor,
+    MinkowskiConvolution,
+    MinkowskiConvolutionTranspose,
+)
+
+
+def _canonical_feats(sparse_tensor):
+    r"""Return the features of ``sparse_tensor`` reordered into a canonical
+    lexicographic coordinate order so that outputs built from differently
+    permuted inputs can be compared row-for-row."""
+    C = sparse_tensor.C.cpu().numpy()
+    F = sparse_tensor.F.detach().cpu()
+    # np.lexsort uses the last key as primary; we want column 0 (batch) primary.
+    keys = tuple(C[:, i] for i in reversed(range(C.shape[1])))
+    order = np.lexsort(keys)
+    return F[order].contiguous(), C[order]
+
+
+def _make_point_set(N, D, in_channels, seed):
+    r"""Build a unique batched coordinate/feature set (2 batches)."""
+    g = torch.Generator().manual_seed(seed)
+    coords = torch.randint(0, 24, (N, D), generator=g)
+    batch = torch.randint(0, 2, (N, 1), generator=g)
+    full = torch.cat([batch, coords], dim=1)
+    full = torch.unique(full, dim=0).int()
+    feats = torch.randn(full.shape[0], in_channels, generator=g)
+    return full, feats
+
+
+class TestDeterminism(unittest.TestCase):
+    def tearDown(self):
+        # Never leak the global flag into other tests.
+        ME.set_deterministic(False)
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_sorted_coordinates_helper(self):
+        # Handles arbitrary D and negative coordinates, keeps feats aligned.
+        coords = torch.IntTensor(
+            [[1, -2, 3, 0], [0, 5, -1, 2], [0, -3, 4, 4], [1, -2, 3, -1]]
+        )
+        feats = torch.arange(4 * 2).view(4, 2).float()
+        st = SparseTensor(feats, coords, device="cuda")
+        srt = ME.sorted_coordinates(st)
+        self.assertEqual(len(srt), len(st))
+        # Rows are lexicographically non-decreasing (batch col most significant).
+        C = srt.C.cpu().numpy()
+        keys = tuple(C[:, i] for i in reversed(range(C.shape[1])))
+        order = np.lexsort(keys)
+        self.assertTrue(np.array_equal(order, np.arange(len(C))))
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_deterministic_conv_stride2(self):
+        in_channels, out_channels, D = 3, 5, 3
+        coords, feats = _make_point_set(4000, D, in_channels, seed=0)
+
+        torch.manual_seed(1234)
+        conv = MinkowskiConvolution(
+            in_channels, out_channels, kernel_size=3, stride=2, dimension=D
+        ).cuda()
+
+        ME.set_deterministic(True)
+        reference = None
+        for trial in range(5):
+            g = torch.Generator().manual_seed(100 + trial)
+            perm = torch.randperm(coords.shape[0], generator=g)
+            inp = SparseTensor(feats[perm].clone(), coords[perm].clone(), device="cuda")
+            out = conv(inp)
+            F_canon, C_canon = _canonical_feats(out)
+            if reference is None:
+                reference = (F_canon, C_canon)
+            else:
+                self.assertTrue(
+                    np.array_equal(C_canon, reference[1]),
+                    "coordinate sets differ across permutations",
+                )
+                self.assertTrue(
+                    torch.equal(F_canon, reference[0]),
+                    f"deterministic stride-2 conv output differs on trial {trial}",
+                )
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_deterministic_conv_transpose(self):
+        in_channels, out_channels, D = 3, 5, 3
+        coords, feats = _make_point_set(3000, D, in_channels, seed=7)
+        # Strided (tensor_stride=2) input so the transpose upsamples to stride 1.
+        coords[:, 1:] = coords[:, 1:] * 2
+
+        torch.manual_seed(4321)
+        conv_tr = MinkowskiConvolutionTranspose(
+            in_channels, out_channels, kernel_size=3, stride=2, dimension=D
+        ).cuda()
+
+        ME.set_deterministic(True)
+        reference = None
+        for trial in range(5):
+            g = torch.Generator().manual_seed(200 + trial)
+            perm = torch.randperm(coords.shape[0], generator=g)
+            inp = SparseTensor(
+                feats[perm].clone(),
+                coords[perm].clone(),
+                tensor_stride=2,
+                device="cuda",
+            )
+            out = conv_tr(inp)
+            F_canon, C_canon = _canonical_feats(out)
+            if reference is None:
+                reference = (F_canon, C_canon)
+            else:
+                self.assertTrue(
+                    np.array_equal(C_canon, reference[1]),
+                    "coordinate sets differ across permutations",
+                )
+                self.assertTrue(
+                    torch.equal(F_canon, reference[0]),
+                    f"deterministic transpose conv output differs on trial {trial}",
+                )
+
+    @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+    def test_flag_off_path_runs(self):
+        # With the flag off we do NOT assert cross-permutation equality (that
+        # would be flaky by design); we only assert the default path still runs
+        # and backprops.
+        in_channels, out_channels, D = 3, 4, 3
+        coords, feats = _make_point_set(2000, D, in_channels, seed=3)
+
+        torch.manual_seed(11)
+        conv = MinkowskiConvolution(
+            in_channels, out_channels, kernel_size=3, stride=2, dimension=D
+        ).cuda()
+
+        ME.set_deterministic(False)
+        self.assertFalse(ME.is_deterministic())
+        inp = SparseTensor(feats.clone(), coords.clone(), device="cuda")
+        inp.F.requires_grad_()
+        out = conv(inp)
+        out.F.sum().backward()
+        self.assertIsNotNone(inp.F.grad)
+        self.assertTrue(torch.isfinite(inp.F.grad).all())
+
+
+if __name__ == "__main__":
+    unittest.main()