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

GlobalAvgPooling for TensorField object. #534

Open Byeol-Haneul opened 1 year ago

Byeol-Haneul commented 1 year ago

Hello, I am trying out the Minkowski-PointNet example provided, but with the GlobalAvgPooling(GAP) layer instead of the GlobalMaxPooling layer. However, when implementing, I found out that the GAP of the TensorField object is not natively supported as shown below.

class MinkowskiGlobalPooling(MinkowskiModuleBase):
   ...
    def forward(
        self,
        input: SparseTensor,
        coordinates: Union[torch.IntTensor, CoordinateMapKey, SparseTensor] = None,
    ):
        # Get a new coordinate map key or extract one from the coordinates
        out_coordinate_map_key = _get_coordinate_map_key(input, coordinates)
        output = self.pooling.apply(
            input.F,
            self.pooling_mode,
            input.coordinate_map_key,
            out_coordinate_map_key,
            input._manager,
        )
        return SparseTensor(
            output,
            coordinate_map_key=out_coordinate_map_key,
            coordinate_manager=input.coordinate_manager,
        )

Adding a snippet of code borrowed from the forward function of MinkowskiGlobalMaxPooling in order to process TensorField object, I found out that the following error rises.

    def forward(
        self,
        input: SparseTensor,
        coordinates: Union[torch.IntTensor, CoordinateMapKey, SparseTensor] = None,
    ):
        # Get a new coordinate map key or extract one from the coordinates
        if isinstance(input, ME.TensorField):
            in_coordinate_map_key = input.coordinate_field_map_key
            out_coordinate_map_key = CoordinateMapKey(
                input.coordinate_field_map_key.get_coordinate_size()
            )
            output = self.pooling.apply(
            input.F,
            self.pooling_mode,
            in_coordinate_map_key,
            out_coordinate_map_key,
            input._manager,
            )
        else:
            out_coordinate_map_key = _get_coordinate_map_key(input, coordinates)
            output = self.pooling.apply(
            input.F,
            self.pooling_mode,
            input.coordinate_map_key,
            out_coordinate_map_key,
            input._manager,
            )

        return SparseTensor(
            output,
            coordinate_map_key=out_coordinate_map_key,
            coordinate_manager=input.coordinate_manager,
        )

MinkowskiEngine/src/coordinate_map_manager.cpp:899, assertion (exists(p_in_map_key)) failed. CoordinateMap not found

Is there a way to bypass this issue? Or is there something that I have missed out? Thank you in advance.

Byeol-Haneul commented 1 year ago

I do not think operating with dense tensors is the focus of this engine, and therefore I believe I can stick with a temporary solution: convert to SparseTensors, but with a very fine sampling rate. This can be done by rescaling the position info.