NVIDIA / MinkowskiEngine

Minkowski Engine is an auto-diff neural network library for high-dimensional sparse tensors
https://nvidia.github.io/MinkowskiEngine
Other
2.47k stars 367 forks source link

add cross-shaped kernel support #437

Closed Karbo123 closed 2 years ago

Karbo123 commented 2 years ago

I previously create an issue (#436) about this feature request, but I finally figure out a simple way to make it support cross-shaped kernel (i.e. HYPER_CROSS). Since the previous version only implements cubic kernel in the coordinate_at functions defined at src/kernel_region.hpp, after writing some codes for the HYPER_CROSS, it supports the cross-shaped kernel now.

test codes for verification:

import torch
from MinkowskiEngine import CoordinateManager
from MinkowskiEngineBackend._C import RegionType, CoordinateMapType

lin = torch.arange(10)
coords_in = torch.cat([torch.zeros([1000, 1], dtype=torch.long),
                       torch.stack(torch.meshgrid(lin, lin, lin), dim=-1).reshape(-1, 3)], dim=1).int().cuda()
coords_out = torch.tensor([[0, 5, 5, 5]]).int().cuda()

cm = CoordinateManager(D=3, coordinate_map_type=CoordinateMapType.CUDA)
ck_in, _ = cm.insert_and_map(coordinates=coords_in, tensor_stride=1, string_id="input")
ck_out, _ = cm.insert_and_map(coordinates=coords_out, tensor_stride=1, string_id="output")

kernel_kwargs = dict(kernel_size=[5, 3, 7], region_type=RegionType.HYPER_CROSS)
ind = torch.cat(list(cm.kernel_map(ck_in, ck_out, **kernel_kwargs).values()), dim=1)[0]
diff = coords_in[ind.long()] - coords_out
print(diff)
print(f"center = {diff.float().mean()} (should be zero)")

the previous commit will output a wrong result like this:

tensor([[ 0,  0,  1, -3],
        [ 0, -1,  1, -3],
        [ 0,  2, -1, -3],
        [ 0,  1, -1, -3],
        [ 0,  0, -1, -3],
        [ 0, -2, -1, -3],
        [ 0, -1, -1, -3],
        [ 0, -2,  0, -3],
        [ 0, -1,  0, -3],
        [ 0,  0,  0, -3],
        [ 0,  1,  0, -3],
        [ 0,  2,  0, -3],
        [ 0, -2,  1, -3]], device='cuda:0', dtype=torch.int32)
center = -0.8461538553237915 (should be zero)

after this commit, the problem is solved, and it prints:

tensor([[ 0,  0,  0, -1],
        [ 0,  0,  0, -2],
        [ 0, -1,  0,  0],
        [ 0, -2,  0,  0],
        [ 0,  2,  0,  0],
        [ 0,  0,  0,  0],
        [ 0,  1,  0,  0],
        [ 0,  0,  1,  0],
        [ 0,  0, -1,  0],
        [ 0,  0,  0,  1],
        [ 0,  0,  0,  2],
        [ 0,  0,  0,  3],
        [ 0,  0,  0, -3]], device='cuda:0', dtype=torch.int32)
center = 0.0 (should be zero)

hopefully this may help.

chrischoy commented 2 years ago

Thanks for the contribution!