Fix MinkowskiInterpolation zero/under-weighted results at sparse boundaries (renormalize over existing neighbors)

MinkowskiEngine/MinkowskiInterpolation.py

@@ -36,6 +36,61 @@
 )
 
 
+def _renormalize_partial_neighborhoods(
+    out_feat: torch.Tensor,
+    out_map: torch.Tensor,
+    weights: torch.Tensor,
+    dimension: int,
+):
+    r"""Renormalize interpolation results over the lattice corners that exist.
+
+    The backend computes, for every query point, multilinear weights over the
+    :math:`2^D` surrounding lattice corners and silently drops corners that
+    are not occupied in the sparse tensor. Without renormalization the
+    remaining weights no longer sum to one, so any query with a partially
+    missing neighborhood is systematically under-weighted -- down to exactly
+    zero when no corner exists (NVIDIA/MinkowskiEngine issues #477, #363,
+    #383). This divides both the output features and the weights of such
+    rows by the sum of the surviving weights.
+
+    Rows with a complete :math:`2^D` neighborhood (whose weights sum to one
+    by construction) and rows with no surviving corner (which stay zero) are
+    returned bit-identical to the backend output.
+    """
+    if out_feat.shape[0] == 0 or weights.numel() == 0:
+        return out_feat, weights
+
+    num_out = out_feat.shape[0]
+    full_neighborhood = 2 ** dimension
+    out_map_long = out_map.long()
+    # accumulate in fp32 (or fp64) even for 16-bit weights
+    w_acc = (
+        weights.float()
+        if weights.dtype in (torch.float16, torch.bfloat16)
+        else weights
+    )
+    weight_sum = torch.zeros(num_out, dtype=w_acc.dtype, device=w_acc.device)
+    weight_sum.scatter_add_(0, out_map_long, w_acc)
+    num_found = torch.zeros(num_out, dtype=torch.long, device=w_acc.device)
+    num_found.scatter_add_(0, out_map_long, torch.ones_like(out_map_long))
+
+    partial = (num_found < full_neighborhood) & (weight_sum > 0)
+    if not bool(partial.any()):
+        return out_feat, weights
+
+    scale = torch.where(
+        partial, weight_sum.reciprocal(), torch.ones_like(weight_sum)
+    )
+    # multiplication by exactly 1.0 leaves full-neighborhood /
+    # exact-on-voxel rows bit-identical
+    out_feat = out_feat * scale.to(out_feat.dtype).unsqueeze(1)
+    # keep the weights consistent with the produced features; the backward
+    # pass and the `return_weights=True` output both use the renormalized
+    # weights
+    weights = weights * scale.to(weights.dtype)[out_map_long]
+    return out_feat, weights
+
+
 class MinkowskiInterpolationFunction(Function):
     @staticmethod
     def forward(
@@ -68,6 +123,9 @@ def forward(
                 in_coordinate_map_key,
                 coordinate_manager._manager,
             )
+            out_feat, weights = _renormalize_partial_neighborhoods(
+                out_feat, out_map, weights, tfield.size(1) - 1
+            )
         ctx.save_for_backward(in_map, out_map, weights)
         ctx.inputs = (
             in_coordinate_map_key,
@@ -100,7 +158,14 @@ def backward(
 
 
 class MinkowskiInterpolation(MinkowskiModuleBase):
-    r"""Sample linearly interpolated features at the provided points."""
+    r"""Sample linearly interpolated features at the provided points.
+
+    Interpolation weights are renormalized over the lattice corners that are
+    actually occupied, so queries at the boundary of the sparse support
+    return a weighted average of the existing neighbors instead of decaying
+    to zero. Queries whose :math:`2^D` neighborhood contains no occupied
+    corner return zeros.
+    """
 
     def __init__(self, return_kernel_map=False, return_weights=False):
         r"""Sample linearly interpolated features at the specified coordinates.

tests/conftest.py

@@ -83,6 +83,11 @@ def _gradcheck_compat(*args, **kwargs):
     "tests/python/coordinate_manager.py::CoordinateManagerTestCase::test_stride_cuda",
     "tests/python/interpolation.py::TestInterpolation::test_gpu",
     "tests/python/interpolation.py::TestInterpolation::test_zero",
+    "tests/python/interpolation_renorm.py::InterpolationRenormTestCase::test_forward_gpu",
+    "tests/python/interpolation_renorm.py::InterpolationRenormTestCase::test_return_weights_gpu",
+    "tests/python/interpolation_renorm.py::InterpolationRenormTestCase::test_gradcheck_gpu",
+    "tests/python/interpolation_renorm.py::InterpolationRenormTestCase::test_half_precision_autocast_gpu",
+    "tests/python/interpolation_renorm.py::InterpolationRenormTestCase::test_features_at_coordinates_gpu",
     "tests/python/norm.py::TestNormalization::test_inst_norm_gpu",
     "tests/python/pruning.py::TestPruning::test_device",
     "tests/python/quantization.py::TestQuantization::test_device",

tests/python/interpolation_renorm.py

@@ -0,0 +1,244 @@
+# Copyright (c) 2020 NVIDIA CORPORATION.
+# Copyright (c) 2018-2020 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 interpolation-weight renormalization at sparse
+boundaries (NVIDIA/MinkowskiEngine issues #477, #363, #383).
+
+Query points whose 2^D lattice neighborhood is only partially occupied used
+to be systematically under-weighted (down to zero) because the weights of
+missing corners were dropped without renormalizing over the surviving ones.
+"""
+import itertools
+import math
+import unittest
+
+import torch
+
+from MinkowskiEngine import (
+    SparseTensor,
+    MinkowskiInterpolation,
+    MinkowskiInterpolationFunction,
+)
+from MinkowskiEngine.MinkowskiCommon import get_minkowski_function
+
+
+def _cube_plus_isolated_voxel(dtype, device):
+    """A fully occupied 2x2x2 cube at (0..1)^3 plus an isolated voxel at
+    (5,5,5), all in batch 0. Features are 1..9."""
+    coords = [[0, i, j, k] for i, j, k in itertools.product(range(2), repeat=3)]
+    coords.append([0, 5, 5, 5])
+    coords = torch.tensor(coords, dtype=torch.int32, device=device)
+    feats = torch.arange(1, len(coords) + 1, dtype=dtype, device=device).unsqueeze(1)
+    return coords, feats
+
+
+def _reference_interpolation(coords, feats, tfield, renormalize=True):
+    """Manual multilinear interpolation over occupied corners only,
+    renormalized by the sum of surviving weights (in fp64)."""
+    fmap = {tuple(c.tolist()): f for c, f in zip(coords.cpu(), feats.cpu().double())}
+    D = tfield.size(1) - 1
+    out = torch.zeros(tfield.size(0), feats.size(1), dtype=torch.float64)
+    for i, q in enumerate(tfield.cpu().double()):
+        b = int(round(q[0].item()))
+        acc = torch.zeros(feats.size(1), dtype=torch.float64)
+        wsum = 0.0
+        lb = [math.floor(q[1 + d].item()) for d in range(D)]
+        for corner in itertools.product(*[(lb[d], lb[d] + 1) for d in range(D)]):
+            w = 1.0
+            for d in range(D):
+                w *= 1 - abs(q[1 + d].item() - corner[d])
+            key = (b,) + corner
+            if key in fmap:
+                acc += w * fmap[key]
+                wsum += w
+        if renormalize and wsum > 0:
+            acc /= wsum
+        out[i] = acc
+    return out
+
+
+def _raw_backend_forward(x, tfield):
+    """Un-renormalized backend output (the pre-fix behavior)."""
+    fw_fn = get_minkowski_function("InterpolationForward", x.F)
+    out_feat, in_map, out_map, weights = fw_fn(
+        x.F, tfield, x.coordinate_map_key, x.coordinate_manager._manager
+    )
+    return out_feat
+
+
+class InterpolationRenormTestCase(unittest.TestCase):
+    def _run_forward_case(self, device, dtype=torch.float32):
+        coords, feats = _cube_plus_isolated_voxel(dtype, device)
+        x = SparseTensor(feats, coordinates=coords)
+        interp = MinkowskiInterpolation()
+
+        # (a) exactly on occupied voxels, (b) interior with a full
+        # neighborhood, (c) partial neighborhoods (1/8 and 4/8 corners).
+        tfield = torch.tensor(
+            [
+                [0, 0, 0, 0],
+                [0, 1, 1, 1],
+                [0, 5, 5, 5],
+                [0, 0.25, 0.5, 0.75],
+                [0, 5.25, 5.25, 5.25],
+                [0, 0.5, 0.5, 1.75],
+                # no occupied corner at all -> must stay exactly zero
+                [0, 20.5, 20.5, 20.5],
+            ],
+            dtype=dtype,
+            device=device,
+        )
+        out = interp(x, tfield)
+        raw = _raw_backend_forward(x, tfield)
+
+        # exact-on-voxel and full-neighborhood rows are bit-identical to the
+        # un-renormalized backend result
+        for row in range(4):
+            self.assertTrue(
+                torch.equal(out[row], raw[row]),
+                f"row {row} must be bit-identical to the backend output",
+            )
+        # exact-on-voxel rows equal the stored features exactly
+        self.assertTrue(torch.equal(out[0], feats[0]))
+        self.assertTrue(torch.equal(out[1], feats[7]))
+        self.assertTrue(torch.equal(out[2], feats[8]))
+
+        # all rows match the renormalized manual reference
+        ref = _reference_interpolation(coords, feats, tfield).to(out)
+        self.assertTrue(
+            torch.allclose(out, ref, rtol=1e-5, atol=1e-6),
+            f"expected {ref.flatten().tolist()}, got {out.flatten().tolist()}",
+        )
+
+        # the bug: boundary rows used to be badly under-weighted
+        # (0.421875 * 9 = 3.796875 and 0.25 * 5 = 1.25); they must now be the
+        # renormalized average of the existing neighbors, NOT near-zero
+        self.assertAlmostEqual(out[4].item(), 9.0, places=4)
+        self.assertAlmostEqual(out[5].item(), 5.0, places=4)
+        # empty neighborhood stays exactly zero
+        self.assertTrue(torch.equal(out[6], torch.zeros_like(out[6])))
+
+    def test_forward_cpu(self):
+        self._run_forward_case("cpu")
+
+    def test_forward_gpu(self):
+        self._run_forward_case("cuda")
+
+    def _run_return_weights_case(self, device):
+        coords, feats = _cube_plus_isolated_voxel(torch.float32, device)
+        x = SparseTensor(feats, coordinates=coords)
+        interp = MinkowskiInterpolation(return_kernel_map=True, return_weights=True)
+        tfield = torch.tensor(
+            [[0, 5.25, 5.25, 5.25], [0, 0.25, 0.5, 0.75]],
+            dtype=torch.float32,
+            device=device,
+        )
+        out, (in_map, out_map), weights = interp(x, tfield)
+        # returned weights are the ones actually applied: they sum to 1 for
+        # every query that found at least one occupied corner
+        wsum = torch.zeros(tfield.size(0), device=device)
+        wsum.scatter_add_(0, out_map.long(), weights)
+        self.assertTrue(torch.allclose(wsum, torch.ones_like(wsum), atol=1e-6))
+
+    def test_return_weights_cpu(self):
+        self._run_return_weights_case("cpu")
+
+    def test_return_weights_gpu(self):
+        self._run_return_weights_case("cuda")
+
+    def _run_gradcheck_case(self, device):
+        # fp64 finite-difference gradient check on a tiny case with a
+        # partial neighborhood (exercises the renormalized backward)
+        coords = torch.tensor(
+            [[0, 0, 0, 0], [0, 1, 0, 0], [0, 5, 5, 5]],
+            dtype=torch.int32,
+            device=device,
+        )
+        feats = torch.randn(3, 2, dtype=torch.float64, device=device)
+        feats.requires_grad_()
+        x = SparseTensor(feats, coordinates=coords)
+        tfield = torch.tensor(
+            [
+                [0, 0.5, 0, 0],  # full 1D-edge neighborhood along x
+                [0, 0.25, 0.75, 0.5],  # partial: 2/8 corners occupied
+                [0, 5.5, 5.25, 5.75],  # partial: 1/8 corners occupied
+            ],
+            dtype=torch.float64,
+            device=device,
+        )
+
+        def interp_only_features(f):
+            return MinkowskiInterpolationFunction.apply(
+                f, tfield, x.coordinate_map_key, x.coordinate_manager
+            )[0]
+
+        self.assertTrue(
+            torch.autograd.gradcheck(interp_only_features, (feats,), eps=1e-6, atol=1e-4)
+        )
+
+    def test_gradcheck_cpu(self):
+        self._run_gradcheck_case("cpu")
+
+    def test_gradcheck_gpu(self):
+        self._run_gradcheck_case("cuda")
+
+    def test_half_precision_autocast_gpu(self):
+        # 16-bit features are promoted to the fp32 coordinate-field dtype
+        # under autocast; boundary queries must match the fp32 renormalized
+        # reference within half-precision tolerance
+        for dtype, tol in ((torch.float16, 1e-3), (torch.bfloat16, 1e-2)):
+            coords, feats = _cube_plus_isolated_voxel(torch.float32, "cuda")
+            x16 = SparseTensor(feats.to(dtype), coordinates=coords)
+            x32 = SparseTensor(feats, coordinates=coords)
+            tfield = torch.tensor(
+                [[0, 5.25, 5.25, 5.25], [0, 0.5, 0.5, 1.75], [0, 0.25, 0.5, 0.75]],
+                dtype=torch.float32,
+                device="cuda",
+            )
+            interp = MinkowskiInterpolation()
+            with torch.autocast("cuda", dtype=dtype):
+                out16 = interp(x16, tfield)
+            out32 = interp(x32, tfield)
+            self.assertEqual(out16.dtype, torch.float32)
+            self.assertTrue(torch.isfinite(out16).all())
+            self.assertTrue(
+                torch.allclose(out16, out32, rtol=tol, atol=tol),
+                f"{dtype}: expected {out32.flatten().tolist()}, "
+                f"got {out16.flatten().tolist()}",
+            )
+
+    def _run_features_at_coordinates_case(self, device):
+        # SparseTensor.features_at_coordinates goes through the same
+        # function and must return renormalized boundary values
+        coords, feats = _cube_plus_isolated_voxel(torch.float32, device)
+        x = SparseTensor(feats, coordinates=coords)
+        q = torch.tensor([[0, 5.25, 5.25, 5.25]], dtype=torch.float32, device=device)
+        out = x.features_at_coordinates(q)
+        self.assertAlmostEqual(out.item(), 9.0, places=4)
+
+    def test_features_at_coordinates_cpu(self):
+        self._run_features_at_coordinates_case("cpu")
+
+    def test_features_at_coordinates_gpu(self):
+        self._run_features_at_coordinates_case("cuda")
+
+
+if __name__ == "__main__":
+    unittest.main()