mit-han-lab / torchsparse

[MICRO'23, MLSys'22] TorchSparse: Efficient Training and Inference Framework for Sparse Convolution on GPUs.
https://torchsparse.mit.edu
MIT License
1.22k stars 143 forks source link

Try to understand spconv behavior #327

Closed nullplay closed 2 months ago

nullplay commented 2 months ago

Hi, I'm trying to learn the behavior of torchsparse's submanifold convolution.

If I understand correctly, a submanifold convolution only performs convolution on the nonzero coordinates. As an example,

# input 
# 1 0 1 0
# 0 1 0 0
# 1 0 0 0
# 0 0 0 1

vvv submanifold conv with 3x3 (weights are initialized with ones) vvv

# expected output after submanifold 3x3 conv
# 2 0 2 0
# 0 4 0 0
# 2 0 0 0
# 0 0 0 1

When I try to confirm my understanding with torchsparse library, it outputs different result than I expected.

import torch
import torchsparse.nn as spnn
from torchsparse import SparseTensor

device = "cuda:0"

# input coords
# 1 0 1 0
# 0 1 0 0
# 1 0 0 0
# 0 0 0 1
coords = torch.tensor(
    [[0,0,0],
     [0,2,0],
     [1,1,0],
     [2,0,0],
     [3,3,0]], dtype=torch.int)
feats = torch.ones((6,1), dtype=torch.float)
input = SparseTensor(coords=coords, feats=feats).to(device)

# 3x3 conv layer initiliazed with ones
conv = spnn.Conv3d(1,1,kernel_size=(3,3,1),stride=1).to(device)
conv.kernel.data.fill_(1.0)

# expected output after submanifold 3x3 conv
# 2 0 2 0
# 0 4 0 0
# 2 0 0 0
# 0 0 0 1
output = conv(input)

# Actual Output
# tensor([[0, 0, 0],
#        [0, 2, 0],
#        [1, 1, 0],
#        [2, 0, 0],
#        [3, 3, 0]], device='cuda:0', dtype=torch.int32)
# tensor([[1.],
#        [1.],
#        [1.],
#        [1.],
#        [1.]], device='cuda:0',
#       grad_fn=<ImplicitGEMMConvolutionFuntionBackward>)
print(output.coords, "\n",  output.feats)

Can you help me with figuring out what I'm missing?

zhijian-liu commented 2 months ago

coords should have a shape of Nx4, where 4 corresponds to (B, X, Y, Z). To achieve this, you need to append an additional column of zeros to the coordinate tensor.

nullplay commented 2 months ago

Awesome. Thanks for the fast response.