Method test_gpu

tests/python/convolution.py:249–298

(self)

class TestConvolutionMode(unittest.TestCase):
    def test_gpu(self):
        print(f"{self.__class__.__name__}: test_gpu")
        if not torch.cuda.is_available():
            return
        in_channels, out_channels, D = 3, 2, 2
        coords, feats, labels = data_loader(in_channels, batch_size=20)
        feats = feats.double()
        feats.requires_grad_()
        device = torch.device("cuda")
        conv = (
            MinkowskiConvolution(
                in_channels,
                out_channels,
                kernel_size=2,
                stride=1,
                bias=False,
                dimension=D,
            )
            .to(device)
            .double()
        )
        # Initialize context
        for mode in [_C.ConvolutionMode.DIRECT_GEMM, _C.ConvolutionMode.COPY_GEMM]:
            conv.convolution_mode = mode
            input = SparseTensor(feats, coordinates=coords, device=device)
            print(mode, input.F.numel(), len(input), input)
            output = conv(input)
            print(output)

            # Check backward
            fn = MinkowskiConvolutionFunction()

            grad = output.F.clone().zero_()
            grad[0] = 1
            output.F.backward(grad)

            self.assertTrue(
                gradcheck(
                    fn,
                    (
                        input.F,
                        conv.kernel,
                        conv.kernel_generator,
                        conv.convolution_mode,
                        input.coordinate_map_key,
                        None,
                        input.coordinate_manager,
                    ),
                )
            )


class TestConvolutionTranspose(unittest.TestCase):