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 365 forks source link

How to recompose sparse tensor after batch-wise decomposition? #339

Open ghost opened 3 years ago

ghost commented 3 years ago

Hi,

I am trying to extend OverlapPredator.Mink to support batch training. Predator is composed of three modules: a Minkowski encoder, a GNN and a Minkowski decoder. After the encoder, I have to perform batch-wise decomposition for the GNN. Then after the GNN, I have to recompose a sparse tensor which shares the coordinate_manager and coordinate_key_map with the sparse tensor from the encoder. How should I implement this?

Thanks.

chrischoy commented 3 years ago

You can use the function ME.utils.batched_coordinates to recreate batched coordinates. For features, you can simply concatenate along the 0th dim.