NVIDIA / MinkowskiEngine

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

Creating the Coordinate Map directly #562

Open HaFred opened 1 year ago

HaFred commented 1 year ago

Is your feature request related to a problem? Please describe. I have an input sparse tensor to a sparse NN, looking like the below:

SparseTensor(
  coordinates=tensor([[  0, 308, 251, 100],
        [  0, 307, 252,  97],
        [  0, 307, 252,  99],
        ...,
        [  0,  83, 249,  90],
        [  0,  83, 249,  89],
        [  0,  83, 246,  97]], device='cuda:0', dtype=torch.int32)
  features=tensor([[ 0.0118,  0.0902,  0.4275],
        [ 0.1137,  0.0588,  0.4431],
        [ 0.1294,  0.2000,  0.4745],
        ...,
        [-0.4275, -0.3961,  0.6941],
        [-0.1137,  0.0353,  0.6549],
        [-0.2157, -0.1294,  0.2784]], device='cuda:0')
  coordinate_map_key=coordinate map key:[1, 1, 1]
  coordinate_manager=CoordinateMapManagerGPU_c10(
    [1, 1, 1, ]:  CoordinateMapGPU:68270x4
    [2, 2, 2, ]:  CoordinateMapGPU:33366x4
    [4, 4, 4, ]:  CoordinateMapGPU:9637x4
    [8, 8, 8, ]:  CoordinateMapGPU:2447x4
    [4, 4, 4, ]->[8, 8, 8, ]:   gpu_kernel_map: number of unique maps:8, kernel map size:9637
    [8, 8, 8, ]->[8, 8, 8, ]:   gpu_kernel_map: number of unique maps:27, kernel map size:38253
    [4, 4, 4, ]->[4, 4, 4, ]:   gpu_kernel_map: number of unique maps:27, kernel map size:141005
    [2, 2, 2, ]->[4, 4, 4, ]:   gpu_kernel_map: number of unique maps:8, kernel map size:33366
    [2, 2, 2, ]->[2, 2, 2, ]:   gpu_kernel_map: number of unique maps:27, kernel map size:406254
    [1, 1, 1, ]->[2, 2, 2, ]:   gpu_kernel_map: number of unique maps:8, kernel map size:68270
    [1, 1, 1, ]->[1, 1, 1, ]:   gpu_kernel_map: number of unique maps:125, kernel map size:1380550
    algorithm=MinkowskiAlgorithm.DEFAULT
  )
  spatial dimension=3)

However, I would like to create such a sparse tensor without passing to a new NN directly, by assigning its coordinate map. It seems like there are no current solutions for this.

Describe alternatives you've considered I tried below:

sparse_tensor = ME.SparseTensor(feat, coords, tensor_stride=4)

But it cannot return the multiple-strides sparse tensor like the snippets above.

Thanks in advance.