From 0218856cf66ac7bc3ccc03c37bed9aa90ad764b5 Mon Sep 17 00:00:00 2001 From: alpsaur Date: Fri, 17 Jul 2026 12:26:03 +0800 Subject: [PATCH] fix: opt-in deterministic conv, pruning empty-mask guard, TensorField backprop test Three correctness items for MinkowskiEngine convolution/pruning/tensor-field. 1. Opt-in deterministic convolution (upstream #554, #504) Conv output depends on input coordinate ordering (atomicAdd accumulation + kernel-map dispatch order), so repeated runs on the same unsorted input differ in low-order bits, worst at stride>1 / transpose. Add an opt-in switch that canonicalizes coordinate order before the backend call: - MinkowskiEngine/utils/determinism.py: sorted_coordinates(sparse_tensor) returns a new SparseTensor with rows in lexicographic coordinate order (stable per-column radix argsort, overflow-free, handles arbitrary D and negative coords), plus a process-wide set_deterministic/is_deterministic flag. - MinkowskiConvolution / MinkowskiConvolutionTranspose forward sort their input through the helper when the flag is on (skipped for the use_mm kernel_volume==1 path). Default behavior unchanged; documented perf cost of one sort per conv. - Exported as ME.set_deterministic / ME.is_deterministic / ME.sorted_coordinates. 2. Pruning empty-mask guard (upstream #579) Reproduced on this fork: an all-False mask through MinkowskiPruning on CUDA already returns a valid empty SparseTensor; downstream conv and backward run cleanly with finite (empty) grads (src/pruning_gpu.cu handles the size-0 case). No code change needed; add a regression test that locks in the behavior. 3. TensorField backprop regression test (upstream #395) The merge_sort illegal-address crash in inverse_mapping backward was fixed upstream by 02fc608 (in our history). Add a field -> sparse -> conv -> slice -> field backward test over several iterations with varying N (~50k, D=3, batch 2) asserting finite gradients. Tests are CUDA-gated (skipUnless) and skip cleanly on CPU CI. --- MinkowskiEngine/MinkowskiConvolution.py | 15 ++ MinkowskiEngine/__init__.py | 6 + MinkowskiEngine/utils/__init__.py | 3 +- MinkowskiEngine/utils/determinism.py | 154 +++++++++++++++++++++ tests/python/determinism.py | 173 ++++++++++++++++++++++++ tests/python/pruning_empty.py | 91 +++++++++++++ tests/python/tensorfield_backward.py | 108 +++++++++++++++ 7 files changed, 549 insertions(+), 1 deletion(-) create mode 100644 MinkowskiEngine/utils/determinism.py create mode 100644 tests/python/determinism.py create mode 100644 tests/python/pruning_empty.py create mode 100644 tests/python/tensorfield_backward.py diff --git a/MinkowskiEngine/MinkowskiConvolution.py b/MinkowskiEngine/MinkowskiConvolution.py index 1f7008ea..65212799 100644 --- a/MinkowskiEngine/MinkowskiConvolution.py +++ b/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 diff --git a/MinkowskiEngine/__init__.py b/MinkowskiEngine/__init__.py index 446fab15..8f8ba7e1 100644 --- a/MinkowskiEngine/__init__.py +++ b/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 ( diff --git a/MinkowskiEngine/utils/__init__.py b/MinkowskiEngine/utils/__init__.py index b5b1ea8e..8d11481f 100644 --- a/MinkowskiEngine/utils/__init__.py +++ b/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 \ No newline at end of file +from .summary import summary +from .determinism import set_deterministic, is_deterministic, sorted_coordinates \ No newline at end of file diff --git a/MinkowskiEngine/utils/determinism.py b/MinkowskiEngine/utils/determinism.py new file mode 100644 index 00000000..525c4dbc --- /dev/null +++ b/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, + ) diff --git a/tests/python/determinism.py b/tests/python/determinism.py new file mode 100644 index 00000000..d5b57c72 --- /dev/null +++ b/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() diff --git a/tests/python/pruning_empty.py b/tests/python/pruning_empty.py new file mode 100644 index 00000000..071d1a61 --- /dev/null +++ b/tests/python/pruning_empty.py @@ -0,0 +1,91 @@ +# 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 test for pruning with an all-False mask on CUDA. +# Upstream issue: https://github.com/NVIDIA/MinkowskiEngine/issues/579 +# +# On this fork the CUDA pruning kernel already emits a valid empty SparseTensor +# for an empty mask (src/pruning_gpu.cu handles the size-0 case), so this test +# documents/locks in that behavior: an all-False mask followed by a downstream +# convolution and backward must not crash, must produce zero-size outputs, and +# must yield finite (empty) gradients. +import unittest + +import torch + +from MinkowskiEngine import ( + SparseTensor, + MinkowskiConvolution, + MinkowskiPruning, +) + + +class TestPruningEmptyMask(unittest.TestCase): + @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") + def test_all_false_mask_downstream_conv_backward(self): + in_channels, out_channels, D = 4, 8, 3 + coords = torch.IntTensor( + [ + [0, 0, 0, 0], + [0, 1, 0, 0], + [0, 0, 1, 0], + [0, 0, 0, 1], + [0, 1, 1, 1], + ] + ) + feats = torch.rand(len(coords), in_channels) + feats.requires_grad_() + inp = SparseTensor(feats, coords, device="cuda") + + mask = torch.zeros(len(inp), dtype=torch.bool, device="cuda") # all False + pruning = MinkowskiPruning() + pruned = pruning(inp, mask) + + # Empty tensor: zero rows, feature width preserved. + self.assertEqual(len(pruned), 0) + self.assertEqual(pruned.F.shape[0], 0) + self.assertEqual(pruned.F.shape[1], in_channels) + # Coordinate access on the empty tensor must not crash. + self.assertEqual(pruned.C.shape[0], 0) + torch.cuda.synchronize() + + # Downstream convolution on the empty tensor. + conv = MinkowskiConvolution( + in_channels, out_channels, kernel_size=3, stride=1, dimension=D + ).cuda() + out = conv(pruned) + self.assertEqual(len(out), 0) + self.assertEqual(out.F.shape[0], 0) + self.assertEqual(out.F.shape[1], out_channels) + torch.cuda.synchronize() + + # Backward through the empty pipeline: grad must exist, be shape-correct + # and finite (vacuously, since it is empty). + out.F.sum().backward() + self.assertIsNotNone(feats.grad) + self.assertEqual(tuple(feats.grad.shape), (len(coords), in_channels)) + self.assertTrue(torch.isfinite(feats.grad).all()) + # No feature survived the all-False prune, so nothing propagated back. + self.assertEqual(float(feats.grad.abs().sum()), 0.0) + torch.cuda.synchronize() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/python/tensorfield_backward.py b/tests/python/tensorfield_backward.py new file mode 100644 index 00000000..15db274e --- /dev/null +++ b/tests/python/tensorfield_backward.py @@ -0,0 +1,108 @@ +# 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 test for TensorField -> sparse -> conv -> slice backprop. +# Upstream issue: https://github.com/NVIDIA/MinkowskiEngine/issues/395 +# +# The reported crash was a merge_sort cudaErrorIllegalAddress in the +# inverse_mapping backward pass, fixed upstream by commit 02fc608 (present in +# this fork's history). This test proves the fix holds: it builds a TensorField +# from tens of thousands of random points, runs field -> sparse -> conv -> +# slice back to a field, computes a loss on the FIELD-level output, and +# backpropagates. It runs several iterations with varying N to catch the +# intermittent illegal-address failure the original bug exhibited. +import unittest + +import torch +import torch.nn as nn + +import MinkowskiEngine as ME +from MinkowskiEngine import ( + TensorField, + MinkowskiConvolution, + MinkowskiConvolutionTranspose, + MinkowskiLinear, + MinkowskiReLU, + MinkowskiToSparseTensor, +) + + +def _make_field(N, D, device, seed): + torch.manual_seed(seed) + n0 = N // 2 + n1 = N - n0 + c0 = torch.rand(n0, D) * 100.0 + c1 = torch.rand(n1, D) * 100.0 + coords = torch.cat( + [ + torch.cat([torch.zeros(n0, 1), c0], dim=1), + torch.cat([torch.ones(n1, 1), c1], dim=1), + ], + dim=0, + ).float() + feats = torch.rand(N, 3) + feats.requires_grad_() + return coords.to(device), feats + + +class TestTensorFieldBackward(unittest.TestCase): + @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") + def test_field_sparse_conv_slice_backward(self): + device = "cuda" + D = 3 + + # Vary N across iterations to catch the intermittent illegal-address + # crash the original #395 bug exhibited. + for it, N in enumerate([50000, 40000, 60000, 30000, 55000]): + coords, feats = _make_field(N, D, device, seed=it) + tfield = TensorField(coordinates=coords, features=feats.to(device)) + + net = nn.Sequential( + MinkowskiLinear(3, 16), + MinkowskiReLU(), + MinkowskiToSparseTensor(), + MinkowskiConvolution(16, 32, kernel_size=3, stride=2, dimension=D), + MinkowskiConvolutionTranspose( + 32, 16, kernel_size=3, stride=2, dimension=D + ), + ).to(device) + + soutput = net(tfield) + # Slice back to the original (field-level) resolution and take the + # loss there, so gradients flow through the field inverse mapping. + ofield = soutput.slice(tfield) + self.assertEqual(len(ofield), len(tfield)) + + loss = ofield.F.pow(2).sum() + loss.backward() + torch.cuda.synchronize() + + self.assertIsNotNone( + feats.grad, f"iter {it} (N={N}): field feature grad is None" + ) + self.assertEqual(tuple(feats.grad.shape), (N, 3)) + self.assertTrue( + torch.isfinite(feats.grad).all(), + f"iter {it} (N={N}): non-finite gradient", + ) + + +if __name__ == "__main__": + unittest.main()