/* * Copyright (c) 2020 NVIDIA Corporation. * Copyright (c) 2018-2020 Chris Choy (chrischoy@ai.stanford.edu). * * Permission is hereby granted, free of charge, to any person obtaining a copy * of this software and associated documentation files (the "Software"), to deal * in the Software without restriction, including without limitation the rights * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell * copies of the Software, and to permit persons to whom the Software is * furnished to do so, subject to the following conditions: * * The above copyright notice and this permission notice shall be included in * all copies or substantial portions of the Software. * * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING * FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS * IN THE SOFTWARE. * * Please cite "4D Spatio-Temporal ConvNets: Minkowski Convolutional Neural * Networks", CVPR'19 (https://arxiv.org/abs/1904.08755) if you use any part * of the code. */ #include "coordinate_map.hpp" #include "coordinate_map_cpu.hpp" #include "coordinate_map_key.hpp" #include "coordinate_map_manager.hpp" #include "errors.hpp" #include "types.hpp" #include "utils.hpp" #ifndef CPU_ONLY #include "allocators.cuh" #include "coordinate_map_gpu.cuh" #include #endif #include #include #include #include #include namespace minkowski { /************************************* * Convolution *************************************/ template at::Tensor ConvolutionForwardCPU(at::Tensor const &in_feat, // at::Tensor const &kernel, // default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offset, // bool const expand_coordinates, // ConvolutionMode::Type const convolution_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // cpu_manager_type *p_map_manager); template std::pair ConvolutionBackwardCPU(at::Tensor const &in_feat, // at::Tensor &grad_out_feat, // at::Tensor const &kernel, // default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offsets, // ConvolutionMode::Type const convolution_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // cpu_manager_type *p_map_manager); #ifndef CPU_ONLY template class TemplatedAllocator> at::Tensor ConvolutionForwardGPU( at::Tensor const &in_feat, // at::Tensor const &kernel, // default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offset, // bool const expand_coordinates, // ConvolutionMode::Type const convolution_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // gpu_manager_type *p_map_manager); template class TemplatedAllocator> std::pair ConvolutionBackwardGPU( at::Tensor const &in_feat, // at::Tensor &grad_out_feat, // at::Tensor const &kernel, // default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offset, // ConvolutionMode::Type const convolution_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // gpu_manager_type *p_map_manager); #endif /************************************* * Convolution Transpose *************************************/ template at::Tensor ConvolutionTransposeForwardCPU( at::Tensor const &in_feat, // at::Tensor const &kernel, // default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offset, // bool const expand_coordinates, // ConvolutionMode::Type const convolution_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // cpu_manager_type *p_map_manager); template std::pair ConvolutionTransposeBackwardCPU( at::Tensor const &in_feat, // at::Tensor const &grad_out_feat, // at::Tensor const &kernel, // default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offsets, // ConvolutionMode::Type const convolution_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // cpu_manager_type *p_map_manager); #ifndef CPU_ONLY template class TemplatedAllocator> at::Tensor ConvolutionTransposeForwardGPU( at::Tensor const &in_feat, // at::Tensor const &kernel, // default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offset, // bool const expand_coordinates, // ConvolutionMode::Type const convolution_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // gpu_manager_type *p_map_manager); template class TemplatedAllocator> std::pair ConvolutionTransposeBackwardGPU( at::Tensor const &in_feat, // at::Tensor const &grad_out_feat, // at::Tensor const &kernel, // default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offset, // ConvolutionMode::Type const convolution_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // gpu_manager_type *p_map_manager); #endif /************************************* * Local Pooling *************************************/ template std::pair LocalPoolingForwardCPU(at::Tensor const &in_feat, default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offset, // PoolingMode::Type pooling_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // cpu_manager_type *p_map_manager); template at::Tensor LocalPoolingBackwardCPU(at::Tensor const &in_feat, // at::Tensor const &grad_out_feat, // at::Tensor const &num_nonzero, // default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offset, // PoolingMode::Type pooling_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // cpu_manager_type *p_map_manager); #ifndef CPU_ONLY template class TemplatedAllocator> std::pair LocalPoolingForwardGPU( at::Tensor const &in_feat, default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offset, // PoolingMode::Type pooling_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // gpu_manager_type *p_map_manager); template class TemplatedAllocator> at::Tensor LocalPoolingBackwardGPU( at::Tensor const &in_feat, // at::Tensor const &grad_out_feat, // at::Tensor const &num_nonzero, // default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offset, // PoolingMode::Type pooling_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // gpu_manager_type *p_map_manager); #endif /************************************* * Local Pooling Transpose *************************************/ template std::pair LocalPoolingTransposeForwardCPU( at::Tensor const &in_feat, default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offset, // bool generate_new_coordinates, // PoolingMode::Type pooling_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // cpu_manager_type *p_map_manager); template at::Tensor LocalPoolingTransposeBackwardCPU( at::Tensor const &in_feat, // at::Tensor const &grad_out_feat, // at::Tensor const &num_nonzero, // default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offset, // PoolingMode::Type pooling_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // cpu_manager_type *p_map_manager); #ifndef CPU_ONLY template class TemplatedAllocator> std::pair LocalPoolingTransposeForwardGPU( at::Tensor const &in_feat, default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offset, // bool generate_new_coordinates, // PoolingMode::Type pooling_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // gpu_manager_type *p_map_manager); template class TemplatedAllocator> at::Tensor LocalPoolingTransposeBackwardGPU( at::Tensor const &in_feat, // at::Tensor const &grad_out_feat, // at::Tensor const &num_nonzero, // default_types::stride_type const &kernel_size, // default_types::stride_type const &kernel_stride, // default_types::stride_type const &kernel_dilation, // RegionType::Type const region_type, // at::Tensor const &offset, // PoolingMode::Type pooling_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // gpu_manager_type *p_map_manager); #endif /************************************* * Global Pooling *************************************/ template std::tuple GlobalPoolingForwardCPU(at::Tensor const &in_feat, PoolingMode::Type const pooling_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // cpu_manager_type *p_map_manager); template at::Tensor GlobalPoolingBackwardCPU(at::Tensor const &in_feat, at::Tensor &grad_out_feat, at::Tensor const &num_nonzero, PoolingMode::Type const pooling_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // cpu_manager_type *p_map_manager); #ifndef CPU_ONLY template class TemplatedAllocator> std::tuple GlobalPoolingForwardGPU( at::Tensor const &in_feat, PoolingMode::Type const pooling_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // gpu_manager_type *p_map_manager); template class TemplatedAllocator> at::Tensor GlobalPoolingBackwardGPU( at::Tensor const &in_feat, // at::Tensor &grad_out_feat, // at::Tensor const &num_nonzero, // PoolingMode::Type const pooling_mode, // CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // gpu_manager_type *p_map_manager); #endif /************************************* * Broadcast *************************************/ template at::Tensor BroadcastForwardCPU(at::Tensor const &in_feat, at::Tensor const &in_feat_glob, BroadcastMode::Type const broadcast_mode, CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_glob_map_key, // cpu_manager_type *p_map_manager); template std::pair BroadcastBackwardCPU(at::Tensor const &in_feat, at::Tensor const &in_feat_glob, at::Tensor const &grad_out_feat, BroadcastMode::Type const op, CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_glob_map_key, // cpu_manager_type *p_map_manager); #ifndef CPU_ONLY template class TemplatedAllocator> at::Tensor BroadcastForwardGPU( at::Tensor const &in_feat, at::Tensor const &in_feat_glob, BroadcastMode::Type const broadcast_mode, CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_glob_map_key, // gpu_manager_type *p_map_manager); template class TemplatedAllocator> std::pair BroadcastBackwardGPU( at::Tensor const &in_feat, at::Tensor const &in_feat_glob, at::Tensor const &grad_out_feat, BroadcastMode::Type const op, CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_glob_map_key, // gpu_manager_type *p_map_manager); #endif /************************************* * Pruning *************************************/ template at::Tensor PruningForwardCPU(at::Tensor const &in_feat, // CPU feat at::Tensor const &keep, // uint8 / bool / byte CPU data CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // cpu_manager_type *p_map_manager); template at::Tensor PruningBackwardCPU(at::Tensor &grad_out_feat, // CPU out feat CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // cpu_manager_type *p_map_manager); #ifndef CPU_ONLY template class TemplatedAllocator> at::Tensor PruningForwardGPU( at::Tensor const &in_feat, // GPU feat at::Tensor const &keep, // uint8 CPU data CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // gpu_manager_type *p_map_manager); template class TemplatedAllocator> at::Tensor PruningBackwardGPU( at::Tensor &grad_out_feat, // GPU out feat CoordinateMapKey *p_in_map_key, // CoordinateMapKey *p_out_map_key, // gpu_manager_type *p_map_manager); #endif /************************************* * Interpolation *************************************/ template std::vector InterpolationForwardCPU(at::Tensor const &in_feat, // at::Tensor const &tfield, // CoordinateMapKey *p_in_map_key, // cpu_manager_type *p_map_manager); template at::Tensor InterpolationBackwardCPU(at::Tensor &grad_out_feat, // at::Tensor const &in_map, // at::Tensor const &out_map, // at::Tensor const &weight, // CoordinateMapKey *p_in_map_key, // cpu_manager_type *p_map_manager); #ifndef CPU_ONLY template class TemplatedAllocator> std::vector InterpolationForwardGPU( at::Tensor const &in_feat, // at::Tensor const &tfield, // CoordinateMapKey *p_in_map_key, // gpu_manager_type *p_map_manager); template class TemplatedAllocator> at::Tensor InterpolationBackwardGPU( at::Tensor &grad_out_feat, // at::Tensor const &in_maps, // at::Tensor const &out_maps, // at::Tensor const &weights, // CoordinateMapKey *p_in_map_key, // gpu_manager_type *p_map_manager); #endif /************************************* * Quantization *************************************/ std::vector quantize_np( py::array_t coords); std::vector quantize_th(at::Tensor &coords); std::vector quantize_label_np( py::array_t coords, py::array_t labels, int invalid_label); std::vector quantize_label_th(at::Tensor coords, at::Tensor labels, int invalid_label); std::pair max_pool_fw(torch::Tensor const &in_map, // torch::Tensor const &out_map, // torch::Tensor const &in_feat, // int const out_nrows, bool const is_sorted); torch::Tensor max_pool_bw(torch::Tensor const &grad_out_feat, // torch::Tensor const &mask_index, // int const in_nrows); #ifndef CPU_ONLY template torch::Tensor coo_spmm(torch::Tensor const &rows, torch::Tensor const &cols, torch::Tensor const &vals, int64_t const dim_i, int64_t const dim_j, torch::Tensor const &mat2, int64_t const spmm_algorithm_id, bool const is_sorted); template std::vector // output, sorted rows, sorted cols, sorted vals. coo_spmm_average(torch::Tensor const &rows, torch::Tensor const &cols, int64_t const dim_i, int64_t const dim_j, torch::Tensor const &mat2, int64_t const spmm_algorithm_id); std::pair get_memory_info(); #endif } // end namespace minkowski namespace py = pybind11; template void instantiate_cpu_func(py::module &m, const std::string &dtypestr) { m.def((std::string("ConvolutionForwardCPU") + dtypestr).c_str(), &minkowski::ConvolutionForwardCPU, py::call_guard()); m.def((std::string("ConvolutionBackwardCPU") + dtypestr).c_str(), &minkowski::ConvolutionBackwardCPU, py::call_guard()); m.def((std::string("ConvolutionTransposeForwardCPU") + dtypestr).c_str(), &minkowski::ConvolutionTransposeForwardCPU, py::call_guard()); m.def((std::string("ConvolutionTransposeBackwardCPU") + dtypestr).c_str(), &minkowski::ConvolutionTransposeBackwardCPU, py::call_guard()); m.def((std::string("LocalPoolingForwardCPU") + dtypestr).c_str(), &minkowski::LocalPoolingForwardCPU, py::call_guard()); m.def((std::string("LocalPoolingBackwardCPU") + dtypestr).c_str(), &minkowski::LocalPoolingBackwardCPU, py::call_guard()); m.def((std::string("LocalPoolingTransposeForwardCPU") + dtypestr).c_str(), &minkowski::LocalPoolingTransposeForwardCPU, py::call_guard()); m.def((std::string("LocalPoolingTransposeBackwardCPU") + dtypestr).c_str(), &minkowski::LocalPoolingTransposeBackwardCPU, py::call_guard()); m.def((std::string("GlobalPoolingForwardCPU") + dtypestr).c_str(), &minkowski::GlobalPoolingForwardCPU, py::call_guard()); m.def((std::string("GlobalPoolingBackwardCPU") + dtypestr).c_str(), &minkowski::GlobalPoolingBackwardCPU, py::call_guard()); m.def((std::string("PruningForwardCPU") + dtypestr).c_str(), &minkowski::PruningForwardCPU, py::call_guard()); m.def((std::string("PruningBackwardCPU") + dtypestr).c_str(), &minkowski::PruningBackwardCPU, py::call_guard()); m.def((std::string("BroadcastForwardCPU") + dtypestr).c_str(), &minkowski::BroadcastForwardCPU, py::call_guard()); m.def((std::string("BroadcastBackwardCPU") + dtypestr).c_str(), &minkowski::BroadcastBackwardCPU, py::call_guard()); m.def((std::string("InterpolationForwardCPU") + dtypestr).c_str(), &minkowski::InterpolationForwardCPU, py::call_guard()); m.def((std::string("InterpolationBackwardCPU") + dtypestr).c_str(), &minkowski::InterpolationBackwardCPU, py::call_guard()); } #ifndef CPU_ONLY template class TemplatedAllocator> void instantiate_gpu_func(py::module &m, const std::string &dtypestr) { m.def((std::string("ConvolutionForwardGPU") + dtypestr).c_str(), &minkowski::ConvolutionForwardGPU, py::call_guard()); m.def((std::string("ConvolutionBackwardGPU") + dtypestr).c_str(), &minkowski::ConvolutionBackwardGPU, py::call_guard()); m.def((std::string("ConvolutionTransposeForwardGPU") + dtypestr).c_str(), &minkowski::ConvolutionTransposeForwardGPU, py::call_guard()); m.def((std::string("ConvolutionTransposeBackwardGPU") + dtypestr).c_str(), &minkowski::ConvolutionTransposeBackwardGPU, py::call_guard()); m.def((std::string("LocalPoolingForwardGPU") + dtypestr).c_str(), &minkowski::LocalPoolingForwardGPU, py::call_guard()); m.def( (std::string("LocalPoolingBackwardGPU") + dtypestr).c_str(), &minkowski::LocalPoolingBackwardGPU, py::call_guard()); m.def((std::string("LocalPoolingTransposeForwardGPU") + dtypestr).c_str(), &minkowski::LocalPoolingTransposeForwardGPU, py::call_guard()); m.def((std::string("LocalPoolingTransposeBackwardGPU") + dtypestr).c_str(), &minkowski::LocalPoolingTransposeBackwardGPU, py::call_guard()); m.def( (std::string("GlobalPoolingForwardGPU") + dtypestr).c_str(), &minkowski::GlobalPoolingForwardGPU, py::call_guard()); m.def( (std::string("GlobalPoolingBackwardGPU") + dtypestr).c_str(), &minkowski::GlobalPoolingBackwardGPU, py::call_guard()); m.def((std::string("PruningForwardGPU") + dtypestr).c_str(), &minkowski::PruningForwardGPU, py::call_guard()); m.def((std::string("PruningBackwardGPU") + dtypestr).c_str(), &minkowski::PruningBackwardGPU, py::call_guard()); m.def((std::string("BroadcastForwardGPU") + dtypestr).c_str(), &minkowski::BroadcastForwardGPU, py::call_guard()); m.def((std::string("BroadcastBackwardGPU") + dtypestr).c_str(), &minkowski::BroadcastBackwardGPU, py::call_guard()); m.def( (std::string("InterpolationForwardGPU") + dtypestr).c_str(), &minkowski::InterpolationForwardGPU, py::call_guard()); m.def( (std::string("InterpolationBackwardGPU") + dtypestr).c_str(), &minkowski::InterpolationBackwardGPU, py::call_guard()); } #endif void non_templated_cpu_func(py::module &m) { m.def("quantize_np", &minkowski::quantize_np); m.def("quantize_th", &minkowski::quantize_th); m.def("quantize_label_np", &minkowski::quantize_label_np); m.def("quantize_label_th", &minkowski::quantize_label_th); m.def("direct_max_pool_fw", &minkowski::max_pool_fw, py::call_guard()); m.def("direct_max_pool_bw", &minkowski::max_pool_bw, py::call_guard()); } #ifndef CPU_ONLY void non_templated_gpu_func(py::module &m) { m.def("coo_spmm_int32", &minkowski::coo_spmm, py::call_guard()); m.def("coo_spmm_average_int32", &minkowski::coo_spmm_average, py::call_guard()); } #endif void initialize_non_templated_classes(py::module &m) { // Enums py::enum_( m, "GPUMemoryAllocatorType") .value("PYTORCH", minkowski::GPUMemoryAllocatorBackend::Type::PYTORCH) .value("CUDA", minkowski::GPUMemoryAllocatorBackend::Type::CUDA) .export_values(); py::enum_(m, "CUDAKernelMapMode") .value("MEMORY_EFFICIENT", minkowski::CUDAKernelMapMode::Mode::MEMORY_EFFICIENT) .value("SPEED_OPTIMIZED", minkowski::CUDAKernelMapMode::Mode::SPEED_OPTIMIZED) .export_values(); py::enum_(m, "MinkowskiAlgorithm") .value("DEFAULT", minkowski::MinkowskiAlgorithm::Mode::DEFAULT) .value("MEMORY_EFFICIENT", minkowski::MinkowskiAlgorithm::Mode::MEMORY_EFFICIENT) .value("SPEED_OPTIMIZED", minkowski::MinkowskiAlgorithm::Mode::SPEED_OPTIMIZED) .export_values(); py::enum_(m, "CoordinateMapType") .value("CPU", minkowski::CoordinateMapBackend::Type::CPU) .value("CUDA", minkowski::CoordinateMapBackend::Type::CUDA) .export_values(); py::enum_(m, "RegionType") .value("HYPER_CUBE", minkowski::RegionType::Type::HYPER_CUBE) .value("HYPER_CROSS", minkowski::RegionType::Type::HYPER_CROSS) .value("CUSTOM", minkowski::RegionType::Type::CUSTOM) .export_values(); py::enum_(m, "PoolingMode") .value("LOCAL_SUM_POOLING", minkowski::PoolingMode::Type::LOCAL_SUM_POOLING) .value("LOCAL_AVG_POOLING", minkowski::PoolingMode::Type::LOCAL_AVG_POOLING) .value("LOCAL_MAX_POOLING", minkowski::PoolingMode::Type::LOCAL_MAX_POOLING) .value("GLOBAL_SUM_POOLING_DEFAULT", minkowski::PoolingMode::Type::GLOBAL_SUM_POOLING_DEFAULT) .value("GLOBAL_AVG_POOLING_DEFAULT", minkowski::PoolingMode::Type::GLOBAL_AVG_POOLING_DEFAULT) .value("GLOBAL_MAX_POOLING_DEFAULT", minkowski::PoolingMode::Type::GLOBAL_MAX_POOLING_DEFAULT) .value("GLOBAL_SUM_POOLING_KERNEL", minkowski::PoolingMode::Type::GLOBAL_SUM_POOLING_KERNEL) .value("GLOBAL_AVG_POOLING_KERNEL", minkowski::PoolingMode::Type::GLOBAL_AVG_POOLING_KERNEL) .value("GLOBAL_MAX_POOLING_KERNEL", minkowski::PoolingMode::Type::GLOBAL_MAX_POOLING_KERNEL) .value("GLOBAL_SUM_POOLING_PYTORCH_INDEX", minkowski::PoolingMode::Type::GLOBAL_SUM_POOLING_PYTORCH_INDEX) .value("GLOBAL_AVG_POOLING_PYTORCH_INDEX", minkowski::PoolingMode::Type::GLOBAL_AVG_POOLING_PYTORCH_INDEX) .value("GLOBAL_MAX_POOLING_PYTORCH_INDEX", minkowski::PoolingMode::Type::GLOBAL_MAX_POOLING_PYTORCH_INDEX) .export_values(); py::enum_(m, "BroadcastMode") .value("ELEMENTWISE_ADDITON", minkowski::BroadcastMode::Type::ELEMENTWISE_ADDITON) .value("ELEMENTWISE_MULTIPLICATION", minkowski::BroadcastMode::Type::ELEMENTWISE_MULTIPLICATION) .export_values(); py::enum_(m, "ConvolutionMode") .value("DEFAULT", minkowski::ConvolutionMode::Type::DEFAULT) .value("DIRECT_GEMM", minkowski::ConvolutionMode::Type::DIRECT_GEMM) .value("COPY_GEMM", minkowski::ConvolutionMode::Type::COPY_GEMM) .export_values(); // Classes py::class_(m, "CoordinateMapKey") .def(py::init()) .def(py::init()) .def("__repr__", &minkowski::CoordinateMapKey::to_string) .def("__hash__", &minkowski::CoordinateMapKey::hash) .def("is_key_set", &minkowski::CoordinateMapKey::is_key_set) .def("get_coordinate_size", &minkowski::CoordinateMapKey::get_coordinate_size) .def("get_key", &minkowski::CoordinateMapKey::get_key) .def("set_key", (void (minkowski::CoordinateMapKey::*)( minkowski::default_types::stride_type, std::string)) & minkowski::CoordinateMapKey::set_key) .def("set_key", (void (minkowski::CoordinateMapKey::*)( minkowski::coordinate_map_key_type const &)) & minkowski::CoordinateMapKey::set_key) .def("get_tensor_stride", &minkowski::CoordinateMapKey::get_tensor_stride) .def("__eq__", [](const minkowski::CoordinateMapKey &self, const minkowski::CoordinateMapKey &other) { return self == other; }); //.def(py::self == py::self); } template void instantiate_manager(py::module &m, const std::string &dtypestr) { py::class_( m, (std::string("CoordinateMapManager") + dtypestr).c_str()) .def(py::init<>()) .def(py::init()) .def("__repr__", py::overload_cast<>(&manager_type::to_string, py::const_)) .def("print_coordinate_map", py::overload_cast( &manager_type::to_string, py::const_)) .def("insert_and_map", &manager_type::insert_and_map) .def("insert_field", &manager_type::insert_field) .def("field_to_sparse_map", &manager_type::field_to_sparse_map) .def("field_to_sparse_insert_and_map", &manager_type::field_to_sparse_insert_and_map) .def("exists_field_to_sparse", py::overload_cast( &manager_type::exists_field_to_sparse, py::const_)) .def("get_field_to_sparse_map", &manager_type::get_field_to_sparse_map) .def("stride", &manager_type::py_stride) .def("origin", &manager_type::py_origin) .def("origin_field", &manager_type::py_origin_field) .def("get_coordinates", &manager_type::get_coordinates) .def("get_coordinate_field", &manager_type::get_coordinate_field) .def("get_coordinate_map_keys", &manager_type::get_coordinate_map_keys) .def("field_to_sparse_keys", &manager_type::field_to_sparse_keys) .def("size", py::overload_cast( &manager_type::size, py::const_)) .def("get_random_string_id", &manager_type::get_random_string_id) .def("origin_map_size", &manager_type::origin_map_size) .def("origin_map", &manager_type::origin_map_th) .def("origin_field_map", &manager_type::origin_field_map_th) .def("union_map", &manager_type::union_map_th) .def("stride_map", &manager_type::stride_map_th) .def("kernel_map", &manager_type::kernel_map_th) .def("interpolation_map_weight", &manager_type::interpolation_map_weight); } bool is_cuda_available() { #ifndef CPU_ONLY return true; #else return false; #endif } int cuda_version() { #if defined(CUDA_VERSION) return CUDA_VERSION; #else return -1; #endif } int cudart_version() { #if defined(CUDART_VERSION) return CUDART_VERSION; #else return -1; #endif } std::pair get_gpu_memory_info() { #ifndef CPU_ONLY return minkowski::get_memory_info(); #else return std::make_pair(0, 0); #endif }