| 40 | |
| 41 | |
| 42 | class MinkowskiConvolutionFunction(Function): |
| 43 | @staticmethod |
| 44 | def forward( |
| 45 | ctx, |
| 46 | input_features: torch.Tensor, |
| 47 | kernel_weights: torch.Tensor, |
| 48 | kernel_generator: KernelGenerator, |
| 49 | convolution_mode: ConvolutionMode, |
| 50 | in_coordinate_map_key: CoordinateMapKey, |
| 51 | out_coordinate_map_key: CoordinateMapKey = None, |
| 52 | coordinate_manager: CoordinateManager = None, |
| 53 | ): |
| 54 | if out_coordinate_map_key is None: |
| 55 | out_coordinate_map_key = CoordinateMapKey( |
| 56 | in_coordinate_map_key.get_coordinate_size() |
| 57 | ) |
| 58 | |
| 59 | input_features = input_features.contiguous() |
| 60 | |
| 61 | ctx.input_features = input_features |
| 62 | ctx.kernel_weights = kernel_weights |
| 63 | ctx.misc = [ |
| 64 | kernel_generator, |
| 65 | convolution_mode, |
| 66 | in_coordinate_map_key, |
| 67 | out_coordinate_map_key, |
| 68 | coordinate_manager, |
| 69 | ] |
| 70 | |
| 71 | fw_fn = get_minkowski_function("ConvolutionForward", input_features) |
| 72 | return fw_fn( |
| 73 | ctx.input_features, |
| 74 | kernel_weights, |
| 75 | kernel_generator.kernel_size, |
| 76 | kernel_generator.kernel_stride, |
| 77 | kernel_generator.kernel_dilation, |
| 78 | kernel_generator.region_type, |
| 79 | kernel_generator.region_offsets, |
| 80 | kernel_generator.expand_coordinates, |
| 81 | convolution_mode, |
| 82 | in_coordinate_map_key, |
| 83 | out_coordinate_map_key, |
| 84 | coordinate_manager._manager, |
| 85 | ) |
| 86 | |
| 87 | @staticmethod |
| 88 | def backward(ctx, grad_out_feat: torch.Tensor): |
| 89 | grad_out_feat = grad_out_feat.contiguous() |
| 90 | ( |
| 91 | kernel_generator, |
| 92 | convolution_mode, |
| 93 | in_coordinate_map_key, |
| 94 | out_coordinate_map_key, |
| 95 | coordinate_manager, |
| 96 | ) = ctx.misc |
| 97 | |
| 98 | bw_fn = get_minkowski_function("ConvolutionBackward", grad_out_feat) |
| 99 | grad_in_feat, grad_kernel = bw_fn( |