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.19k stars 139 forks source link

Usage of generalized convolution #215

Open 96lives opened 1 year ago

96lives commented 1 year ago

Hello, I'm trying to use the generalized convolution of sparse tensors and want to ask you guys how to use it, cause the issues that I've looked seems to support generalized convolution as in MinkowskiEngine.

In MinkowskiEngine, the generalized convolution is supported by taking coordinates as inputs during the forward() on the convolution operation as in here, which allows us to compute the features on the coordinates that is not directly in the input coordinates. Looking at the torch-sparse issues in github, generlized convolution seems to be implemented, but I can't find the right documentation for it :(

Could you help me with this issue? I want to compute the features that are within the neighborhood of the input coordinates.

Thanks a milliions for the wonderful repo!

kentang-mit commented 1 year ago

Hi @96lives,

To enable generative transposed sparseconv, you may simply pass in the generative=True argument to the convolution module:

def make_up_block(in_channels, out_channels, generative=False):
    conv = partial(spnn.Conv3d, transposed=True, generative=generative)
    return nn.Sequential(
        conv(in_channels, out_channels, kernel_size=3, stride=2),
        spnn.BatchNorm(out_channels),
        spnn.ReLU(inplace=True),
    )

This is a sample code snippet in which you can get a generative transposed sparseconv block by passing in generative=True.

I will release another example codebase for indoor object detection that utilizes this operator soon.

Best, Haotian

96lives commented 1 year ago

@kentang-mit Thanks for the quick reply! However, to the best of my knowledge, the spnn.Conv3d module does not take in the generative argument currently. Is there a quick way to fix this? or should I wait for the update?

Best, Dongsu

kentang-mit commented 1 year ago

Hi @96lives,

The source code of torchsparse 2.1.0 is still under internal quality check, so what you see is the 2.0.0 source code which does not have that argument. But for the library you installed from the wheels, the generative argument is there. Feel free to directly add it.

Best, Haotian

96lives commented 1 year ago

@kentang-mit Thanks for the reply!! Is works. Still I'm quite not sure as how to obtain the features of the generative convolution. For example, I wish to obtain the neighborhood features from the following convolution afther the UNet (code snippet is a simple modification of the examples directory of the repo):

import numpy as np
import torch
from torch import nn

from torchsparse import SparseTensor
from torchsparse.backbones import SparseResNet21D, SparseResUNet42
from torchsparse.utils.quantize import sparse_quantize
from torchsparse import nn as spnn
import torchsparse.nn.functional as F
F.set_kmap_mode("hashmap")

@torch.no_grad()
def main() -> None:
    device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

    for backbone in [SparseResUNet42]:
        print(f'{backbone.__name__}:')
        model: nn.Module = backbone(in_channels=4, width_multiplier=1.0)
        model = model.to(device).eval()

        conv3d = spnn.Conv3d(
            in_channels=model.decoder_channels[-1],
            out_channels=1,
            transposed=True,
            generative=True
        )
        conv3d = conv3d.to(device)

        # generate data
        input_size, voxel_size = 100, 0.2
        inputs = np.random.uniform(-100, 100, size=(input_size, 4))
        pcs, feats = inputs[:, :3], inputs
        pcs -= np.min(pcs, axis=0, keepdims=True)
        pcs, indices = sparse_quantize(pcs, voxel_size, return_index=True)
        coords = np.zeros((pcs.shape[0], 4))
        coords[:, -3:] = pcs[:, :3]
        coords[:, 0] = 0
        coords = torch.as_tensor(coords, dtype=torch.int)
        feats = torch.as_tensor(feats[indices], dtype=torch.float)
        input = SparseTensor(coords=coords, feats=feats).to(device)

        # forward
        outputs = model(input)
        out_conv = conv3d(outputs[-1])
        #  @kentang-mit: expect out_conv to have coordinates more than 100
        # To be exact, the neighborhood of coords, with kernel size 3, but currently, the out_conv.C.shape[0] == 100

        # print feature shapes
        for k, output in enumerate(outputs):
            print(f'output[{k}].F.shape = {output.feats.shape}')

if __name__ == '__main__':
    main()

I expect out_conv to have coordinates that are activated points due to sparse convolution of out with kernel size 3. However, the number of points in out_conv match those of out, which means that the convolution does not generate new points.

Could you help me with the generative argument? I'm not quite sure how to use it without the documentation. I really appreciate your help :)