About the kNN-based Transformer

[sz3623]

This is a new concept in the sparse convolution framework. There is a question about the point cloud data structure in this framework.

Point clouds belonging to different batches form a tensor in the sparse convolution framework, which will cause errors in the selection of k-nearest neighbors. In your work, taking the batch index into account for the k-nearest neighbor distance calculation is an option.

However, weighting the batch index number with different scales (1024, 1024x1204, ...) is a more effective way. In addition, separation and k-nearest neighbor processing for each batch are also considered. Do you have a better way to deal with it?

[3dpcc]

If parallel training is needed, I recommend reshaping the sparse tensor before the KNN operation into coordinates (shape: b, n, 3) and features (shape: b, n, c), and then performing the PyTorch3D KNN operation. The KNN indices obtained on the coordinates can be used simultaneously for both the coordinates and the features. After the KNN attention is completed, the features and coordinates can be recombined into a sparse tensor.