Upload folder using huggingface_hub

.gitattributes

@@ -1 +1,15 @@
 *.ckpt filter=lfs diff=lfs merge=lfs -text
+MinkowskiEngine/docs/images/classification_3d_net.png filter=lfs diff=lfs merge=lfs -text
+MinkowskiEngine/docs/images/conv_dense.gif filter=lfs diff=lfs merge=lfs -text
+MinkowskiEngine/docs/images/conv_generalized.gif filter=lfs diff=lfs merge=lfs -text
+MinkowskiEngine/docs/images/conv_sparse.gif filter=lfs diff=lfs merge=lfs -text
+MinkowskiEngine/docs/images/conv_sparse_conv.gif filter=lfs diff=lfs merge=lfs -text
+MinkowskiEngine/docs/images/detection_3d_net.png filter=lfs diff=lfs merge=lfs -text
+MinkowskiEngine/docs/images/generative_3d_results.gif filter=lfs diff=lfs merge=lfs -text
+MinkowskiEngine/docs/images/kernel_map.gif filter=lfs diff=lfs merge=lfs -text
+MinkowskiEngine/docs/images/segmentation.png filter=lfs diff=lfs merge=lfs -text
+MinkowskiEngine/docs/images/segmentation_3d_net.png filter=lfs diff=lfs merge=lfs -text
+diff-gaussian-rasterization-modified-main/third_party/glm/doc/manual/frontpage1.png filter=lfs diff=lfs merge=lfs -text
+diff-gaussian-rasterization-modified-main/third_party/glm/doc/manual/frontpage2.png filter=lfs diff=lfs merge=lfs -text
+diff-gaussian-rasterization-modified-main/third_party/glm/doc/manual.pdf filter=lfs diff=lfs merge=lfs -text
+visualization/depth_comparison.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,180 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+output.txt
+checkpoints/*
+outputs/*
+models/*
+datasets/*
+pretrained/*
+.idea/*
+
+test.sh
+test.ipynb
+test.py
+# test/
+
+*.zip
+final/
+wandb/
+*.mp4

.hydra/config.yaml

@@ -0,0 +1,193 @@
+dataset:
+  view_sampler:
+    name: boundedv2
+    num_target_views: 8
+    num_context_views: 6
+    min_distance_between_context_views: 20
+    max_distance_between_context_views: 50
+    max_distance_to_context_views: 0
+    context_gap_warm_up_steps: 10000
+    target_gap_warm_up_steps: 0
+    initial_min_distance_between_context_views: 15
+    initial_max_distance_between_context_views: 30
+    initial_max_distance_to_context_views: 0
+    extra_views_sampling_strategy: farthest_point
+    target_views_replace_sample: false
+  name: dl3dv
+  roots:
+  - datasets/dl3dv
+  make_baseline_1: false
+  augment: true
+  image_shape:
+  - 270
+  - 480
+  background_color:
+  - 0.0
+  - 0.0
+  - 0.0
+  cameras_are_circular: false
+  baseline_epsilon: 0.001
+  max_fov: 100.0
+  skip_bad_shape: true
+  near: 0.5
+  far: 200.0
+  baseline_scale_bounds: false
+  shuffle_val: true
+  test_len: -1
+  test_chunk_interval: 1
+  sort_target_index: true
+  sort_context_index: true
+  train_times_per_scene: 1
+  test_times_per_scene: 1
+  ori_image_shape:
+  - 270
+  - 480
+  overfit_max_views: 148
+  use_index_to_load_chunk: false
+  mix_tartanair: false
+  no_mix_test_set: true
+  load_depth: false
+  overfit_to_scene: null
+  min_views: 1
+  max_views: 6
+  highres: false
+model:
+  encoder:
+    name: depthsplat
+    num_depth_candidates: 128
+    num_surfaces: 1
+    gaussians_per_pixel: 1
+    gaussian_adapter:
+      gaussian_scale_min: 1.0e-10
+      gaussian_scale_max: 3.0
+      sh_degree: 2
+    d_feature: 128
+    visualizer:
+      num_samples: 8
+      min_resolution: 256
+      export_ply: false
+    unimatch_weights_path: pretrained/gmdepth-scale1-resumeflowthings-scannet-5d9d7964.pth
+    multiview_trans_attn_split: 2
+    costvolume_unet_feat_dim: 128
+    costvolume_unet_channel_mult:
+    - 1
+    - 1
+    - 1
+    costvolume_unet_attn_res:
+    - 4
+    depth_unet_feat_dim: 32
+    depth_unet_attn_res:
+    - 16
+    depth_unet_channel_mult:
+    - 1
+    - 1
+    - 1
+    - 1
+    - 1
+    downscale_factor: 4
+    shim_patch_size: 16
+    local_mv_match: 2
+    monodepth_vit_type: vitb
+    supervise_intermediate_depth: true
+    return_depth: true
+    num_scales: 2
+    upsample_factor: 4
+    lowest_feature_resolution: 8
+    depth_unet_channels: 128
+    grid_sample_disable_cudnn: false
+    large_gaussian_head: false
+    color_large_unet: false
+    init_sh_input_img: true
+    feature_upsampler_channels: 64
+    gaussian_regressor_channels: 64
+    train_depth_only: false
+  decoder:
+    name: splatting_cuda
+loss:
+  mse:
+    weight: 1.0
+  lpips:
+    weight: 0.05
+    apply_after_step: 0
+wandb:
+  project: depthsplat
+  entity: placeholder
+  name: dl3dv
+  mode: disabled
+  id: null
+  tags:
+  - dl3dv
+  - 270x480
+mode: train
+data_loader:
+  train:
+    num_workers: 10
+    persistent_workers: true
+    batch_size: 1
+    seed: 1234
+  test:
+    num_workers: 4
+    persistent_workers: false
+    batch_size: 1
+    seed: 2345
+  val:
+    num_workers: 1
+    persistent_workers: true
+    batch_size: 1
+    seed: 3456
+optimizer:
+  lr: 0.0002
+  lr_monodepth: 2.0e-06
+  warm_up_steps: 2000
+  weight_decay: 0.01
+checkpointing:
+  load: null
+  every_n_train_steps: 5000
+  save_top_k: 5
+  pretrained_model: pretrained/depthsplat-gs-base-re10k-256x448-view2-76a0605a.pth
+  pretrained_monodepth: null
+  pretrained_mvdepth: null
+  pretrained_depth: null
+  no_strict_load: false
+  resume: false
+train:
+  depth_mode: null
+  extended_visualization: false
+  print_log_every_n_steps: 100
+  eval_model_every_n_val: 2
+  eval_data_length: 999999
+  eval_deterministic: false
+  eval_time_skip_steps: 3
+  eval_save_model: true
+  l1_loss: false
+  intermediate_loss_weight: 0.9
+  no_viz_video: false
+  viz_depth: false
+  forward_depth_only: false
+  train_ignore_large_loss: 0.0
+  no_log_projections: false
+test:
+  output_path: outputs/test
+  compute_scores: true
+  eval_time_skip_steps: 0
+  save_image: false
+  save_video: false
+  save_gt_image: false
+  save_input_images: false
+  save_depth: false
+  save_depth_npy: false
+  save_depth_concat_img: false
+  save_gaussian: false
+  render_chunk_size: null
+  stablize_camera: false
+  stab_camera_kernel: 50
+  dec_chunk_size: 30
+seed: 111123
+trainer:
+  max_steps: 100000
+  val_check_interval: 0.5
+  gradient_clip_val: 0.5
+  num_sanity_val_steps: 2
+  num_nodes: 1
+output_dir: checkpoints/dl3dv-256x448-depthsplat-base-randview2-6
+use_plugins: false

.hydra/hydra.yaml

@@ -0,0 +1,178 @@
+hydra:
+  run:
+    dir: /mnt/pfs/users/wangweijie/yeqing/BEV-Splat
+  sweep:
+    dir: multirun/${now:%Y-%m-%d}/${now:%H-%M-%S}
+    subdir: ${hydra.job.num}
+  launcher:
+    _target_: hydra._internal.core_plugins.basic_launcher.BasicLauncher
+  sweeper:
+    _target_: hydra._internal.core_plugins.basic_sweeper.BasicSweeper
+    max_batch_size: null
+    params: null
+  help:
+    app_name: ${hydra.job.name}
+    header: '${hydra.help.app_name} is powered by Hydra.
+
+      '
+    footer: 'Powered by Hydra (https://hydra.cc)
+
+      Use --hydra-help to view Hydra specific help
+
+      '
+    template: '${hydra.help.header}
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (group=option)
+
+
+      $APP_CONFIG_GROUPS
+
+
+      == Config ==
+
+      Override anything in the config (foo.bar=value)
+
+
+      $CONFIG
+
+
+      ${hydra.help.footer}
+
+      '
+  hydra_help:
+    template: 'Hydra (${hydra.runtime.version})
+
+      See https://hydra.cc for more info.
+
+
+      == Flags ==
+
+      $FLAGS_HELP
+
+
+      == Configuration groups ==
+
+      Compose your configuration from those groups (For example, append hydra/job_logging=disabled
+      to command line)
+
+
+      $HYDRA_CONFIG_GROUPS
+
+
+      Use ''--cfg hydra'' to Show the Hydra config.
+
+      '
+    hydra_help: ???
+  hydra_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][HYDRA] %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+    root:
+      level: INFO
+      handlers:
+      - console
+    loggers:
+      logging_example:
+        level: DEBUG
+    disable_existing_loggers: false
+  job_logging:
+    version: 1
+    formatters:
+      simple:
+        format: '[%(asctime)s][%(name)s][%(levelname)s] - %(message)s'
+    handlers:
+      console:
+        class: logging.StreamHandler
+        formatter: simple
+        stream: ext://sys.stdout
+      file:
+        class: logging.FileHandler
+        formatter: simple
+        filename: ${hydra.runtime.output_dir}/${hydra.job.name}.log
+    root:
+      level: INFO
+      handlers:
+      - console
+      - file
+    disable_existing_loggers: false
+  env: {}
+  mode: RUN
+  searchpath: []
+  callbacks: {}
+  output_subdir: .hydra
+  overrides:
+    hydra:
+    - hydra.run.dir="/mnt/pfs/users/wangweijie/yeqing/BEV-Splat"
+    - hydra.job.name=train_ddp_process_1
+    - hydra.mode=RUN
+    task:
+    - +experiment=dl3dv
+    - data_loader.train.batch_size=1
+    - dataset.roots=["datasets/dl3dv"]
+    - dataset.view_sampler.num_target_views=8
+    - dataset.view_sampler.num_context_views=6
+    - dataset.min_views=1
+    - dataset.max_views=6
+    - trainer.max_steps=100000
+    - trainer.num_nodes=1
+    - model.encoder.num_scales=2
+    - model.encoder.upsample_factor=4
+    - model.encoder.lowest_feature_resolution=8
+    - model.encoder.monodepth_vit_type=vitb
+    - checkpointing.pretrained_model=pretrained/depthsplat-gs-base-re10k-256x448-view2-76a0605a.pth
+    - wandb.project=depthsplat
+    - output_dir=checkpoints/dl3dv-256x448-depthsplat-base-randview2-6
+  job:
+    name: train_ddp_process_1
+    chdir: null
+    override_dirname: +experiment=dl3dv,checkpointing.pretrained_model=pretrained/depthsplat-gs-base-re10k-256x448-view2-76a0605a.pth,data_loader.train.batch_size=1,dataset.max_views=6,dataset.min_views=1,dataset.roots=["datasets/dl3dv"],dataset.view_sampler.num_context_views=6,dataset.view_sampler.num_target_views=8,model.encoder.lowest_feature_resolution=8,model.encoder.monodepth_vit_type=vitb,model.encoder.num_scales=2,model.encoder.upsample_factor=4,output_dir=checkpoints/dl3dv-256x448-depthsplat-base-randview2-6,trainer.max_steps=100000,trainer.num_nodes=1,wandb.project=depthsplat
+    id: ???
+    num: ???
+    config_name: main
+    env_set: {}
+    env_copy: []
+    config:
+      override_dirname:
+        kv_sep: '='
+        item_sep: ','
+        exclude_keys: []
+  runtime:
+    version: 1.3.2
+    version_base: '1.3'
+    cwd: /mnt/pfs/users/wangweijie/yeqing/BEV-Splat
+    config_sources:
+    - path: hydra.conf
+      schema: pkg
+      provider: hydra
+    - path: /mnt/pfs/users/wangweijie/yeqing/BEV-Splat/config
+      schema: file
+      provider: main
+    - path: ''
+      schema: structured
+      provider: schema
+    output_dir: /mnt/pfs/users/wangweijie/yeqing/BEV-Splat
+    choices:
+      experiment: dl3dv
+      model/decoder: splatting_cuda
+      model/encoder: depthsplat
+      dataset: dl3dv
+      dataset/view_sampler: boundedv2_360
+      dataset/view_sampler_dataset_specific_config: boundedv2_360_dl3dv
+      hydra/env: default
+      hydra/callbacks: null
+      hydra/job_logging: default
+      hydra/hydra_logging: default
+      hydra/hydra_help: default
+      hydra/help: default
+      hydra/sweeper: basic
+      hydra/launcher: basic
+      hydra/output: default
+  verbose: false

.hydra/overrides.yaml

@@ -0,0 +1,16 @@
+- +experiment=dl3dv
+- data_loader.train.batch_size=1
+- dataset.roots=["datasets/dl3dv"]
+- dataset.view_sampler.num_target_views=8
+- dataset.view_sampler.num_context_views=6
+- dataset.min_views=1
+- dataset.max_views=6
+- trainer.max_steps=100000
+- trainer.num_nodes=1
+- model.encoder.num_scales=2
+- model.encoder.upsample_factor=4
+- model.encoder.lowest_feature_resolution=8
+- model.encoder.monodepth_vit_type=vitb
+- checkpointing.pretrained_model=pretrained/depthsplat-gs-base-re10k-256x448-view2-76a0605a.pth
+- wandb.project=depthsplat
+- output_dir=checkpoints/dl3dv-256x448-depthsplat-base-randview2-6

.vscode/launch.json

@@ -0,0 +1,18 @@
+
+{
+    // 使用 IntelliSense 了解相关属性。 
+    // 悬停以查看现有属性的描述。
+    // 欲了解更多信息，请访问: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "sh_file_debug",
+            "type": "debugpy",
+            "request": "attach",
+            "connect": {
+            "host": "localhost",
+            "port": 9326}
+        }
+    ]
+}
+

CLAUDE.md

@@ -0,0 +1,161 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Project Overview
+
+This is BEV-Splat, a 3D computer vision project based on DepthSplat that uses Gaussian splatting for novel view synthesis. The codebase is built with PyTorch Lightning and uses Hydra for configuration management. It incorporates MinkowskiEngine for sparse 3D convolution operations and integrates depth prediction models (Depth Anything V2) with multi-view depth estimation.
+
+## Development Environment Setup
+
+### Environment Creation
+```bash
+# Create conda environment from environment.yml 
+conda env create -f environment.yml
+conda activate scube
+
+# Alternative: Use depthsplat environment
+conda activate depthsplat
+conda install -c conda-forge openblas
+```
+
+### MinkowskiEngine Installation
+This project requires MinkowskiEngine for sparse 3D convolution. For CUDA versions >11, follow the specific installation guide referenced in the README.md.
+
+Test installation:
+```bash
+python -c "import MinkowskiEngine as ME; print(ME.__version__)"
+```
+
+## Common Development Commands
+
+### Training
+Basic training command structure:
+```bash
+python -m src.main +experiment=<dataset_name> \
+  data_loader.train.batch_size=<batch_size> \
+  trainer.max_steps=<steps> \
+  output_dir=<checkpoint_dir>
+```
+
+For DL3DV dataset training:
+```bash
+CUDA_VISIBLE_DEVICES=<gpu_id> python -m src.main +experiment=dl3dv \
+  data_loader.train.batch_size=1 \
+  'dataset.roots'='["datasets/dl3dv"]' \
+  dataset.view_sampler.num_target_views=8 \
+  dataset.view_sampler.num_context_views=6 \
+  trainer.max_steps=100000 \
+  model.encoder.monodepth_vit_type=vitb \
+  checkpointing.pretrained_monodepth=pretrained/pretrained_weights/depth_anything_v2_vitb.pth \
+  output_dir=checkpoints/<experiment_name>
+```
+
+### Testing/Inference
+```bash
+python -m src.main mode=test \
+  checkpointing.pretrained_model=<model_path> \
+  test.output_path=outputs/test \
+  test.save_image=true
+```
+
+### Dataset Processing
+Data conversion scripts are located in `src/scripts/`:
+- `convert_dl3dv_test.py` - Convert DL3DV test set
+- `convert_dl3dv_train.py` - Convert DL3DV training set  
+- `generate_dl3dv_index.py` - Generate index.json files
+
+## Architecture Overview
+
+### Core Components
+
+**Main Entry Point**: `src/main.py`
+- Hydra-based configuration management
+- PyTorch Lightning trainer setup
+- Multi-GPU DDP support with custom strategy for frozen parameters
+- Pretrained model loading for depth networks
+
+**Model Architecture**:
+- **Encoder** (`src/model/encoder/`): DepthSplat encoder with configurable scales and depth prediction
+- **Decoder** (`src/model/decoder/`): CUDA-based Gaussian splatting renderer
+- **ModelWrapper** (`src/model/model_wrapper.py`): Lightning module wrapper handling optimization and training logic
+
+**Data Pipeline**:
+- **DataModule** (`src/dataset/data_module.py`): Lightning data module
+- **View Samplers** (`src/dataset/view_sampler/`): Different sampling strategies (bounded, arbitrary, evaluation)
+- **Dataset Classes**: Support for RealEstate10K, DL3DV datasets with torch chunk format
+
+### Configuration System
+
+Uses Hydra with YAML configs in `config/`:
+- `main.yaml`: Base configuration  
+- `experiment/`: Dataset-specific experiment configs
+- `model/encoder/`: Encoder configurations
+- `model/decoder/`: Decoder configurations  
+- `dataset/`: Dataset and view sampler configurations
+
+Key config sections:
+- `wandb`: Experiment tracking (using SwanLab logger)
+- `optimizer`: Learning rates, warm-up, weight decay
+- `checkpointing`: Model loading/saving, pretrained weights
+- `trainer`: PyTorch Lightning trainer settings
+
+### Key Features
+
+**Pretrained Model Integration**:
+- Supports loading pretrained depth models (Depth Anything V2)
+- Multi-view depth estimation integration  
+- Parameter freezing/unfreezing during training
+- Custom callback for gradual parameter unfreezing
+
+**Multi-GPU Training**:
+- DDP strategy with `find_unused_parameters=True` for frozen parameters
+- Global seed management for distributed training
+- Custom parameter update checking
+
+**Visualization**:
+- Video rendering capabilities
+- Depth visualization
+- Camera trajectory visualization
+- Integration with visualization tools
+
+## Dataset Structure
+
+Expected dataset format (torch chunks):
+```
+datasets/
+├── re10k/
+│   ├── train/
+│   │   ├── 000000.torch
+│   │   └── index.json
+│   └── test/
+├── dl3dv/
+│   ├── train/
+│   └── test/
+```
+
+Use symbolic links to point to actual dataset locations:
+```bash
+ln -s /path/to/your/datasets datasets
+```
+
+## Important Implementation Notes
+
+### Camera Conventions
+The codebase uses specific camera coordinate conventions - refer to camera convention documentation when adding new datasets.
+
+### Memory Management  
+- Batch sizes are typically small (1-4) due to memory constraints
+- Use gradient accumulation for effective larger batch sizes
+- CUDA memory optimization for Gaussian splatting operations
+
+### Distributed Training Considerations
+- Custom DDP strategy handles frozen parameters correctly
+- Global seed setting with rank-specific PyTorch seeds
+- Proper barrier synchronization for distributed operations
+
+### Logging and Checkpointing
+- Uses SwanLab instead of WandB for experiment tracking
+- Automatic checkpoint saving every N steps
+- Supports resuming from latest checkpoint
+- Model parameter loading reports with detailed statistics

DATASETS.md

@@ -0,0 +1,78 @@
+# Datasets
+
+For view synthesis experiments with Gaussian splatting, we mainly use [RealEstate10K](https://google.github.io/realestate10k/index.html) and [DL3DV](https://github.com/DL3DV-10K/Dataset) datasets. We provide the data processing scripts to convert the original datasets to pytorch chunk files which can be directly loaded with this codebase. 
+
+Expected folder structure:
+
+```
+├── datasets
+│   ├── re10k
+│   ├── ├── train
+│   ├── ├── ├── 000000.torch
+│   ├── ├── ├── ...
+│   ├── ├── ├── index.json
+│   ├── ├── test
+│   ├── ├── ├── 000000.torch
+│   ├── ├── ├── ...
+│   ├── ├── ├── index.json
+│   ├── dl3dv
+│   ├── ├── train
+│   ├── ├── ├── 000000.torch
+│   ├── ├── ├── ...
+│   ├── ├── ├── index.json
+│   ├── ├── test
+│   ├── ├── ├── 000000.torch
+│   ├── ├── ├── ...
+│   ├── ├── ├── index.json
+```
+
+It's recommended to create a symbolic link from `YOUR_DATASET_PATH` to `datasets` using
+```
+ln -s YOUR_DATASET_PATH datasets
+```
+
+Or you can specify your dataset path with `dataset.roots=[YOUR_DATASET_PATH]/re10k` and `dataset.roots=[YOUR_DATASET_PATH]/dl3dv` in the config.
+
+We also provide instructions to convert additional datasets to the desired format.
+
+
+## RealEstate10K
+
+For experiments on RealEstate10K, we primarily follow [pixelSplat](https://github.com/dcharatan/pixelsplat) and [MVSplat](https://github.com/donydchen/mvsplat) to train and evaluate on the 256x256 resolution.
+
+Please refer to [pixelSplat repo](https://github.com/dcharatan/pixelsplat?tab=readme-ov-file#acquiring-datasets) for acquiring the processed 360p (360x640) dataset.
+
+If you would like to train and evaluate on the high-resolution RealEstate10K dataset, you will need to download the 720p (720x1280) version. Please refer to [here](https://github.com/yilundu/cross_attention_renderer/tree/master/data_download) for the downloading script. Note that the script by default downloads the 360p videos, you will need to modify the`360p` to `720p` in [this line of code](https://github.com/yilundu/cross_attention_renderer/blob/master/data_download/generate_realestate.py#L137) to download the 720p videos.
+
+After downloading the 720p dataset, you can use the scripts [here](https://github.com/dcharatan/real_estate_10k_tools/tree/main/src) to convert the dataset to the desired format in this codebase.
+
+Considering the full 720p dataset is quite large and may take time to download and process, we provide a preprocessed subset in `.torch` chunks ([download](https://huggingface.co/datasets/haofeixu/depthsplat/resolve/main/re10k_720p_test_subset.zip)) containing two test scenes to quickly run inference with our model.
+
+## DL3DV
+
+For experiments on DL3DV, we primarily train and evaluate at a resolution of 256×448. Additionally, we train high-resolution models (448×768) for qualitative results.
+
+For the test set, we use the [DL3DV-Benchmark](https://huggingface.co/datasets/DL3DV/DL3DV-Benchmark) split, which contains 140 scenes for evaluation. You can first use the script [src/scripts/convert_dl3dv_test.py](src/scripts/convert_dl3dv_test.py) to convert the test set, and then run [src/scripts/generate_dl3dv_index.py](src/scripts/generate_dl3dv_index.py) to generate the `index.json` file for the test set.
+
+For the training set, we use the [DL3DV-480p](https://huggingface.co/datasets/DL3DV/DL3DV-ALL-480P) dataset (270x480 resolution), where the 140 scenes in the test set are excluded during processing the training set. After downloading the [DL3DV-480p](https://huggingface.co/datasets/DL3DV/DL3DV-ALL-480P) dataset, You can first use the script [src/scripts/convert_dl3dv_train.py](src/scripts/convert_dl3dv_train.py) to convert the training set, and then run [src/scripts/generate_dl3dv_index.py](src/scripts/generate_dl3dv_index.py) to generate the `index.json` file for the training set.
+
+Please note that you will need to update the dataset paths in the aforementioned processing scripts.
+
+If you would like to train and evaluate on the high-resolution DL3DV dataset, you will need to download the [DL3DV-960P](https://huggingface.co/datasets/DL3DV/DL3DV-ALL-960P) version (540x960 resolution). Simply follow the same procedure for data processing, but update the `images_8` folder to `images_4`.
+
+Please follow the [DL3DV license](https://github.com/DL3DV-10K/Dataset/blob/main/License.md) if you use this dataset in your project and kindly [reference the DL3DV paper](https://github.com/DL3DV-10K/Dataset?tab=readme-ov-file#bibtex).
+
+Considering the full 960p dataset is quite large and may take time to download and process, we provide a preprocessed subset in `.torch` chunks ([download](https://huggingface.co/datasets/haofeixu/depthsplat/resolve/main/dl3dv_960p_test_subset.zip)) containing two test scenes to quickly run inference with our model. Please note that this released subset is intended solely for research purposes. We disclaim any responsibility for the misuse, inappropriate use, or unethical application of the dataset by individuals or entities who download or access it. We kindly ask users to adhere to the [DL3DV license](https://github.com/DL3DV-10K/Dataset/blob/main/License.md).
+
+
+## ACID
+
+
+We also evaluate our generalization on the [ACID](https://infinite-nature.github.io/) dataset. Note that we do not use the training set; you only need to [download the test set](http://schadenfreude.csail.mit.edu:8000/re10k_test_only.zip) (provided by [pixelSplat repo](https://github.com/dcharatan/pixelsplat?tab=readme-ov-file#acquiring-datasets)) for evaluation.
+
+
+
+## Additional Datasets
+
+If you would like to train and/or evaluate on additional datasets, just modify the [data processing scripts](src/scripts) to convert the dataset format. Kindly note the [camera conventions](https://github.com/cvg/depthsplat/tree/main?tab=readme-ov-file#camera-conventions) used in this codebase.
+

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Haofei Xu
+
+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.

MODEL_ZOO.md

@@ -0,0 +1,46 @@
+# Model Zoo
+
+- We provide pre-trained models for view synthesis with 3D Gaussian splatting and scale-consistent depth estimation from multi-view posed images.
+
+- We assume that the downloaded weights are stored in the `pretrained` directory. It's recommended to create a symbolic link from `YOUR_MODEL_PATH` to `pretrained` using
+```
+ln -s YOUR_MODEL_PATH pretrained
+```
+
+- To verify the integrity of downloaded files, each model on this page includes its [sha256sum](https://sha256sum.com/) prefix in the file name, which can be checked using the command `sha256sum filename`.
+
+
+## Gaussian Splatting
+
+- The models are trained on RealEstate10K (re10k) and/or DL3DV (dl3dv) datasets at resolutions of 256x256, 256x448, and 448x768. The number of training views ranges from 2 to 10.
+
+- The "→" symbol indicates that the models are trained in two stages. For example, "re10k → (re10k+dl3dv)" means the model is firstly trained on the RealEstate10K dataset and then fine-tuned using a combination of the RealEstate10K and DL3DV datasets.
+
+
+| Model                                                        |       Training Data        |  Training Resolution  | Training Views | Params (M) |                           Download                           |
+| ------------------------------------------------------------ | :------------------------: | :-------------------: | :------------: | :--------: | :----------------------------------------------------------: |
+| depthsplat-gs-small-re10k-256x256-view2-cfeab6b1.pth         |           re10k            |        256x256        |       2        |     37     | [download](https://huggingface.co/haofeixu/depthsplat/resolve/main/depthsplat-gs-small-re10k-256x256-view2-cfeab6b1.pth) |
+| depthsplat-gs-base-re10k-256x256-view2-ca7b6795.pth          |           re10k            |        256x256        |       2        |    117     | [download](https://huggingface.co/haofeixu/depthsplat/resolve/main/depthsplat-gs-base-re10k-256x256-view2-ca7b6795.pth) |
+| depthsplat-gs-large-re10k-256x256-view2-e0f0f27a.pth         |           re10k            |        256x256        |       2        |    360     | [download](https://huggingface.co/haofeixu/depthsplat/resolve/main/depthsplat-gs-large-re10k-256x256-view2-e0f0f27a.pth) |
+| depthsplat-gs-base-re10k-256x448-view2-fea94f65.pth          |           re10k            |        256x448        |       2        |    117     | [download](https://huggingface.co/haofeixu/depthsplat/resolve/main/depthsplat-gs-base-re10k-256x448-view2-fea94f65.pth) |
+| depthsplat-gs-base-dl3dv-256x448-randview2-6-02c7b19d.pth    |     re10k → dl3dv     |        256x448        |      2-6       |    117     | [download](https://huggingface.co/haofeixu/depthsplat/resolve/main/depthsplat-gs-base-dl3dv-256x448-randview2-6-02c7b19d.pth) |
+| depthsplat-gs-small-re10kdl3dv-448x768-randview4-10-c08188db.pth | re10k → (re10k+dl3dv) | 256x448 →448x768 |      4-10      |     37     | [download](https://huggingface.co/haofeixu/depthsplat/resolve/main/depthsplat-gs-small-re10kdl3dv-448x768-randview4-10-c08188db.pth) |
+| depthsplat-gs-base-re10kdl3dv-448x768-randview2-6-f8ddd845.pth | re10k → (re10k+dl3dv) | 256x448 →448x768 |      2-6       |    117     | [download](https://huggingface.co/haofeixu/depthsplat/resolve/main/depthsplat-gs-base-re10kdl3dv-448x768-randview2-6-f8ddd845.pth) |
+
+
+
+## Depth Prediction
+
+- The depth models are trained with the following procedure:
+  - Initialize the monocular feature with Depth Anything V2 and the multi-view Transformer with UniMatch.
+  - Train the full DepthSplat model end-to-end on the mixed RealEstate10K and DL3DV datasets.
+  - Fine-tune the pre-trained depth model on the depth datasets with ground truth depth supervision. The depth datasets used for fine-tuning include ScanNet, TartanAir, and VKITTI2.
+- The depth models are fine-tuned with random numbers (2-8) of input images, and the training image resolution is 352x640.
+- The scale of the predicted depth is aligned with the scale of camera pose's translation.
+
+| Model                                                   |                  Training Data                   |  Training Resolution   | Training Views | Params (M) |                           Download                           |
+| ------------------------------------------------------- | :----------------------------------------------: | :--------------------: | :------------: | :--------: | :----------------------------------------------------------: |
+| depthsplat-depth-small-352x640-randview2-8-e807bd82.pth | (re10k+dl3dv) → (scannet+tartanair+vkitti2) | 448x768 → 352x640 |      2-8       |     36     | [download](https://huggingface.co/haofeixu/depthsplat/resolve/main/depthsplat-depth-small-352x640-randview2-8-e807bd82.pth) |
+| depthsplat-depth-base-352x640-randview2-8-65a892c5.pth  | (re10k+dl3dv) → (scannet+tartanair+vkitti2) | 448x768 → 352x640 |      2-8       |    111     | [download](https://huggingface.co/haofeixu/depthsplat/resolve/main/depthsplat-depth-base-352x640-randview2-8-65a892c5.pth) |
+
+

MinkowskiEngine/.gitignore

@@ -0,0 +1,11 @@
+*.so
+*.[o,d,a]
+*.swo
+*.swp
+*.swn
+*.pyc
+
+build/
+objs/
+dist/
+MinkowskiEngine.egg-info/

MinkowskiEngine/CHANGELOG.md

@@ -0,0 +1,319 @@
+# Change Log
+
+## [0.5.5]
+
+- MKL compilation fix (#358)
+- Example updates for `has_bias` (PR #356)
+- Add thrust exception handling (issue #357)
+- `min_coordinate` to device in `dense()`
+- `cpu_kernel_region` `to_gpu` now uses `shared_ptr` allocator for gpu memory
+- Docker installation instruction added
+- Fix for GPU `coo_spmm` when nnz == 0
+    - Fix MinkowskiInterpolationGPU for invalid samples (issue #383)
+- gradcheck wrap func 1.9
+- Handles `dense` for an empty tensor (issue #384)
+- Handles an emtpy tensor for convolution (issue #384)
+- When `coordinates` is provided in `MinkowskiToSparseTensor`, `remove_zeros` will be ignored (issue #387)
+- Pybind equality error from the package (issue #414)
+- Fix undefined coordinate merge for multiple coordinate unions
+- Add cross-shaped kernel support (issue #436)
+
+## [0.5.4]
+
+- Fix `TensorField.sparse()` for no duplicate coordinates
+- Skip unnecessary spmm if `SparseTensor.initialize_coordinates()` has no duplicate coordinates
+- Model summary utility function added
+- TensorField.splat function for splat features to a sparse tensor
+- SparseTensor.interpolate function for extracting interpolated features
+- `coordinate_key` property function for `SparseTensor` and `TensorField`
+- Fix .dense() for GPU tensors. (PR #319)
+
+## [0.5.3]
+
+- Updated README for pytorch 1.8.1 support
+- Use custom `gpu_storage` instead of thrust vector for faster constructors
+- pytorch installation instruction updates
+- fix transpose kernel map with `kernel_size == stride_size`
+- Update reconstruction and vae examples for v0.5 API
+- `stack_unet.py` example, API updates
+- `MinkowskiToFeature` layer
+
+## [0.5.2]
+
+- spmm average cuda function
+- SparseTensor list operators (cat, mean, sum, var)
+- MinkowskiStack containers
+- Replace all at::cuda::getCurrentCUDASparseHandle with custom getCurrentCUDASparseHandle (issue #308)
+- fix coordinate manager kernel map python function
+- direct max pool
+    - SparseTensorQuantizationMode.MAX_POOL
+- TensorField global max pool
+    - origin field
+    - origin field map
+    - MinkowskiGlobalMaxPool CPU/GPU updates for a field input
+- SparseTensor.dense() raises a value error when a coordinate is negative rather than subtracting the minimum coordinate from a sparse tensor. (issue #316)
+- Added `to_sparse()` that removes zeros. (issue #317)
+    - Previous `to_sparse()` was renamed to `to_sparse_all()`
+    - `MinkowskiToSparseTensor` takes an optional `remove_zeros` boolean argument.
+- Fix global max pool with batch size 1
+- Use separate memory chunks for in, out map, and kernel indices for `gpu_kernel_map` for gpu memory misaligned error
+
+
+## [0.5.1]
+
+- v0.5 documentation updates
+- Nonlinear functionals and modules
+- Warning when using cuda without ME cuda support
+- diagnostics test
+- TensorField slice
+    - Cache the unique map and inverse map pair in the coordinate manager
+    - generate inverse_mapping on the fly
+- CoordinateManager
+    - `field_to_sparse_insert_and_map`
+    - `exists_field_to_sparse`
+    - `get_field_to_sparse_map`
+    - fix `kernel_map` with empty coordinate maps
+- CoordiateFieldMap
+    - `quantize_coordinates`
+- TensorField binary ops fix
+- MinkowskiSyncBatchNorm
+    - Support tfield
+    - conver sync batchnorm updates
+- TensorField to sparse with coordinate map key
+- Sparse matrix multiplication
+    - force contiguous matrix
+- Fix AveragePooling cudaErrorMisalignedAddress error for CUDA 10 (#246)
+
+## [0.5.0] - 2020-12-24
+
+## [0.5.0a] - 2020-08-05
+
+### Changed
+
+- Remove Makefile for installation as pytorch supports multithreaded compilation
+- GPU coordinate map support
+- Coordinate union
+- Sparse tensor binary operators
+- CUDA 11.1 support
+- quantization function updates
+- Multi GPU examples
+- Pytorch-lightning multi-gpu example
+- Transpose pooling layers
+- TensorField updates
+- Batch-wise decomposition
+- inverse_map when sparse() called (slice)
+- ChannelwiseConvolution
+- TensorField support for non-linearities
+
+## [0.4.3] - 2020-05-29
+
+### Changed
+
+- Use `CPU_ONLY` compile when `torch` fails to detect a GPU (Issue #105)
+- Fix `get_kernel_map` for `CPU_ONLY` (Issue #107)
+- Update `get_union_map` doc (Issue #108)
+- Abstract getattr minkowski backend functions
+- Add `coordinates_and_features_at(batch_index)` function in the SparseTensor class.
+- Add `MinkowskiChannelwiseConvolution` (Issue #92)
+- Update `MinkowskiPruning` to generate an empty sparse tensor as output (Issue #102)
+- Add `return_index` for `sparse_quantize`
+- Templated CoordsManager for coords to int and coords to vector classes
+- Sparse tensor quantization mode
+    - Features at duplicated coordinates will be averaged automatically with `quantization_mode=ME.SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE`
+- `SparseTensor.slice()` slicing features on discrete coordinates to continuous coordinates
+- CoordsManager.getKernelMapGPU returns long type tensors (Issue #125)
+- SyncBatchNorm error fix (Issue #129)
+- Sparse Tensor `dense()` doc update (Issue #126)
+- Installation arguments `--cuda_home=<value>`, `--force_cuda`, `--blas_include_dirs=<comma_separated_values>`, and '--blas_library_dirs=<comma_separated_values>`. (Issue #135)
+- SparseTensor query by coordinates `features_at_coords` (Issue #137)
+- Memory manager control. CUDA | Pytorch memory manager for cuda malloc
+
+
+## [0.4.2] - 2020-03-13
+
+### Added
+
+- Completion and VAE examples
+- GPU version of getKernelMap: getKernelMapGPU
+
+### Changed
+
+- Fix dtype double to float on the multi-gpu example
+- Remove the dimension input argument on GlobalPooling, Broadcast functions
+- Kernel map generation has tensor stride > 0 check
+- Fix `SparseTensor.set_tensor_stride`
+- Track whether the batch indices are set first when initializing coords, The initial batch indices will be used throughout the lifetime of a sparse tensor
+- Add a memory warning on ModelNet40 training example (Issue #86)
+- Update the readme, definition
+- Fix an error in examples.convolution
+- Changed `features_at`, `coordinates_at` to take a batch index not the index of the unique batch indices. (Issue #100)
+- Fix an error torch.range --> torch.arange in `sparse_quantize` (Issue #101)
+- Fix BLAS installation link error (Issue #94)
+- Fix `MinkowskiBroadcast` and `MinkowskiBroadcastConcatenation` to use arbitrary channel sizes
+- Fix `pointnet.py` example (Issue #103)
+
+
+## [0.4.1] - 2020-01-28
+
+### Changed
+
+- Kernel maps with region size 1 do not require `Region` class initialization.
+- Faster union map with out map initialization
+- Batch index order hot fix on `dense()`, `sparse()`
+
+
+## [0.4.0] - 2020-01-26
+
+### Added
+
+- Add `MinkowskiGlobalSumPooling`, `MinkowskiGlobalAvgPooling`
+- Add `examples/convolution.py` to showcase various usages
+- Add `examples/sparse_tensor_basic.py` and a SparseTensor tutorial page
+- Add convolution, kernel map gifs
+- Add batch decomposition functions
+    - Add `SparseTensor.decomposed_coordinates`
+    - Add `SparseTensor.decomposed_features`
+    - Add `SparseTensor.coordinates_at(batch_index)`
+    - Add `SparseTensor.features_at(batch_index)`
+    - Add `CoordsManager.get_row_indices_at(coords_key, batch_index)`
+
+
+### Changed
+
+- `SparseTensor` additional coords.device guard
+- `MinkowskiConvolution`, `Minkowski*Pooling` output coordinates will be equal to the input coordinates if stride == 1. Before this change, they generated output coordinates previously defined for a specific tensor stride.
+- `MinkowskiUnion` and `Ops.cat` will take a variable number of sparse tensors not a list of sparse tensors
+- Namespace cleanup
+- Fix global in out map with uninitialized global map
+- `getKernelMap` now can generate new kernel map if it doesn't exist
+- `MinkowskiPruning` initialization takes no argument
+- Batched coordinates with batch indices prepended before coordinates
+
+
+## [0.3.3] - 2020-01-07
+
+### Added
+
+- Add `get_coords_map` on `CoordsManager`.
+- Add `get_coords_map` on `MinkowskiEngine.utils`.
+- Sparse Tensor Sparse Tensor binary operations `(+,-,*,/)`
+    - Binary operations between sparse tensors or sparse tensor + pytorch tensor
+    - Inplace operations for the same coords key
+- Sparse Tensor operation mode
+    - Add `set_sparse_tensor_operation_mode` sharing the global coords manager by default
+
+### Changed
+
+- Minor changes on `setup.py` for torch installation check and system assertions.
+- Update BLAS installation configuration.
+- Update union kernel map and union coords to use reference wrappers.
+- namespace `minkowski` for all cpp, cu files
+- `MinkowskiConvolution` and `MinkowskiConvolutionTranspose` now support output coordinate specification on the function call.
+- `Minkowski[Avg|Max|Sum]Pooling` and `Minkowski[Avg|Max|Sum]PoolingTranspose` now support output coordinate specification on the function call.
+
+
+## [0.3.2] - 2019-12-25
+
+### Added
+- Synchronized Batch Norm: `ME.MinkowskiSyncBatchNorm`
+    - `ME.MinkowskiSyncBatchNorm.convert_sync_batchnorm` converts a MinkowskiNetwork automatically to use synched batch norm.
+- `examples/multigpu.py` update for `ME.MinkowskiSynchBatchNorm`.
+- Add `MinkowskiUnion`
+- Add CoordsManager functions
+    - `get_batch_size`
+    - `get_batch_indices`
+    - `set_origin_coords_key`
+- Add `quantize_th`, `quantize_label_th`
+- Add MANIFEST
+
+### Changed
+
+- Update `MinkowskiUnion`, `MinkowskiPruning` docs
+- Update multigpu documentation
+- Update GIL release
+- Use cudaMalloc instead of `at::Tensor` for GPU memory management for illegal memory access, invalid arg.
+- Minor error fixes on `examples/modelnet40.py`
+- CoordsMap size initialization updates
+- Region hypercube iterator with even numbered kernel
+- Fix global reduction in-out map with non contiguous batch indices
+- GlobalPooling with torch reduction
+- Update CoordsManager function `get_row_indices_per_batch` to return a list of `torch.LongTensor` for mapping indices. The corresponding batch indices is accessible by `get_batch_indices`.
+- Update `MinkowskiBroadcast`, `MinkowskiBroadcastConcatenation` to use row indices per batch (`getRowIndicesPerBatch`)
+- Update `SparseTensor`
+    - `allow_duplicate_coords` argument support
+    - update documentation, add unittest
+- Update the training demo and documentation.
+- Update `MinkowskiInstanceNorm`: no `dimension` argument.
+- Fix CPU only build
+
+
+## [0.3.1] - 2019-12-15
+
+- Cache in-out mapping on device
+- Robinhood unordered map for coordinate management
+- hash based quantization to C++ CoordsManager based quantization with label collision
+- CUDA compilation to support older devices (compute_30, 35)
+- OMP_NUM_THREADS to initialize the number of threads
+
+
+## [0.3.0] - 2019-12-08
+
+- Change the hash map from google-sparsehash to Threading Building Blocks (TBB) `concurrent_unordered_map`.
+    - Optional duplicate coords (CoordsMap.initialize, TODO: make mapping more user-friendly)
+    - Explicit coords generation (`CoordsMap.stride`, `CoordsMap.reduce`, `CoordsMap.transposed_stride`)
+    - Speed up for pooling with `kernel_size == stride_size`.
+- Faster `SparseTensor.dense` function.
+- Force scratch memory space to be contiguous.
+- CUDA error checks
+- Update Makefile
+    - Architecture and sm updates for CUDA > 10.0
+    - Optional cblas
+
+
+## [0.2.9] - 2019-11-17
+
+- Pytorch 1.3 support
+    - Update torch cublas, cusparse handles.
+- Global max pooling layers.
+- Minor error fix in the coordinate manager
+    - Fix cases to return `in_coords_key` when stride is identity.
+
+
+## [0.2.8] - 2019-10-18
+
+- ModelNet40 training.
+- open3d v0.8 update.
+- Dynamic coordinate generation.
+
+
+## [0.2.7] - 2019-09-04
+
+Use `std::vector` for all internal coordinates to support arbitrary spatial dimensions.
+
+- Vectorized coordinates to support arbitrary spatial dimensions.
+- Removed all dimension template instantiation.
+- Use assertion macro for cleaner exception handling.
+
+
+## [0.2.6] - 2019-08-28
+
+Use OpenMP for multi-threaded kernel map generation and minor renaming and explicit coordinate management for future upgrades.
+
+- Major speed up
+    - Suboptimal kernels were introduced, and optimal kernels removed for faulty cleanup in v0.2.5. CUDA kernels were re-introduced and major speed up was restored.
+- Minor name changes in `CoordsManager`.
+- `CoordsManager` saves all coordinates for future updates.
+- `CoordsManager` functions `createInOutPerKernel` and `createInOutPerKernelTranspose` now support multi-threaded kernel map generation by default using OpenMP.
+    - Thus, all manual thread functions such as `createInOutPerKernelInThreads`, `initialize_nthreads` removed.
+        - Use `export OMP_NUM_THREADS` to control the number of threads.
+
+
+## [0.2.5a0] - 2019-07-12
+
+- Added the `MinkowskiBroadcast` and `MinkowskiBroadcastConcatenation` module.
+
+
+## [0.2.5] - 2019-07-02
+
+- Better GPU memory management:
+    - GPU Memory management is now delegated to pytorch. Before the change, we need to cleanup the GPU cache that pytorch created to call `cudaMalloc`, which not only is slow but also hampers the long-running training that dies due to Out Of Memory (OOM).

MinkowskiEngine/LICENSE

@@ -0,0 +1,26 @@
+MIT License
+
+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.

MinkowskiEngine/MANIFEST.in

@@ -0,0 +1,3 @@
+include Makefile
+recursive-include src *.hpp *.cpp *.cu *.cuh *.h
+recursive-include pybind *.hpp *.cpp

MinkowskiEngine/Makefile

@@ -0,0 +1,178 @@
+###############################################################################
+# Uncomment for debugging
+# DEBUG := 1
+# Pretty build
+Q ?= @
+
+# Uncomment for CPU only build. From the command line, `python setup.py install --cpu_only`
+# CPU_ONLY := 1
+
+CXX ?= g++
+PYTHON ?= python
+
+EXTENSION_NAME := minkowski
+
+# BLAS choice:
+# atlas for ATLAS
+# blas for default blas
+# mkl for MKL. For conda, conda install -c intel mkl mkl-include
+# openblas for OpenBlas (default)
+BLAS ?= openblas
+CUDA_HOME ?= $(shell $(PYTHON) -c 'from torch.utils.cpp_extension import _find_cuda_home; print(_find_cuda_home())')
+
+# Custom (MKL/ATLAS/OpenBLAS) include and lib directories.
+# Leave commented to accept the defaults for your choice of BLAS
+# (which should work)!
+# BLAS_INCLUDE_DIRS ?=
+# BLAS_LIBRARY_DIRS ?=
+
+###############################################################################
+# PYTHON Header path
+PYTHON_HEADER_DIR := $(shell $(PYTHON) -c 'from distutils.sysconfig import get_python_inc; print(get_python_inc())')
+PYTORCH_INCLUDES := $(shell $(PYTHON) -c 'from torch.utils.cpp_extension import include_paths; [print(p) for p in include_paths()]')
+PYTORCH_LIBRARIES := $(shell $(PYTHON) -c 'from torch.utils.cpp_extension import library_paths; [print(p) for p in library_paths()]')
+
+# HEADER DIR is in pythonX.Xm folder
+INCLUDE_DIRS := $(PYTHON_HEADER_DIR)
+INCLUDE_DIRS += $(PYTHON_HEADER_DIR)/..
+INCLUDE_DIRS += $(PYTORCH_INCLUDES)
+LIBRARY_DIRS := $(PYTORCH_LIBRARIES)
+
+# Determine ABI support
+WITH_ABI := $(shell $(PYTHON) -c 'import torch; print(int(torch._C._GLIBCXX_USE_CXX11_ABI))')
+
+# Determine platform
+UNAME := $(shell uname -s)
+ifeq ($(UNAME), Linux)
+	LINUX := 1
+else ifeq ($(UNAME), Darwin)
+	OSX := 1
+	OSX_MAJOR_VERSION := $(shell sw_vers -productVersion | cut -f 1 -d .)
+	OSX_MINOR_VERSION := $(shell sw_vers -productVersion | cut -f 2 -d .)
+
+	CXX := /usr/local/opt/llvm/bin/clang
+	# brew install llvm libomp
+	INCLUDE_DIRS += /usr/local/opt/llvm/include
+	LIBRARY_DIRS += /usr/local/opt/llvm/lib
+endif
+
+ifneq ($(CPU_ONLY), 1)
+	# CUDA ROOT DIR that contains bin/ lib64/ and include/
+	# CUDA_HOME := /usr/local/cuda
+	
+	NVCC ?= $(CUDA_HOME)/bin/nvcc
+	INCLUDE_DIRS += ./ $(CUDA_HOME)/include
+	LIBRARY_DIRS += $(CUDA_HOME)/lib64
+endif
+
+SRC_DIR := ./src
+OBJ_DIR := ./objs
+CPP_SRCS := $(wildcard $(SRC_DIR)/*.cpp)
+CU_SRCS := $(wildcard $(SRC_DIR)/*.cu)
+OBJS := $(patsubst $(SRC_DIR)/%.cpp,$(OBJ_DIR)/%.o,$(CPP_SRCS))
+CU_OBJS := $(patsubst $(SRC_DIR)/%.cu,$(OBJ_DIR)/cuda/%.o,$(CU_SRCS))
+STATIC_LIB := $(OBJ_DIR)/lib$(EXTENSION_NAME).a
+
+# We will also explicitly add stdc++ to the link target.
+LIBRARIES := stdc++ c10 caffe2 torch torch_python _C
+ifneq ($(CPU_ONLY), 1)
+	LIBRARIES += cudart cublas cusparse caffe2_gpu c10_cuda
+	CUDA_ARCH := -gencode arch=compute_30,code=sm_30 \
+			-gencode arch=compute_35,code=sm_35 \
+			-gencode=arch=compute_50,code=sm_50 \
+			-gencode=arch=compute_52,code=sm_52 \
+			-gencode=arch=compute_60,code=sm_60 \
+			-gencode=arch=compute_61,code=sm_61 \
+			-gencode=arch=compute_70,code=sm_70 \
+			-gencode=arch=compute_75,code=sm_75 \
+			-gencode=arch=compute_75,code=compute_75
+endif
+
+# BLAS configuration: mkl, atlas, open, blas
+BLAS ?= openblas
+ifeq ($(BLAS), mkl)
+	# MKL
+	LIBRARIES += mkl_rt
+	COMMON_FLAGS += -DUSE_MKL
+	MKLROOT ?= /opt/intel/mkl
+	BLAS_INCLUDE_DIRS ?= $(MKLROOT)/include
+	BLAS_LIBRARY_DIRS ?= $(MKLROOT)/lib $(MKLROOT)/lib/intel64
+else ifeq ($(BLAS), openblas)
+	# OpenBLAS
+	LIBRARIES += openblas
+else ifeq ($(BLAS), blas)
+	# OpenBLAS
+	LIBRARIES += blas
+else
+	# ATLAS
+	LIBRARIES += atlas
+	ATLAS_PATH := $(shell $(PYTHON) -c "import numpy.distutils.system_info as si; ai = si.atlas_info(); [print(p) for p in ai.get_lib_dirs()]")
+	BLAS_LIBRARY_DIRS += $(ATLAS_PATH)
+endif
+
+INCLUDE_DIRS += ./src/3rdparty
+INCLUDE_DIRS += $(BLAS_INCLUDE_DIRS)
+LIBRARY_DIRS += $(BLAS_LIBRARY_DIRS)
+
+# Debugging
+ifeq ($(DEBUG), 1)
+	COMMON_FLAGS += -DDEBUG -g -O0
+	# https://gcoe-dresden.de/reaching-the-shore-with-a-fog-warning-my-eurohack-day-4-morning-session/
+	NVCCFLAGS := -g -G # -rdc true
+else
+	COMMON_FLAGS += -DNDEBUG -O3
+endif
+
+WARNINGS := -Wall -Wcomment -Wno-sign-compare -Wno-deprecated-declarations
+
+# Automatic dependency generation (nvcc is handled separately)
+CXXFLAGS += -MMD -MP
+
+# Fast math
+CXXFLAGS += -ffast-math -funsafe-math-optimizations -fno-math-errno
+
+# BATCH FIRST
+CXXFLAGS += -DBATCH_FIRST=1
+
+# Complete build flags.
+COMMON_FLAGS += $(foreach includedir,$(INCLUDE_DIRS),-I$(includedir)) \
+	     -DTORCH_API_INCLUDE_EXTENSION_H -DTORCH_EXTENSION_NAME=$(EXTENSION_NAME) \
+	     -D_GLIBCXX_USE_CXX11_ABI=$(WITH_ABI)
+
+CXXFLAGS += -fopenmp -fPIC -fwrapv -std=c++14 $(COMMON_FLAGS) $(WARNINGS)
+NVCCFLAGS += -std=c++14 -ccbin=$(CXX) -Xcompiler -fPIC $(COMMON_FLAGS)
+LINKFLAGS += -pthread -fPIC $(WARNINGS) -Wl,-rpath=$(PYTHON_LIB_DIR) -Wl,--no-as-needed -Wl,--sysroot=/
+LDFLAGS += $(foreach librarydir,$(LIBRARY_DIRS),-L$(librarydir)) \
+	   $(foreach library,$(LIBRARIES),-l$(library))
+
+ifeq ($(CPU_ONLY), 1)
+	ALL_OBJS := $(OBJS)
+	CXXFLAGS += -DCPU_ONLY
+else
+	ALL_OBJS := $(OBJS) $(CU_OBJS)
+endif
+
+all: $(STATIC_LIB)
+	$(RM) -rf build dist
+
+$(OBJ_DIR):
+	@ mkdir -p $@
+	@ mkdir -p $@/cuda
+
+$(OBJ_DIR)/%.o: $(SRC_DIR)/%.cpp | $(OBJ_DIR)
+	@ echo CXX $<
+	$(Q)$(CXX) $< $(CXXFLAGS) -c -o $@
+
+$(OBJ_DIR)/cuda/%.o: $(SRC_DIR)/%.cu | $(OBJ_DIR)
+	@ echo NVCC $<
+	$(Q)$(NVCC) $(NVCCFLAGS) $(CUDA_ARCH) -M $< -o ${@:.o=.d} \
+		-odir $(@D)
+	$(Q)$(NVCC) $(NVCCFLAGS) $(CUDA_ARCH) -c $< -o $@
+
+$(STATIC_LIB): $(ALL_OBJS) | $(OBJ_DIR)
+	$(RM) -f $(STATIC_LIB)
+	@ echo LD -o $@
+	ar rc $(STATIC_LIB) $(ALL_OBJS)
+
+clean:
+	@- $(RM) -rf $(OBJ_DIR) build dist

MinkowskiEngine/MinkowskiEngine/MinkowskiBroadcast.py

@@ -0,0 +1,253 @@
+# 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.
+from typing import Union
+
+import torch
+from torch.nn import Module
+from torch.autograd import Function
+
+from MinkowskiEngineBackend._C import CoordinateMapKey, RegionType, BroadcastMode
+from MinkowskiSparseTensor import SparseTensor, _get_coordinate_map_key
+from MinkowskiCoordinateManager import CoordinateManager
+from MinkowskiCommon import (
+    MinkowskiModuleBase,
+    get_minkowski_function,
+)
+
+
+class MinkowskiBroadcastFunction(Function):
+    @staticmethod
+    def forward(
+        ctx,
+        input_features: torch.Tensor,
+        input_features_global: torch.Tensor,
+        operation_type: BroadcastMode,
+        in_coords_key: CoordinateMapKey,
+        glob_coords_key: CoordinateMapKey,
+        coords_manager: CoordinateManager,
+    ):
+        assert isinstance(operation_type, BroadcastMode)
+
+        ctx.saved_vars = (
+            input_features,
+            input_features_global,
+            operation_type,
+            in_coords_key,
+            glob_coords_key,
+            coords_manager,
+        )
+
+        fw_fn = get_minkowski_function("BroadcastForward", input_features)
+        return fw_fn(
+            input_features,
+            input_features_global,
+            operation_type,
+            in_coords_key,
+            glob_coords_key,
+            coords_manager._manager,
+        )
+
+    @staticmethod
+    def backward(ctx, grad_out_feat):
+        if not grad_out_feat.is_contiguous():
+            grad_out_feat = grad_out_feat.contiguous()
+
+        (
+            input_features,
+            input_features_global,
+            operation_type,
+            in_coords_key,
+            glob_coords_key,
+            coords_manager,
+        ) = ctx.saved_vars
+
+        bw_fn = get_minkowski_function("BroadcastBackward", grad_out_feat)
+        grad_in_feat, grad_in_feat_glob = bw_fn(
+            input_features,
+            input_features_global,
+            grad_out_feat,
+            operation_type,
+            in_coords_key,
+            glob_coords_key,
+            coords_manager._manager,
+        )
+        return grad_in_feat, grad_in_feat_glob, None, None, None, None
+
+
+class MinkowskiBroadcastBase(MinkowskiModuleBase):
+    def __init__(self, operation_type):
+        MinkowskiModuleBase.__init__(self)
+        assert isinstance(operation_type, BroadcastMode)
+
+        self.operation_type = operation_type
+
+        self.broadcast = MinkowskiBroadcastFunction()
+
+    def forward(self, input: SparseTensor, input_glob: SparseTensor):
+        assert isinstance(input, SparseTensor)
+
+        output = self.broadcast.apply(
+            input.F,
+            input_glob.F,
+            self.operation_type,
+            input.coordinate_map_key,
+            input_glob.coordinate_map_key,
+            input.coordinate_manager,
+        )
+        return SparseTensor(
+            output,
+            coordinate_map_key=input.coordinate_map_key,
+            coordinate_manager=input.coordinate_manager,
+        )
+
+    def __repr__(self):
+        return self.__class__.__name__
+
+
+class MinkowskiBroadcastAddition(MinkowskiBroadcastBase):
+    r"""Broadcast the reduced features to all input coordinates.
+
+    .. math::
+
+        \mathbf{y}_\mathbf{u} = \mathbf{x}_{1, \mathbf{u}} + \mathbf{x}_2
+        \; \text{for} \; \mathbf{u} \in \mathcal{C}^\text{in}
+
+
+    For all input :math:`\mathbf{x}_\mathbf{u}`, add :math:`\mathbf{x}_2`. The
+    output coordinates will be the same as the input coordinates
+    :math:`\mathcal{C}^\text{in} = \mathcal{C}^\text{out}`.
+
+    .. note::
+        The first argument takes a sparse tensor; the second argument takes
+        features that are reduced to the origin. This can be typically done with
+        the global reduction such as the :attr:`MinkowskiGlobalPooling`.
+
+    """
+
+    def __init__(self):
+        MinkowskiBroadcastBase.__init__(self, BroadcastMode.ELEMENTWISE_ADDITON)
+
+
+class MinkowskiBroadcastMultiplication(MinkowskiBroadcastBase):
+    r"""Broadcast reduced features to all input coordinates.
+
+    .. math::
+
+        \mathbf{y}_\mathbf{u} = \mathbf{x}_{1, \mathbf{u}} \times \mathbf{x}_2
+        \; \text{for} \; \mathbf{u} \in \mathcal{C}^\text{in}
+
+
+    For all input :math:`\mathbf{x}_\mathbf{u}`, multiply :math:`\mathbf{x}_2`
+    element-wise. The output coordinates will be the same as the input
+    coordinates :math:`\mathcal{C}^\text{in} = \mathcal{C}^\text{out}`.
+
+    .. note::
+        The first argument takes a sparse tensor; the second argument takes
+        features that are reduced to the origin. This can be typically done with
+        the global reduction such as the :attr:`MinkowskiGlobalPooling`.
+
+    """
+
+    def __init__(self):
+        MinkowskiBroadcastBase.__init__(self, BroadcastMode.ELEMENTWISE_MULTIPLICATION)
+
+
+class MinkowskiBroadcast(Module):
+    r"""Broadcast reduced features to all input coordinates.
+
+    .. math::
+
+        \mathbf{y}_\mathbf{u} = \mathbf{x}_2 \; \text{for} \; \mathbf{u} \in
+        \mathcal{C}^\text{in}
+
+
+    For all input :math:`\mathbf{x}_\mathbf{u}`, copy value :math:`\mathbf{x}_2`
+    element-wise. The output coordinates will be the same as the input
+    coordinates :math:`\mathcal{C}^\text{in} = \mathcal{C}^\text{out}`. The
+    first input :math:`\mathbf{x}_1` is only used for defining the output
+    coordinates.
+
+    .. note::
+        The first argument takes a sparse tensor; the second argument takes
+        features that are reduced to the origin. This can be typically done with
+        the global reduction such as the :attr:`MinkowskiGlobalPooling`.
+
+    """
+
+    def __repr__(self):
+        return self.__class__.__name__
+
+    def forward(self, input: SparseTensor, input_glob: SparseTensor):
+        assert isinstance(input, SparseTensor)
+        assert isinstance(input_glob, SparseTensor)
+
+        broadcast_feat = input.F.new(len(input), input_glob.size()[1])
+        batch_indices, batch_rows = input.coordinate_manager.origin_map(input.coordinate_map_key)
+        for b, rows in zip(batch_indices, batch_rows):
+            broadcast_feat[rows] = input_glob.F[b]
+
+        return SparseTensor(
+            broadcast_feat,
+            coordinate_map_key=input.coordinate_map_key,
+            coordinate_manager=input.coordinate_manager,
+        )
+
+
+class MinkowskiBroadcastConcatenation(MinkowskiBroadcast):
+    r"""Broadcast reduced features to all input coordinates and concatenate to the input.
+
+    .. math::
+
+        \mathbf{y}_\mathbf{u} = [\mathbf{x}_{1,\mathbf{u}}, \mathbf{x}_2] \;
+        \text{for} \; \mathbf{u} \in \mathcal{C}^\text{in}
+
+
+    For all input :math:`\mathbf{x}_\mathbf{u}`, concatenate vector
+    :math:`\mathbf{x}_2`. :math:`[\cdot, \cdot]` is a concatenation operator.
+    The output coordinates will be the same as the input coordinates
+    :math:`\mathcal{C}^\text{in} = \mathcal{C}^\text{out}`.
+
+    .. note::
+        The first argument takes a sparse tensor; the second argument takes
+        features that are reduced to the origin. This can be typically done with
+        the global reduction such as the :attr:`MinkowskiGlobalPooling`.
+
+    """
+
+    def forward(self, input: SparseTensor, input_glob: SparseTensor):
+        assert isinstance(input, SparseTensor)
+        assert isinstance(input_glob, SparseTensor)
+
+        broadcast_feat = input.F.new(len(input), input_glob.size()[1])
+        batch_indices, batch_rows = input.coordinate_manager.origin_map(input.coordinate_map_key)
+        for b, row_ind in zip(batch_indices, batch_rows):
+            broadcast_feat[row_ind] = input_glob.F[b]
+
+        broadcast_cat = torch.cat((input.F, broadcast_feat), dim=1)
+        return SparseTensor(
+            broadcast_cat,
+            coordinate_map_key=input.coordinate_map_key,
+            coordinate_manager=input.coordinate_manager,
+        )

MinkowskiEngine/MinkowskiEngine/MinkowskiChannelwiseConvolution.py

@@ -0,0 +1,215 @@
+# 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.
+import math
+from typing import Union
+
+import torch
+from torch.nn import Parameter
+
+from MinkowskiSparseTensor import SparseTensor
+from MinkowskiEngineBackend._C import CoordinateMapKey, RegionType
+from MinkowskiCommon import MinkowskiModuleBase
+from MinkowskiKernelGenerator import KernelGenerator
+
+
+class MinkowskiChannelwiseConvolution(MinkowskiModuleBase):
+
+    __slots__ = (
+        "in_channels",
+        "out_channels",
+        "kernel_generator",
+        "dimension",
+        "kernel",
+        "bias",
+        "conv",
+    )
+
+    r"""Channelwise (Depthwise) Convolution layer for a sparse tensor.
+
+
+    .. math::
+
+        \mathbf{x}_\mathbf{u} = \sum_{\mathbf{i} \in \mathcal{N}^D(\mathbf{u}, K) \cap
+        \mathcal{C}^\text{in}} W_\mathbf{i} \odot \mathbf{x}_{\mathbf{i} +
+        \mathbf{u}} \;\text{for} \; \mathbf{u} \in \mathcal{C}^\text{out}
+
+    where :math:`K` is the kernel size and :math:`\mathcal{N}^D(\mathbf{u}, K)
+    \cap \mathcal{C}^\text{in}` is the set of offsets that are at most :math:`\left
+    \lceil{\frac{1}{2}(K - 1)} \right \rceil` away from :math:`\mathbf{u}`
+    defined in :math:`\mathcal{S}^\text{in}`. :math:`\odot` indicates the
+    elementwise product.
+
+    .. note::
+        For even :math:`K`, the kernel offset :math:`\mathcal{N}^D`
+        implementation is different from the above definition. The offsets
+        range from :math:`\mathbf{i} \in [0, K)^D, \; \mathbf{i} \in
+        \mathbb{Z}_+^D`.
+
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        kernel_size=-1,
+        stride=1,
+        dilation=1,
+        bias=False,
+        kernel_generator=None,
+        dimension=-1,
+    ):
+        r"""convolution on a sparse tensor
+
+        Args:
+            :attr:`in_channels` (int): the number of input channels in the
+            input tensor.
+
+            :attr:`kernel_size` (int, optional): the size of the kernel in the
+            output tensor. If not provided, :attr:`region_offset` should be
+            :attr:`RegionType.CUSTOM` and :attr:`region_offset` should be a 2D
+            matrix with size :math:`N\times D` such that it lists all :math:`N`
+            offsets in D-dimension.
+
+            :attr:`stride` (int, or list, optional): stride size of the
+            convolution layer. If non-identity is used, the output coordinates
+            will be at least :attr:`stride` :math:`\times` :attr:`tensor_stride`
+            away. When a list is given, the length must be D; each element will
+            be used for stride size for the specific axis.
+
+            :attr:`dilation` (int, or list, optional): dilation size for the
+            convolution kernel. When a list is given, the length must be D and
+            each element is an axis specific dilation. All elements must be > 0.
+
+            :attr:`bias` (bool, optional): if True, the convolution layer
+            has a bias.
+
+            :attr:`kernel_generator` (:attr:`MinkowskiEngine.KernelGenerator`,
+            optional): defines the custom kernel shape.
+
+            :attr:`dimension` (int): the spatial dimension of the space where
+            all the inputs and the network are defined. For example, images are
+            in a 2D space, meshes and 3D shapes are in a 3D space.
+
+        """
+
+        super(MinkowskiChannelwiseConvolution, self).__init__()
+        assert (
+            dimension > 0
+        ), f"Invalid dimension. Please provide a valid dimension argument. dimension={dimension}"
+
+        if kernel_generator is None:
+            kernel_generator = KernelGenerator(
+                kernel_size=kernel_size,
+                stride=stride,
+                dilation=dilation,
+                dimension=dimension,
+            )
+
+        self.kernel_generator = kernel_generator
+
+        self.in_channels = in_channels
+        self.dimension = dimension
+
+        self.kernel_shape = (kernel_generator.kernel_volume, self.in_channels)
+
+        Tensor = torch.FloatTensor
+        self.kernel = Parameter(Tensor(*self.kernel_shape))
+        self.bias = Parameter(Tensor(1, in_channels)) if bias else None
+
+        self.reset_parameters()
+
+    def forward(
+        self,
+        input: SparseTensor,
+        coords: Union[torch.IntTensor, CoordinateMapKey, SparseTensor] = None,
+    ):
+        r"""
+        :attr:`input` (`MinkowskiEngine.SparseTensor`): Input sparse tensor to apply a
+        convolution on.
+
+        :attr:`coords` ((`torch.IntTensor`, `MinkowskiEngine.CoordinateMapKey`,
+        `MinkowskiEngine.SparseTensor`), optional): If provided, generate
+        results on the provided coordinates. None by default.
+
+        """
+        assert isinstance(input, SparseTensor)
+        assert input.D == self.dimension
+        assert (
+            self.in_channels == input.shape[1]
+        ), f"Channel size mismatch {self.in_channels} != {input.shape[1]}"
+
+        # Create a region_offset
+        region_type_, region_offset_, _ = self.kernel_generator.get_kernel(
+            input.tensor_stride, False
+        )
+
+        cm = input.coordinate_manager
+        in_key = input.coordinate_map_key
+
+        out_key = cm.stride(in_key, self.kernel_generator.kernel_stride)
+        N_out = cm.size(out_key)
+        out_F = input._F.new(N_out, self.in_channels).zero_()
+
+        kernel_map = cm.kernel_map(
+            in_key,
+            out_key,
+            self.kernel_generator.kernel_stride,
+            self.kernel_generator.kernel_size,
+            self.kernel_generator.kernel_dilation,
+            region_type=region_type_,
+            region_offset=region_offset_,
+        )
+
+        for k, in_out in kernel_map.items():
+            in_out = in_out.long().to(input.device)
+            out_F[in_out[1]] += input.F[in_out[0]] * self.kernel[k]
+
+        if self.bias is not None:
+            out_F += self.bias
+
+        return SparseTensor(out_F, coordinate_map_key=out_key, coordinate_manager=cm)
+
+    def reset_parameters(self, is_transpose=False):
+        with torch.no_grad():
+            n = (
+                self.out_channels if is_transpose else self.in_channels
+            ) * self.kernel_generator.kernel_volume
+            stdv = 1.0 / math.sqrt(n)
+            self.kernel.data.uniform_(-stdv, stdv)
+            if self.bias is not None:
+                self.bias.data.uniform_(-stdv, stdv)
+
+    def __repr__(self):
+        s = "(in={}, region_type={}, ".format(
+            self.in_channels, self.kernel_generator.region_type
+        )
+        if self.kernel_generator.region_type in [RegionType.CUSTOM]:
+            s += "kernel_volume={}, ".format(self.kernel_generator.kernel_volume)
+        else:
+            s += "kernel_size={}, ".format(self.kernel_generator.kernel_size)
+        s += "stride={}, dilation={})".format(
+            self.kernel_generator.kernel_stride,
+            self.kernel_generator.kernel_dilation,
+        )
+        return self.__class__.__name__ + s

MinkowskiEngine/MinkowskiEngine/MinkowskiCommon.py

@@ -0,0 +1,120 @@
+# 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.
+from collections.abc import Sequence
+import numpy as np
+from typing import Union
+
+import torch
+
+from torch.nn import Module
+
+import MinkowskiEngineBackend._C as MEB
+
+
+StrideType = Union[int, Sequence, np.ndarray, torch.IntTensor]
+
+
+def convert_to_int_list(
+    arg: Union[int, Sequence, np.ndarray, torch.Tensor], dimension: int
+):
+    if isinstance(arg, list):
+        assert len(arg) == dimension
+        return arg
+
+    if isinstance(arg, (Sequence, np.ndarray, torch.Tensor)):
+        tmp = [i for i in arg]
+        assert len(tmp) == dimension
+    elif np.isscalar(arg):  # Assume that it is a scalar
+        tmp = [int(arg) for i in range(dimension)]
+    else:
+        raise ValueError("Input must be a scalar or a sequence")
+
+    return tmp
+
+
+def convert_to_int_tensor(
+    arg: Union[int, Sequence, np.ndarray, torch.IntTensor], dimension: int
+):
+    if isinstance(arg, torch.IntTensor):
+        assert arg.numel() == dimension
+        return arg
+
+    if isinstance(arg, (Sequence, np.ndarray)):
+        tmp = torch.IntTensor([i for i in arg])
+        assert tmp.numel() == dimension
+    elif np.isscalar(arg):  # Assume that it is a scalar
+        tmp = torch.IntTensor([int(arg) for i in range(dimension)])
+    else:
+        raise ValueError("Input must be a scalar or a sequence")
+
+    return tmp
+
+
+def prep_args(
+    tensor_stride: Union[int, Sequence, np.ndarray, torch.IntTensor],
+    stride: Union[int, Sequence, np.ndarray, torch.IntTensor],
+    kernel_size: Union[int, Sequence, np.ndarray, torch.IntTensor],
+    dilation: Union[int, Sequence, np.ndarray, torch.IntTensor],
+    region_type: Union[int, MEB.RegionType],
+    D=-1,
+):
+    assert torch.prod(
+        kernel_size > 0
+    ), f"kernel_size must be a positive integer, provided {kernel_size}"
+    assert D > 0, f"dimension must be a positive integer, {D}"
+    tensor_stride = convert_to_int_tensor(tensor_stride, D)
+    stride = convert_to_int_tensor(stride, D)
+    kernel_size = convert_to_int_tensor(kernel_size, D)
+    dilation = convert_to_int_tensor(dilation, D)
+    region_type = int(region_type)
+    return (
+        tensor_stride,
+        stride,
+        kernel_size,
+        dilation,
+        region_type,
+    )
+
+
+def get_postfix(tensor: torch.Tensor):
+    postfix = "GPU" if tensor.is_cuda else "CPU"
+    return postfix
+
+
+class MinkowskiModuleBase(Module):
+    pass
+
+
+def get_minkowski_function(name, variable):
+    fn_name = name + get_postfix(variable)
+    if hasattr(MEB, fn_name):
+        return getattr(MEB, fn_name)
+    else:
+        if variable.is_cuda:
+            raise ValueError(
+                f"Function {fn_name} not available. Please compile MinkowskiEngine with `torch.cuda.is_available()` is `True`."
+            )
+        else:
+            raise ValueError(f"Function {fn_name} not available.")

MinkowskiEngine/MinkowskiEngine/MinkowskiConvolution.py

@@ -0,0 +1,634 @@
+# 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.
+import math
+from typing import Union
+
+import torch
+from torch.autograd import Function
+from torch.nn import Parameter
+
+from MinkowskiEngineBackend._C import CoordinateMapKey, RegionType, ConvolutionMode
+from MinkowskiSparseTensor import SparseTensor, _get_coordinate_map_key
+from MinkowskiCommon import (
+    MinkowskiModuleBase,
+    get_minkowski_function,
+)
+from MinkowskiCoordinateManager import CoordinateManager
+from MinkowskiKernelGenerator import KernelGenerator
+
+
+class MinkowskiConvolutionFunction(Function):
+    @staticmethod
+    def forward(
+        ctx,
+        input_features: torch.Tensor,
+        kernel_weights: torch.Tensor,
+        kernel_generator: KernelGenerator,
+        convolution_mode: ConvolutionMode,
+        in_coordinate_map_key: CoordinateMapKey,
+        out_coordinate_map_key: CoordinateMapKey = None,
+        coordinate_manager: CoordinateManager = None,
+    ):
+        if out_coordinate_map_key is None:
+            out_coordinate_map_key = CoordinateMapKey(
+                in_coordinate_map_key.get_coordinate_size()
+            )
+
+        input_features = input_features.contiguous()
+
+        ctx.input_features = input_features
+        ctx.kernel_weights = kernel_weights
+        ctx.misc = [
+            kernel_generator,
+            convolution_mode,
+            in_coordinate_map_key,
+            out_coordinate_map_key,
+            coordinate_manager,
+        ]
+
+        fw_fn = get_minkowski_function("ConvolutionForward", input_features)
+        return fw_fn(
+            ctx.input_features,
+            kernel_weights,
+            kernel_generator.kernel_size,
+            kernel_generator.kernel_stride,
+            kernel_generator.kernel_dilation,
+            kernel_generator.region_type,
+            kernel_generator.region_offsets,
+            kernel_generator.expand_coordinates,
+            convolution_mode,
+            in_coordinate_map_key,
+            out_coordinate_map_key,
+            coordinate_manager._manager,
+        )
+
+    @staticmethod
+    def backward(ctx, grad_out_feat: torch.Tensor):
+        grad_out_feat = grad_out_feat.contiguous()
+        (
+            kernel_generator,
+            convolution_mode,
+            in_coordinate_map_key,
+            out_coordinate_map_key,
+            coordinate_manager,
+        ) = ctx.misc
+
+        bw_fn = get_minkowski_function("ConvolutionBackward", grad_out_feat)
+        grad_in_feat, grad_kernel = bw_fn(
+            ctx.input_features,
+            grad_out_feat,
+            ctx.kernel_weights,
+            kernel_generator.kernel_size,
+            kernel_generator.kernel_stride,
+            kernel_generator.kernel_dilation,
+            kernel_generator.region_type,
+            kernel_generator.region_offsets,
+            convolution_mode,
+            in_coordinate_map_key,
+            out_coordinate_map_key,
+            coordinate_manager._manager,
+        )
+        return (
+            grad_in_feat,
+            grad_kernel,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+class MinkowskiConvolutionTransposeFunction(Function):
+    @staticmethod
+    def forward(
+        ctx,
+        input_features: torch.Tensor,
+        kernel_weights: torch.Tensor,
+        kernel_generator: KernelGenerator,
+        convolution_mode: ConvolutionMode,
+        in_coordinate_map_key: CoordinateMapKey,
+        out_coordinate_map_key: CoordinateMapKey = None,
+        coordinate_manager: CoordinateManager = None,
+    ):
+        if out_coordinate_map_key is None:
+            out_coordinate_map_key = CoordinateMapKey(
+                in_coordinate_map_key.get_coordinate_size()
+            )
+        input_features = input_features.contiguous()
+        ctx.input_features = input_features
+        ctx.kernel_weights = kernel_weights
+        ctx.misc = (
+            kernel_generator,
+            convolution_mode,
+            in_coordinate_map_key,
+            out_coordinate_map_key,
+            coordinate_manager,
+        )
+
+        fw_fn = get_minkowski_function("ConvolutionTransposeForward", input_features)
+        return fw_fn(
+            ctx.input_features,
+            kernel_weights,
+            kernel_generator.kernel_size,
+            kernel_generator.kernel_stride,
+            kernel_generator.kernel_dilation,
+            kernel_generator.region_type,
+            kernel_generator.region_offsets,
+            kernel_generator.expand_coordinates,
+            convolution_mode,
+            in_coordinate_map_key,
+            out_coordinate_map_key,
+            coordinate_manager._manager,
+        )
+
+    @staticmethod
+    def backward(ctx, grad_out_feat: torch.Tensor):
+        grad_out_feat = grad_out_feat.contiguous()
+        (
+            kernel_generator,
+            convolution_mode,
+            in_coordinate_map_key,
+            out_coordinate_map_key,
+            coordinate_manager,
+        ) = ctx.misc
+
+        bw_fn = get_minkowski_function("ConvolutionTransposeBackward", grad_out_feat)
+        grad_in_feat, grad_kernel = bw_fn(
+            ctx.input_features,
+            grad_out_feat,
+            ctx.kernel_weights,
+            kernel_generator.kernel_size,
+            kernel_generator.kernel_stride,
+            kernel_generator.kernel_dilation,
+            kernel_generator.region_type,
+            kernel_generator.region_offsets,
+            convolution_mode,
+            in_coordinate_map_key,
+            out_coordinate_map_key,
+            coordinate_manager._manager,
+        )
+        return (
+            grad_in_feat,
+            grad_kernel,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+class MinkowskiConvolutionBase(MinkowskiModuleBase):
+
+    __slots__ = (
+        "in_channels",
+        "out_channels",
+        "is_transpose",
+        "kernel_generator",
+        "dimension",
+        "use_mm",
+        "kernel",
+        "bias",
+        "conv",
+    )
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=-1,
+        stride=1,
+        dilation=1,
+        bias=False,
+        kernel_generator=None,
+        is_transpose=False,  # only the base class has this argument
+        expand_coordinates=False,
+        convolution_mode=ConvolutionMode.DEFAULT,
+        dimension=-1,
+    ):
+        r"""
+
+        .. note::
+
+           When the kernel generator is provided, all kernel related arguments
+           (kernel_size, stride, dilation) will be ignored.
+
+        """
+        super(MinkowskiConvolutionBase, self).__init__()
+        assert (
+            dimension > 0
+        ), f"Invalid dimension. Please provide a valid dimension argument. dimension={dimension}"
+
+        if kernel_generator is None:
+            kernel_generator = KernelGenerator(
+                kernel_size=kernel_size,
+                stride=stride,
+                dilation=dilation,
+                expand_coordinates=expand_coordinates,
+                dimension=dimension,
+            )
+        else:
+            kernel_generator.expand_coordinates = expand_coordinates
+
+        self.is_transpose = is_transpose
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.kernel_generator = kernel_generator
+        self.dimension = dimension
+        self.use_mm = False  # use matrix multiplication when kernel_volume is 1
+
+        Tensor = torch.FloatTensor
+        if (
+            self.kernel_generator.kernel_volume == 1
+            and self.kernel_generator.requires_strided_coordinates
+        ):
+            kernel_shape = (self.in_channels, self.out_channels)
+            self.use_mm = True
+        else:
+            kernel_shape = (
+                self.kernel_generator.kernel_volume,
+                self.in_channels,
+                self.out_channels,
+            )
+
+        self.kernel = Parameter(Tensor(*kernel_shape))
+        self.bias = Parameter(Tensor(1, out_channels)) if bias else None
+        self.convolution_mode = convolution_mode
+        self.conv = (
+            MinkowskiConvolutionTransposeFunction()
+            if is_transpose
+            else MinkowskiConvolutionFunction()
+        )
+
+    def forward(
+        self,
+        input: SparseTensor,
+        coordinates: Union[torch.Tensor, CoordinateMapKey, SparseTensor] = None,
+    ):
+        r"""
+        :attr:`input` (`MinkowskiEngine.SparseTensor`): Input sparse tensor to apply a
+        convolution on.
+
+        :attr:`coordinates` ((`torch.IntTensor`, `MinkowskiEngine.CoordinateMapKey`,
+        `MinkowskiEngine.SparseTensor`), optional): If provided, generate
+        results on the provided coordinates. None by default.
+
+        """
+        assert isinstance(input, SparseTensor)
+        assert input.D == self.dimension
+
+        if self.use_mm:
+            # If the kernel_size == 1, the convolution is simply a matrix
+            # multiplication
+            out_coordinate_map_key = input.coordinate_map_key
+            outfeat = input.F.mm(self.kernel)
+        else:
+            # Get a new coordinate_map_key or extract one from the coords
+            out_coordinate_map_key = _get_coordinate_map_key(
+                input, coordinates, self.kernel_generator.expand_coordinates
+            )
+            outfeat = self.conv.apply(
+                input.F,
+                self.kernel,
+                self.kernel_generator,
+                self.convolution_mode,
+                input.coordinate_map_key,
+                out_coordinate_map_key,
+                input._manager,
+            )
+        if self.bias is not None:
+            outfeat += self.bias
+
+        return SparseTensor(
+            outfeat,
+            coordinate_map_key=out_coordinate_map_key,
+            coordinate_manager=input._manager,
+        )
+
+    def reset_parameters(self, is_transpose=False):
+        with torch.no_grad():
+            n = (
+                self.out_channels if is_transpose else self.in_channels
+            ) * self.kernel_generator.kernel_volume
+            stdv = 1.0 / math.sqrt(n)
+            self.kernel.data.uniform_(-stdv, stdv)
+            if self.bias is not None:
+                self.bias.data.uniform_(-stdv, stdv)
+
+    def __repr__(self):
+        s = "(in={}, out={}, ".format(
+            self.in_channels,
+            self.out_channels,
+        )
+        if self.kernel_generator.region_type in [RegionType.CUSTOM]:
+            s += "region_type={}, kernel_volume={}, ".format(
+                self.kernel_generator.region_type, self.kernel_generator.kernel_volume
+            )
+        else:
+            s += "kernel_size={}, ".format(self.kernel_generator.kernel_size)
+        s += "stride={}, dilation={})".format(
+            self.kernel_generator.kernel_stride,
+            self.kernel_generator.kernel_dilation,
+        )
+        return self.__class__.__name__ + s
+
+
+class MinkowskiConvolution(MinkowskiConvolutionBase):
+    r"""Convolution layer for a sparse tensor.
+
+
+    .. math::
+
+        \mathbf{x}_\mathbf{u} = \sum_{\mathbf{i} \in \mathcal{N}^D(\mathbf{u}, K,
+        \mathcal{C}^\text{in})} W_\mathbf{i} \mathbf{x}_{\mathbf{i} +
+        \mathbf{u}} \;\text{for} \; \mathbf{u} \in \mathcal{C}^\text{out}
+
+    where :math:`K` is the kernel size and :math:`\mathcal{N}^D(\mathbf{u}, K,
+    \mathcal{C}^\text{in})` is the set of offsets that are at most :math:`\left
+    \lceil{\frac{1}{2}(K - 1)} \right \rceil` away from :math:`\mathbf{u}`
+    definied in :math:`\mathcal{S}^\text{in}`.
+
+    .. note::
+        For even :math:`K`, the kernel offset :math:`\mathcal{N}^D`
+        implementation is different from the above definition. The offsets
+        range from :math:`\mathbf{i} \in [0, K)^D, \; \mathbf{i} \in
+        \mathbb{Z}_+^D`.
+
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=-1,
+        stride=1,
+        dilation=1,
+        bias=False,
+        kernel_generator=None,
+        expand_coordinates=False,
+        convolution_mode=ConvolutionMode.DEFAULT,
+        dimension=None,
+    ):
+        r"""convolution on a sparse tensor
+
+        Args:
+            :attr:`in_channels` (int): the number of input channels in the
+            input tensor.
+
+            :attr:`out_channels` (int): the number of output channels in the
+            output tensor.
+
+            :attr:`kernel_size` (int, optional): the size of the kernel in the
+            output tensor. If not provided, :attr:`region_offset` should be
+            :attr:`RegionType.CUSTOM` and :attr:`region_offset` should be a 2D
+            matrix with size :math:`N\times D` such that it lists all :math:`N`
+            offsets in D-dimension.
+
+            :attr:`stride` (int, or list, optional): stride size of the
+            convolution layer. If non-identity is used, the output coordinates
+            will be at least :attr:`stride` :math:`\times` :attr:`tensor_stride`
+            away. When a list is given, the length must be D; each element will
+            be used for stride size for the specific axis.
+
+            :attr:`dilation` (int, or list, optional): dilation size for the
+            convolution kernel. When a list is given, the length must be D and
+            each element is an axis specific dilation. All elements must be > 0.
+
+            :attr:`bias` (bool, optional): if True, the convolution layer
+            has a bias.
+
+            :attr:`kernel_generator` (:attr:`MinkowskiEngine.KernelGenerator`,
+            optional): defines custom kernel shape.
+
+            :attr:`expand_coordinates` (bool, optional): Force generation of
+            new coordinates. When True, the output coordinates will be the
+            outer product of the kernel shape and the input coordinates.
+            `False` by default.
+
+            :attr:`dimension` (int): the spatial dimension of the space where
+            all the inputs and the network are defined. For example, images are
+            in a 2D space, meshes and 3D shapes are in a 3D space.
+
+        """
+        MinkowskiConvolutionBase.__init__(
+            self,
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            dilation,
+            bias,
+            kernel_generator,
+            is_transpose=False,
+            expand_coordinates=expand_coordinates,
+            convolution_mode=convolution_mode,
+            dimension=dimension,
+        )
+        self.reset_parameters()
+
+
+class MinkowskiConvolutionTranspose(MinkowskiConvolutionBase):
+    r"""A generalized sparse transposed convolution or deconvolution layer."""
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=-1,
+        stride=1,
+        dilation=1,
+        bias=False,
+        kernel_generator=None,
+        expand_coordinates=False,
+        convolution_mode=ConvolutionMode.DEFAULT,
+        dimension=None,
+    ):
+        r"""a generalized sparse transposed convolution layer.
+
+        Args:
+            :attr:`in_channels` (int): the number of input channels in the
+            input tensor.
+
+            :attr:`out_channels` (int): the number of output channels in the
+            output tensor.
+
+            :attr:`kernel_size` (int, optional): the size of the kernel in the
+            output tensor. If not provided, :attr:`region_offset` should be
+            :attr:`RegionType.CUSTOM` and :attr:`region_offset` should be a 2D
+            matrix with size :math:`N\times D` such that it lists all :math:`N`
+            offsets in D-dimension.
+
+            :attr:`stride` (int, or list, optional): stride size that defines
+            upsampling rate. If non-identity is used, the output coordinates
+            will be :attr:`tensor_stride` / :attr:`stride` apart.  When a list is
+            given, the length must be D; each element will be used for stride
+            size for the specific axis.
+
+            :attr:`dilation` (int, or list, optional): dilation size for the
+            convolution kernel. When a list is given, the length must be D and
+            each element is an axis specific dilation. All elements must be > 0.
+
+            :attr:`bias` (bool, optional): if True, the convolution layer
+            has a bias.
+
+            :attr:`kernel_generator` (:attr:`MinkowskiEngine.KernelGenerator`,
+            optional): defines custom kernel shape.
+
+            :attr:`expand_coordinates` (bool, optional): Force generation of
+            new coordinates. When True, the output coordinates will be the
+            outer product of the kernel shape and the input coordinates.
+            `False` by default.
+
+            :attr:`dimension` (int): the spatial dimension of the space where
+            all the inputs and the network are defined. For example, images are
+            in a 2D space, meshes and 3D shapes are in a 3D space.
+
+        .. note:
+            TODO: support `kernel_size` > `stride`.
+
+        """
+        if kernel_generator is None:
+            kernel_generator = KernelGenerator(
+                kernel_size=kernel_size,
+                stride=stride,
+                dilation=dilation,
+                dimension=dimension,
+            )
+
+        MinkowskiConvolutionBase.__init__(
+            self,
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            dilation,
+            bias,
+            kernel_generator,
+            is_transpose=True,
+            expand_coordinates=expand_coordinates,
+            convolution_mode=convolution_mode,
+            dimension=dimension,
+        )
+        self.reset_parameters(True)
+
+
+class MinkowskiGenerativeConvolutionTranspose(MinkowskiConvolutionBase):
+    r"""A generalized sparse transposed convolution or deconvolution layer that
+    generates new coordinates.
+    """
+
+    def __init__(
+        self,
+        in_channels,
+        out_channels,
+        kernel_size=-1,
+        stride=1,
+        dilation=1,
+        bias=False,
+        kernel_generator=None,
+        convolution_mode=ConvolutionMode.DEFAULT,
+        dimension=None,
+    ):
+        r"""a generalized sparse transposed convolution layer that creates new coordinates.
+
+        Please refer to `Generative Sparse Detection Networks for 3D Single-shot Object Detection <https://arxiv.org/abs/2006.12356>`_ for more detail. Also, please cite the following paper if you use this function.
+
+        >> @inproceedings{gwak2020gsdn,
+        >>   title={Generative Sparse Detection Networks for 3D Single-shot Object Detection},
+        >>   author={Gwak, JunYoung and Choy, Christopher B and Savarese, Silvio},
+        >>   booktitle={European conference on computer vision},
+        >>   year={2020}
+        >> }
+
+        Args:
+            :attr:`in_channels` (int): the number of input channels in the
+            input tensor.
+
+            :attr:`out_channels` (int): the number of output channels in the
+            output tensor.
+
+            :attr:`kernel_size` (int, optional): the size of the kernel in the
+            output tensor. If not provided, :attr:`region_offset` should be
+            :attr:`RegionType.CUSTOM` and :attr:`region_offset` should be a 2D
+            matrix with size :math:`N\times D` such that it lists all :math:`N`
+            offsets in D-dimension.
+
+            :attr:`stride` (int, or list, optional): stride size that defines
+            upsampling rate. If non-identity is used, the output coordinates
+            will be :attr:`tensor_stride` / :attr:`stride` apart.  When a list is
+            given, the length must be D; each element will be used for stride
+            size for the specific axis.
+
+            :attr:`dilation` (int, or list, optional): dilation size for the
+            convolution kernel. When a list is given, the length must be D and
+            each element is an axis specific dilation. All elements must be > 0.
+
+            :attr:`bias` (bool, optional): if True, the convolution layer
+            has a bias.
+
+            :attr:`kernel_generator` (:attr:`MinkowskiEngine.KernelGenerator`,
+            optional): defines custom kernel shape.
+
+            :attr:`expand_coordinates` (bool, optional): Force generation of
+            new coordinates. When True, the output coordinates will be the
+            outer product of the kernel shape and the input coordinates.
+            `False` by defaul.
+
+            :attr:`dimension` (int): the spatial dimension of the space where
+            all the inputs and the network are defined. For example, images are
+            in a 2D space, meshes and 3D shapes are in a 3D space.
+
+        .. note:
+            TODO: support `kernel_size` > `stride`.
+
+        """
+        if kernel_generator is None:
+            kernel_generator = KernelGenerator(
+                kernel_size=kernel_size,
+                stride=stride,
+                dilation=dilation,
+                expand_coordinates=True,
+                dimension=dimension,
+            )
+        else:
+            kernel_generator.expand_coordinates = True
+
+        MinkowskiConvolutionBase.__init__(
+            self,
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            dilation,
+            bias,
+            kernel_generator,
+            is_transpose=True,
+            expand_coordinates=True,
+            convolution_mode=convolution_mode,
+            dimension=dimension,
+        )
+        self.reset_parameters(True)

MinkowskiEngine/MinkowskiEngine/MinkowskiCoordinateManager.py

@@ -0,0 +1,498 @@
+# 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.
+import os
+import numpy as np
+from collections.abc import Sequence
+from typing import Union, List, Tuple
+import warnings
+
+import torch
+from MinkowskiCommon import convert_to_int_list, convert_to_int_tensor
+import MinkowskiEngineBackend._C as _C
+from MinkowskiEngineBackend._C import (
+    CoordinateMapKey,
+    GPUMemoryAllocatorType,
+    CoordinateMapType,
+    MinkowskiAlgorithm,
+    RegionType,
+)
+
+CPU_COUNT = os.cpu_count()
+if "OMP_NUM_THREADS" in os.environ:
+    CPU_COUNT = int(os.environ["OMP_NUM_THREADS"])
+
+_allocator_type = GPUMemoryAllocatorType.PYTORCH
+_coordinate_map_type = (
+    CoordinateMapType.CUDA if _C.is_cuda_available() else CoordinateMapType.CPU
+)
+_minkowski_algorithm = MinkowskiAlgorithm.DEFAULT
+
+
+def set_coordinate_map_type(coordinate_map_type: CoordinateMapType):
+    r"""Set the default coordinate map type.
+
+    The MinkowskiEngine automatically set the coordinate_map_type to CUDA if
+    a NVIDIA GPU is available. To control the
+    """
+    global _coordinate_map_type
+    _coordinate_map_type = coordinate_map_type
+
+
+def set_gpu_allocator(backend: GPUMemoryAllocatorType):
+    r"""Set the GPU memory allocator
+
+    By default, the Minkowski Engine will use the pytorch memory pool to
+    allocate temporary GPU memory slots. This allows the pytorch backend to
+    effectively reuse the memory pool shared between the pytorch backend and
+    the Minkowski Engine. It tends to allow training with larger batch sizes
+    given a fixed GPU memory. However, pytorch memory manager tend to be slower
+    than allocating GPU directly using raw CUDA calls.
+
+    By default, the Minkowski Engine uses
+    :attr:`ME.GPUMemoryAllocatorType.PYTORCH` for memory management.
+
+    Example::
+
+       >>> import MinkowskiEngine as ME
+       >>> # Set the GPU memory manager backend to raw CUDA calls
+       >>> ME.set_gpu_allocator(ME.GPUMemoryAllocatorType.CUDA)
+       >>> # Set the GPU memory manager backend to the pytorch c10 allocator
+       >>> ME.set_gpu_allocator(ME.GPUMemoryAllocatorType.PYTORCH)
+
+    """
+    assert isinstance(
+        backend, GPUMemoryAllocatorType
+    ), f"Input must be an instance of GPUMemoryAllocatorType not {backend}"
+    global _allocator_type
+    _allocator_type = backend
+
+
+def set_memory_manager_backend(backend: GPUMemoryAllocatorType):
+    r"""Alias for set_gpu_allocator. Deprecated and will be removed."""
+    warnings.warn(
+        "`set_memory_manager_backend` has been deprecated. Use `set_gpu_allocator` instead."
+    )
+    set_gpu_allocator(backend)
+
+
+class CoordsManager:
+    def __init__(*args, **kwargs):
+        raise DeprecationWarning(
+            "`CoordsManager` has been deprecated. Use `CoordinateManager` instead."
+        )
+
+
+class CoordinateManager:
+    def __init__(
+        self,
+        D: int = 0,
+        num_threads: int = -1,
+        coordinate_map_type: CoordinateMapType = None,
+        allocator_type: GPUMemoryAllocatorType = None,
+        minkowski_algorithm: MinkowskiAlgorithm = None,
+    ):
+        r"""
+
+        :attr:`D`: The order, or dimension of the coordinates.
+        """
+        global _coordinate_map_type, _allocator_type, _minkowski_algorithm
+        if D < 1:
+            raise ValueError(f"Invalid rank D > 0, D = {D}.")
+        if num_threads < 0:
+            num_threads = min(CPU_COUNT, 20)
+        if coordinate_map_type is None:
+            coordinate_map_type = _coordinate_map_type
+        if allocator_type is None:
+            allocator_type = _allocator_type
+        if minkowski_algorithm is None:
+            minkowski_algorithm = _minkowski_algorithm
+
+        postfix = ""
+        if coordinate_map_type == CoordinateMapType.CPU:
+            postfix = "CPU"
+        else:
+            assert (
+                _C.is_cuda_available()
+            ), "The MinkowskiEngine was compiled with CPU_ONLY flag. If you want to compile with CUDA support, make sure `torch.cuda.is_available()` is True when you install MinkowskiEngine."
+
+            postfix = "GPU" + (
+                "_default" if allocator_type == GPUMemoryAllocatorType.CUDA else "_c10"
+            )
+
+        self.D = D
+        self.minkowski_algorithm = minkowski_algorithm
+        self._CoordinateManagerClass = getattr(_C, "CoordinateMapManager" + postfix)
+        self._manager = self._CoordinateManagerClass(minkowski_algorithm, num_threads)
+
+    # TODO: insert without remap, unique_map, inverse_mapa
+    #
+    # def insert() -> CoordinateMapKey
+
+    def insert_and_map(
+        self,
+        coordinates: torch.Tensor,
+        tensor_stride: Union[int, Sequence, np.ndarray] = 1,
+        string_id: str = "",
+    ) -> Tuple[CoordinateMapKey, Tuple[torch.IntTensor, torch.IntTensor]]:
+        r"""create a new coordinate map and returns (key, (map, inverse_map)).
+
+        :attr:`coordinates`: `torch.Tensor` (Int tensor. `CUDA` if
+        coordinate_map_type == `CoordinateMapType.GPU`) that defines the
+        coordinates.
+
+        :attr:`tensor_stride` (`list`): a list of `D` elements that defines the
+        tensor stride for the new order-`D + 1` sparse tensor.
+
+        Example::
+
+           >>> manager = CoordinateManager(D=1)
+           >>> coordinates = torch.IntTensor([[0, 0], [0, 0], [0, 1], [0, 2]])
+           >>> key, (unique_map, inverse_map) = manager.insert(coordinates, [1])
+           >>> print(key) # key is tensor_stride, string_id [1]:""
+           >>> torch.all(coordinates[unique_map] == manager.get_coordinates(key)) # True
+           >>> torch.all(coordinates == coordinates[unique_map][inverse_map]) # True
+
+        """
+        tensor_stride = convert_to_int_list(tensor_stride, self.D)
+        return self._manager.insert_and_map(coordinates, tensor_stride, string_id)
+
+    def insert_field(
+        self,
+        coordinates: torch.Tensor,
+        tensor_stride: Sequence,
+        string_id: str = "",
+    ) -> Tuple[CoordinateMapKey, Tuple[torch.IntTensor, torch.IntTensor]]:
+        r"""create a new coordinate map and returns
+
+        :attr:`coordinates`: `torch.FloatTensor` (`CUDA` if coordinate_map_type
+        == `CoordinateMapType.GPU`) that defines the coordinates.
+
+        :attr:`tensor_stride` (`list`): a list of `D` elements that defines the
+        tensor stride for the new order-`D + 1` sparse tensor.
+
+
+        Example::
+
+           >>> manager = CoordinateManager(D=1)
+           >>> coordinates = torch.FloatTensor([[0, 0.1], [0, 2.3], [0, 1.2], [0, 2.4]])
+           >>> key, (unique_map, inverse_map) = manager.insert(coordinates, [1])
+           >>> print(key) # key is tensor_stride, string_id [1]:""
+           >>> torch.all(coordinates[unique_map] == manager.get_coordinates(key)) # True
+           >>> torch.all(coordinates == coordinates[unique_map][inverse_map]) # True
+
+        """
+        return self._manager.insert_field(coordinates, tensor_stride, string_id)
+
+    def field_to_sparse_insert_and_map(
+        self,
+        field_map_key: CoordinateMapKey,
+        sparse_tensor_stride: Union[int, Sequence, np.ndarray],
+        sparse_tensor_string_id: str = "",
+    ) -> Tuple[CoordinateMapKey, Tuple[torch.IntTensor, torch.IntTensor]]:
+
+        r"""Create a sparse tensor coordinate map with the tensor stride.
+
+        :attr:`field_map_key` (`CoordinateMapKey`): field map that a new sparse
+        tensor will be created from.
+
+        :attr:`tensor_stride` (`list`): a list of `D` elements that defines the
+        tensor stride for the new order-`D + 1` sparse tensor.
+
+        :attr:`string_id` (`str`): string id of the new sparse tensor coordinate map key.
+
+        Example::
+
+           >>> manager = CoordinateManager(D=1)
+           >>> coordinates = torch.FloatTensor([[0, 0.1], [0, 2.3], [0, 1.2], [0, 2.4]])
+           >>> key, (unique_map, inverse_map) = manager.insert(coordinates, [1])
+
+        """
+        return self._manager.field_to_sparse_insert_and_map(
+            field_map_key, sparse_tensor_stride, sparse_tensor_string_id
+        )
+
+    def exists_field_to_sparse(
+        self, field_map_key: CoordinateMapKey, sparse_map_key: CoordinateMapKey
+    ):
+        return self._manager.exists_field_to_sparse(field_map_key, sparse_map_key)
+
+    def field_to_sparse_keys(self, field_map_key: CoordinateMapKey):
+        return self._manager.field_to_sparse_keys(field_map_key.get_key())
+
+    def get_field_to_sparse_map(
+        self, field_map_key: CoordinateMapKey, sparse_map_key: CoordinateMapKey
+    ):
+        return self._manager.get_field_to_sparse_map(field_map_key, sparse_map_key)
+
+    def field_to_sparse_map(
+        self, field_map_key: CoordinateMapKey, sparse_map_key: CoordinateMapKey
+    ):
+        return self._manager.field_to_sparse_map(field_map_key, sparse_map_key)
+
+    def stride(
+        self,
+        coordinate_map_key: CoordinateMapKey,
+        stride: Union[int, Sequence, np.ndarray, torch.Tensor],
+        string_id: str = "",
+    ) -> CoordinateMapKey:
+        r"""Generate a new coordinate map and returns the key.
+
+        :attr:`coordinate_map_key` (:attr:`MinkowskiEngine.CoordinateMapKey`):
+        input map to generate the strided map from.
+
+        :attr:`stride`: stride size.
+        """
+        stride = convert_to_int_list(stride, self.D)
+        return self._manager.stride(coordinate_map_key, stride, string_id)
+
+    def origin(self) -> CoordinateMapKey:
+        return self._manager.origin()
+
+    def origin_field(self) -> CoordinateMapKey:
+        return self._manager.origin_field()
+
+    def size(self, coordinate_map_key: CoordinateMapKey) -> int:
+        return self._manager.size(coordinate_map_key)
+
+    # def transposed_stride(
+    #     self,
+    #     coords_key: CoordsKey,
+    #     stride: Union[int, Sequence, np.ndarray, torch.Tensor],
+    #     kernel_size: Union[int, Sequence, np.ndarray, torch.Tensor],
+    #     dilation: Union[int, Sequence, np.ndarray, torch.Tensor],
+    #     force_creation: bool = False,
+    # ):
+    #     assert isinstance(coords_key, CoordsKey)
+    #     stride = convert_to_int_list(stride, self.D)
+    #     kernel_size = convert_to_int_list(kernel_size, self.D)
+    #     dilation = convert_to_int_list(dilation, self.D)
+    #     region_type = 0
+    #     region_offset = torch.IntTensor()
+
+    #     strided_key = CoordsKey(self.D)
+    #     tensor_stride = coords_key.getTensorStride()
+    #     strided_key.setTensorStride([int(t / s) for t, s in zip(tensor_stride, stride)])
+
+    #     strided_key.setKey(
+    #         self.CPPCoordsManager.createTransposedStridedRegionCoords(
+    #             coords_key.getKey(),
+    #             coords_key.getTensorStride(),
+    #             stride,
+    #             kernel_size,
+    #             dilation,
+    #             region_type,
+    #             region_offset,
+    #             force_creation,
+    #         )
+    #     )
+    #     return strided_key
+
+    def _get_coordinate_map_key(self, key_or_tensor_strides) -> CoordinateMapKey:
+        r"""Helper function that retrieves the first coordinate map key for the given tensor stride."""
+        assert isinstance(key_or_tensor_strides, CoordinateMapKey) or isinstance(
+            key_or_tensor_strides, (Sequence, np.ndarray, torch.IntTensor, int)
+        ), f"The input must be either a CoordinateMapKey or tensor_stride of type (int, list, tuple, array, Tensor). Invalid: {key_or_tensor_strides}"
+        if isinstance(key_or_tensor_strides, CoordinateMapKey):
+            # Do nothing and return the input
+            return key_or_tensor_strides
+        else:
+            tensor_strides = convert_to_int_list(key_or_tensor_strides, self.D)
+            keys = self._manager.get_coordinate_map_keys(tensor_strides)
+            assert len(keys) > 0
+            return keys[0]
+
+    def get_coordinates(self, coords_key_or_tensor_strides) -> torch.Tensor:
+        key = self._get_coordinate_map_key(coords_key_or_tensor_strides)
+        return self._manager.get_coordinates(key)
+
+    def get_coordinate_field(self, coords_key_or_tensor_strides) -> torch.Tensor:
+        key = self._get_coordinate_map_key(coords_key_or_tensor_strides)
+        return self._manager.get_coordinate_field(key)
+
+    def number_of_unique_batch_indices(self) -> int:
+        return self._manager.origin_map_size()
+
+    def get_unique_coordinate_map_key(
+        self, tensor_stride: Union[int, list]
+    ) -> CoordinateMapKey:
+        """
+        Returns a unique coordinate_map_key for a given tensor stride.
+
+        :attr:`tensor_stride` (`list`): a list of `D` elements that defines the
+        tensor stride for the new order-`D + 1` sparse tensor.
+
+        """
+        return self._manager.get_random_string_id(tensor_stride, "")
+
+    def get_kernel_map(
+        self,
+        in_key: CoordinateMapKey,
+        out_key: CoordinateMapKey,
+        stride=1,
+        kernel_size=3,
+        dilation=1,
+        region_type=RegionType.HYPER_CUBE,
+        region_offset=None,
+        is_transpose=False,
+        is_pool=False,
+    ) -> dict:
+        r"""Alias of :attr:`CoordinateManager.kernel_map`. Will be deprecated in the next version."""
+        warnings.warn(
+            "`get_kernel_map` will be deprecated. Please use `kernel_map` instead."
+        )
+        return self.kernel_map(
+            in_key,
+            out_key,
+            stride,
+            kernel_size,
+            dilation,
+            region_type,
+            region_offset,
+            is_transpose,
+            is_pool,
+        )
+
+    def kernel_map(
+        self,
+        in_key: CoordinateMapKey,
+        out_key: CoordinateMapKey,
+        stride=1,
+        kernel_size=3,
+        dilation=1,
+        region_type=RegionType.HYPER_CUBE,
+        region_offset=None,
+        is_transpose=False,
+        is_pool=False,
+    ) -> dict:
+        r"""Get kernel in-out maps for the specified coords keys or tensor strides.
+
+        returns dict{kernel_index: in_out_tensor} where in_out_tensor[0] is the input row indices that correspond to in_out_tensor[1], which is the row indices for output.
+        """
+        # region type 1 iteration with kernel_size 1 is invalid
+        if isinstance(kernel_size, torch.Tensor):
+            assert (kernel_size > 0).all(), f"Invalid kernel size: {kernel_size}"
+            if (kernel_size == 1).all() == 1:
+                region_type = RegionType.HYPER_CUBE
+        elif isinstance(kernel_size, int):
+            assert kernel_size > 0, f"Invalid kernel size: {kernel_size}"
+            if kernel_size == 1:
+                region_type = RegionType.HYPER_CUBE
+
+        in_key = self._get_coordinate_map_key(in_key)
+        out_key = self._get_coordinate_map_key(out_key)
+
+        if region_offset is None:
+            region_offset = torch.IntTensor()
+
+        kernel_map = self._manager.kernel_map(
+            in_key,
+            out_key,
+            convert_to_int_list(kernel_size, self.D),  #
+            convert_to_int_list(stride, self.D),  #
+            convert_to_int_list(dilation, self.D),  #
+            region_type,
+            region_offset,
+            is_transpose,
+            is_pool,
+        )
+
+        return kernel_map
+
+    def origin_map(self, key: CoordinateMapKey):
+        return self._manager.origin_map(key)
+
+    def origin_field_map(self, key: CoordinateMapKey):
+        return self._manager.origin_field_map(key)
+
+    def stride_map(self, in_key: CoordinateMapKey, stride_key: CoordinateMapKey):
+        return self._manager.stride_map(in_key, stride_key)
+
+    def union_map(self, in_keys: list, out_key):
+        return self._manager.union_map(in_keys, out_key)
+
+    def interpolation_map_weight(
+        self,
+        key: CoordinateMapKey,
+        samples: torch.Tensor,
+    ):
+        return self._manager.interpolation_map_weight(samples, key)
+
+    # def get_union_map(self, in_keys: List[CoordsKey], out_key: CoordsKey):
+    #     r"""Generates a union of coordinate sets and returns the mapping from input sets to the new output coordinates.
+
+    #     Args:
+    #         :attr:`in_keys` (List[CoordsKey]): A list of coordinate keys to
+    #         create a union on.
+
+    #         :attr:`out_key` (CoordsKey): the placeholder for the coords key of
+    #         the generated union coords hash map.
+
+    #     Returns:
+    #         :attr:`in_maps` (List[Tensor[int]]): A list of long tensors that contain mapping from inputs to the union output. Please see the example for more details.
+    #         :attr:`out_maps` (List[Tensor[int]]): A list of long tensors that contain a mapping from input to the union output. Please see the example for more details.
+
+    #     Example::
+
+    #         >>> # Adding two sparse tensors: A, B
+    #         >>> out_key = CoordsKey(coords_man.D)
+    #         >>> ins, outs = coords_man.get_union_map((A.coords_key, B.coords_key), out_key)
+    #         >>> N = coords_man.get_coords_size_by_coords_key(out_key)
+    #         >>> out_F = torch.zeros((N, A.F.size(1)), dtype=A.dtype)
+    #         >>> out_F[outs[0]] = A.F[ins[0]]
+    #         >>> out_F[outs[1]] += B.F[ins[1]]
+
+    #     """
+    #     return self.CPPCoordsManager.getUnionMap(
+    #         [key.CPPCoordsKey for key in in_keys], out_key.CPPCoordsKey
+    #     )
+
+    # def permute_label(
+    #     self, label, max_label, target_tensor_stride, label_tensor_stride=1
+    # ):
+    #     if target_tensor_stride == label_tensor_stride:
+    #         return label
+
+    #     label_coords_key = self._get_coordinate_map_key(label_tensor_stride)
+    #     target_coords_key = self._get_coordinate_map_key(target_tensor_stride)
+
+    #     permutation = self.get_mapping_by_coords_key(
+    #         label_coords_key, target_coords_key
+    #     )
+    #     nrows = self.get_coords_size_by_coords_key(target_coords_key)
+
+    #     label = label.contiguous().numpy()
+    #     permutation = permutation.numpy()
+
+    #     counter = np.zeros((nrows, max_label), dtype="int32")
+    #     np.add.at(counter, (permutation, label), 1)
+    #     return torch.from_numpy(np.argmax(counter, 1))
+
+    def __repr__(self):
+        return (
+            self._CoordinateManagerClass.__name__
+            + "(\n"
+            + str(self._manager)
+            + f"\talgorithm={self.minkowski_algorithm}\n  )\n"
+        )

MinkowskiEngine/MinkowskiEngine/MinkowskiFunctional.py

@@ -0,0 +1,232 @@
+# Copyright (c) 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.
+import torch.nn.functional as F
+
+from MinkowskiSparseTensor import SparseTensor
+from MinkowskiTensorField import TensorField
+
+
+def _wrap_tensor(input, F):
+    if isinstance(input, TensorField):
+        return TensorField(
+            F,
+            coordinate_field_map_key=input.coordinate_field_map_key,
+            coordinate_manager=input.coordinate_manager,
+            quantization_mode=input.quantization_mode,
+        )
+    else:
+        return SparseTensor(
+            F,
+            coordinate_map_key=input.coordinate_map_key,
+            coordinate_manager=input.coordinate_manager,
+        )
+
+
+# Activations
+def threshold(input, *args, **kwargs):
+    return _wrap_tensor(input, F.threshold(input.F, *args, **kwargs))
+
+
+def relu(input, *args, **kwargs):
+    return _wrap_tensor(input, F.relu(input.F, *args, **kwargs))
+
+
+def hardtanh(input, *args, **kwargs):
+    return _wrap_tensor(input, F.hardtanh(input.F, *args, **kwargs))
+
+
+def hardswish(input, *args, **kwargs):
+    return _wrap_tensor(input, F.hardswish(input.F, *args, **kwargs))
+
+
+def relu6(input, *args, **kwargs):
+    return _wrap_tensor(input, F.relu6(input.F, *args, **kwargs))
+
+
+def elu(input, *args, **kwargs):
+    return _wrap_tensor(input, F.elu(input.F, *args, **kwargs))
+
+
+def selu(input, *args, **kwargs):
+    return _wrap_tensor(input, F.selu(input.F, *args, **kwargs))
+
+
+def celu(input, *args, **kwargs):
+    return _wrap_tensor(input, F.celu(input.F, *args, **kwargs))
+
+
+def leaky_relu(input, *args, **kwargs):
+    return _wrap_tensor(input, F.leaky_relu(input.F, *args, **kwargs))
+
+
+def prelu(input, *args, **kwargs):
+    return _wrap_tensor(input, F.prelu(input.F, *args, **kwargs))
+
+
+def rrelu(input, *args, **kwargs):
+    return _wrap_tensor(input, F.rrelu(input.F, *args, **kwargs))
+
+
+def glu(input, *args, **kwargs):
+    return _wrap_tensor(input, F.glu(input.F, *args, **kwargs))
+
+
+def gelu(input, *args, **kwargs):
+    return _wrap_tensor(input, F.gelu(input.F, *args, **kwargs))
+
+
+def logsigmoid(input, *args, **kwargs):
+    return _wrap_tensor(input, F.logsigmoid(input.F, *args, **kwargs))
+
+
+def hardshrink(input, *args, **kwargs):
+    return _wrap_tensor(input, F.hardshrink(input.F, *args, **kwargs))
+
+
+def tanhshrink(input, *args, **kwargs):
+    return _wrap_tensor(input, F.tanhshrink(input.F, *args, **kwargs))
+
+
+def softsign(input, *args, **kwargs):
+    return _wrap_tensor(input, F.softsign(input.F, *args, **kwargs))
+
+
+def softplus(input, *args, **kwargs):
+    return _wrap_tensor(input, F.softplus(input.F, *args, **kwargs))
+
+
+def softmin(input, *args, **kwargs):
+    return _wrap_tensor(input, F.softmin(input.F, *args, **kwargs))
+
+
+def softmax(input, *args, **kwargs):
+    return _wrap_tensor(input, F.softmax(input.F, *args, **kwargs))
+
+
+def softshrink(input, *args, **kwargs):
+    return _wrap_tensor(input, F.softshrink(input.F, *args, **kwargs))
+
+
+def gumbel_softmax(input, *args, **kwargs):
+    return _wrap_tensor(input, F.gumbel_softmax(input.F, *args, **kwargs))
+
+
+def log_softmax(input, *args, **kwargs):
+    return _wrap_tensor(input, F.log_softmax(input.F, *args, **kwargs))
+
+
+def tanh(input, *args, **kwargs):
+    return _wrap_tensor(input, F.tanh(input.F, *args, **kwargs))
+
+
+def sigmoid(input, *args, **kwargs):
+    return _wrap_tensor(input, F.sigmoid(input.F, *args, **kwargs))
+
+
+def hardsigmoid(input, *args, **kwargs):
+    return _wrap_tensor(input, F.hardsigmoid(input.F, *args, **kwargs))
+
+
+def silu(input, *args, **kwargs):
+    return _wrap_tensor(input, F.silu(input.F, *args, **kwargs))
+
+
+# Normalization
+def batch_norm(input, *args, **kwargs):
+    return _wrap_tensor(input, F.batch_norm(input.F, *args, **kwargs))
+
+
+def normalize(input, *args, **kwargs):
+    return _wrap_tensor(input, F.normalize(input.F, *args, **kwargs))
+
+
+# Linear
+def linear(input, *args, **kwargs):
+    return _wrap_tensor(input, F.linear(input.F, *args, **kwargs))
+
+
+# Dropouts
+def dropout(input, *args, **kwargs):
+    return _wrap_tensor(input, F.dropout(input.F, *args, **kwargs))
+
+
+def alpha_dropout(input, *args, **kwargs):
+    return _wrap_tensor(input, F.alpha_dropout(input.F, *args, **kwargs))
+
+
+# Loss functions
+def binary_cross_entropy(input, target, *args, **kwargs):
+    return F.binary_cross_entropy(input.F, target, *args, **kwargs)
+
+
+def binary_cross_entropy_with_logits(input, target, *args, **kwargs):
+    return F.binary_cross_entropy_with_logits(input.F, target, *args, **kwargs)
+
+
+def poisson_nll_loss(input, target, *args, **kwargs):
+    return F.poisson_nll_loss(input.F, target, *args, **kwargs)
+
+
+def cross_entropy(input, target, *args, **kwargs):
+    return F.cross_entropy(input.F, target, *args, **kwargs)
+
+
+def hinge_embedding_loss(input, target, *args, **kwargs):
+    return F.hinge_embedding_loss(input.F, target, *args, **kwargs)
+
+
+def kl_div(input, target, *args, **kwargs):
+    return F.kl_div(input.F, target, *args, **kwargs)
+
+
+def l1_loss(input, target, *args, **kwargs):
+    return F.l1_loss(input.F, target, *args, **kwargs)
+
+
+def mse_loss(input, target, *args, **kwargs):
+    return F.mse_loss(input.F, target, *args, **kwargs)
+
+
+def multilabel_margin_loss(input, target, *args, **kwargs):
+    return F.multilabel_margin_loss(input.F, target, *args, **kwargs)
+
+
+def multilabel_soft_margin_loss(input, target, *args, **kwargs):
+    return F.multilabel_soft_margin_loss(input.F, target, *args, **kwargs)
+
+
+def multi_margin_loss(input, target, *args, **kwargs):
+    return F.multi_margin_loss(input.F, target, *args, **kwargs)
+
+
+def nll_loss(input, target, *args, **kwargs):
+    return F.nll_loss(input.F, target, *args, **kwargs)
+
+
+def smooth_l1_loss(input, target, *args, **kwargs):
+    return F.smooth_l1_loss(input.F, target, *args, **kwargs)
+
+
+def soft_margin_loss(input, target, *args, **kwargs):
+    return F.soft_margin_loss(input.F, target, *args, **kwargs)

MinkowskiEngine/MinkowskiEngine/MinkowskiInterpolation.py

@@ -0,0 +1,131 @@
+# 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.
+from typing import Union
+
+import torch
+from torch.autograd import Function
+
+from MinkowskiEngineBackend._C import CoordinateMapKey
+from MinkowskiSparseTensor import SparseTensor
+from MinkowskiCoordinateManager import CoordinateManager
+from MinkowskiCommon import (
+    MinkowskiModuleBase,
+    get_minkowski_function,
+)
+
+
+class MinkowskiInterpolationFunction(Function):
+    @staticmethod
+    def forward(
+        ctx,
+        input_features: torch.Tensor,
+        tfield: torch.Tensor,
+        in_coordinate_map_key: CoordinateMapKey,
+        coordinate_manager: CoordinateManager = None,
+    ):
+        input_features = input_features.contiguous()
+        # in_map, out_map, weights = coordinate_manager.interpolation_map_weight(
+        #     in_coordinate_map_key, tfield)
+        fw_fn = get_minkowski_function("InterpolationForward", input_features)
+        out_feat, in_map, out_map, weights = fw_fn(
+            input_features,
+            tfield,
+            in_coordinate_map_key,
+            coordinate_manager._manager,
+        )
+        ctx.save_for_backward(in_map, out_map, weights)
+        ctx.inputs = (
+            in_coordinate_map_key,
+            coordinate_manager,
+        )
+        return out_feat, in_map, out_map, weights
+
+    @staticmethod
+    def backward(
+        ctx, grad_out_feat=None, grad_in_map=None, grad_out_map=None, grad_weights=None
+    ):
+        grad_out_feat = grad_out_feat.contiguous()
+        bw_fn = get_minkowski_function("InterpolationBackward", grad_out_feat)
+        (
+            in_coordinate_map_key,
+            coordinate_manager,
+        ) = ctx.inputs
+        in_map, out_map, weights = ctx.saved_tensors
+
+        grad_in_feat = bw_fn(
+            grad_out_feat,
+            in_map,
+            out_map,
+            weights,
+            in_coordinate_map_key,
+            coordinate_manager._manager,
+        )
+        return grad_in_feat, None, None, None
+
+
+class MinkowskiInterpolation(MinkowskiModuleBase):
+    r"""Sample linearly interpolated features at the provided points."""
+
+    def __init__(self, return_kernel_map=False, return_weights=False):
+        r"""Sample linearly interpolated features at the specified coordinates.
+
+        Args:
+            :attr:`return_kernel_map` (bool): In addition to the sampled
+            features, the layer returns the kernel map as a pair of input row
+            indices and output row indices. False by default.
+
+            :attr:`return_weights` (bool): When True, return the linear
+            interpolation weights. False by default.
+        """
+        MinkowskiModuleBase.__init__(self)
+        self.return_kernel_map = return_kernel_map
+        self.return_weights = return_weights
+        self.interp = MinkowskiInterpolationFunction()
+
+    def forward(
+        self,
+        input: SparseTensor,
+        tfield: torch.Tensor,
+    ):
+        # Get a new coordinate map key or extract one from the coordinates
+        out_feat, in_map, out_map, weights = self.interp.apply(
+            input.F,
+            tfield,
+            input.coordinate_map_key,
+            input._manager,
+        )
+
+        return_args = [out_feat]
+        if self.return_kernel_map:
+            return_args.append((in_map, out_map))
+        if self.return_weights:
+            return_args.append(weights)
+        if len(return_args) > 1:
+            return tuple(return_args)
+        else:
+            return out_feat
+
+    def __repr__(self):
+        return self.__class__.__name__ + "()"

MinkowskiEngine/MinkowskiEngine/MinkowskiKernelGenerator.py

@@ -0,0 +1,395 @@
+# 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.
+import math
+from collections import namedtuple
+from collections.abc import Sequence
+from functools import reduce
+import numpy as np
+from typing import Union
+
+import torch
+from MinkowskiCommon import convert_to_int_list
+from MinkowskiEngineBackend._C import CoordinateMapKey, RegionType
+from MinkowskiCoordinateManager import CoordinateManager
+
+
+def get_kernel_volume(region_type, kernel_size, region_offset, axis_types, dimension):
+    """
+    when center is True, the custom region_offset will be centered at the
+    origin. Currently, for HYPER_CUBE, HYPER_CROSS with odd kernel sizes cannot
+    use center=False.
+    """
+    if region_type == RegionType.HYPER_CUBE:
+        assert reduce(
+            lambda k1, k2: k1 > 0 and k2 > 0, kernel_size
+        ), "kernel_size must be positive"
+        assert (
+            region_offset is None
+        ), "Region offset must be None when region_type is given"
+        assert axis_types is None, "Axis types must be None when region_type is given"
+        # Typical convolution kernel
+
+        # Convolution kernel with even numbered kernel size not defined.
+        kernel_volume = torch.prod(torch.IntTensor(kernel_size)).item()
+
+    elif region_type == RegionType.HYPER_CROSS:
+        assert reduce(
+            lambda k1, k2: k1 > 0 and k2 > 0, kernel_size
+        ), "kernel_size must be positive"
+        assert (
+            torch.IntTensor(kernel_size) % 2
+        ).prod().item() == 1, "kernel_size must be odd for region_type HYPER_CROSS"
+        # 0th: itself, (1, 2) for 0th dim neighbors, (3, 4) for 1th dim ...
+        kernel_volume = (torch.sum(torch.IntTensor(kernel_size) - 1) + 1).item()
+
+    # elif region_type == RegionType.HYBRID:
+    #     assert reduce(
+    #         lambda k1, k2: k1 > 0 and k2 > 0, kernel_size
+    #     ), "kernel_size must be positive"
+    #     assert (
+    #         region_offset is None
+    #     ), "region_offset must be None when region_type is HYBRID"
+    #     kernel_size_list = kernel_size.tolist()
+    #     kernel_volume = 1
+    #     # First HYPER_CUBE
+    #     for axis_type, curr_kernel_size, d in zip(
+    #         axis_types, kernel_size_list, range(dimension)
+    #     ):
+    #         if axis_type == RegionType.HYPER_CUBE:
+    #             kernel_volume *= curr_kernel_size
+
+    #     # Second, HYPER_CROSS
+    #     for axis_type, curr_kernel_size, d in zip(
+    #         axis_types, kernel_size_list, range(dimension)
+    #     ):
+    #         if axis_type == RegionType.HYPER_CROSS:
+    #             kernel_volume += curr_kernel_size - 1
+
+    elif region_type == RegionType.CUSTOM:
+        assert (
+            region_offset.numel() > 0
+        ), "region_offset must be non empty when region_type is CUSTOM"
+        assert (
+            region_offset.size(1) == dimension
+        ), "region_offset must have the same dimension as the network"
+        kernel_volume = int(region_offset.size(0))
+
+    else:
+        raise NotImplementedError()
+
+    return kernel_volume
+
+
+def convert_region_type(
+    region_type: RegionType,
+    tensor_stride: Union[Sequence, np.ndarray, torch.IntTensor],
+    kernel_size: Union[Sequence, np.ndarray, torch.IntTensor],
+    up_stride: Union[Sequence, np.ndarray, torch.IntTensor],
+    dilation: Union[Sequence, np.ndarray, torch.IntTensor],
+    region_offset: Union[Sequence, np.ndarray, torch.IntTensor],
+    axis_types: Union[Sequence, np.ndarray, torch.IntTensor],
+    dimension: int,
+    center: bool = True,
+):
+    """
+    when center is True, the custom region_offset will be centered at the
+    origin. Currently, for HYPER_CUBE, HYPER_CROSS with odd kernel sizes cannot
+    use center=False.
+
+    up_stride: stride for conv_transpose, otherwise set it as 1
+    """
+    if region_type == RegionType.HYPER_CUBE:
+        if isinstance(region_offset, torch.Tensor):
+            assert (
+                region_offset.numel() == 0
+            ), "Region offset must be empty when region_type is given"
+        else:
+            assert (
+                region_offset is None
+            ), "Region offset must be None when region_type is given"
+
+        assert axis_types is None, "Axis types must be None when region_type is given"
+        # Typical convolution kernel
+        assert reduce(
+            lambda k1, k2: k1 > 0 and k2 > 0, kernel_size
+        ), "kernel_size must be positive"
+        # assert torch.unique(dilation).numel() == 1
+        kernel_volume = reduce(lambda k1, k2: k1 * k2, kernel_size)
+
+    elif region_type == RegionType.HYPER_CROSS:
+        assert reduce(
+            lambda k1, k2: k1 > 0 and k2 > 0, kernel_size
+        ), "kernel_size must be positive"
+        assert (
+            kernel_size % 2
+        ).prod() == 1, "kernel_size must be odd for region_type HYPER_CROSS"
+        # 0th: itself, (1, 2) for 0th dim neighbors, (3, 4) for 1th dim ...
+        kernel_volume = (
+            reduce(lambda k1, k2: k1 + k2, map(lambda k: k - 1, kernel_size)) + 1
+        )
+
+    elif region_type == RegionType.HYBRID:
+        assert reduce(
+            lambda k1, k2: k1 > 0 and k2 > 0, kernel_size
+        ), "kernel_size must be positive"
+        if isinstance(region_offset, torch.Tensor):
+            assert (
+                region_offset.numel() == 0
+            ), "Region offset must be empty when region_type is given"
+        else:
+            assert (
+                region_offset is None
+            ), "Region offset must be None when region_type is given"
+
+        region_offset = [
+            [
+                0,
+            ]
+            * dimension
+        ]
+        kernel_size_list = kernel_size.tolist()
+        # First HYPER_CUBE
+        for axis_type, curr_kernel_size, d in zip(
+            axis_types, kernel_size_list, range(dimension)
+        ):
+            new_offset = []
+            if axis_type == RegionType.HYPER_CUBE:
+                for offset in region_offset:
+                    for curr_offset in range(curr_kernel_size):
+                        off_center = (
+                            int(math.floor((curr_kernel_size - 1) / 2)) if center else 0
+                        )
+                        offset = offset.copy()  # Do not modify the original
+                        # Exclude the coord (0, 0, ..., 0)
+                        if curr_offset == off_center:
+                            continue
+                        offset[d] = (
+                            (curr_offset - off_center)
+                            * dilation[d]
+                            * (tensor_stride[d] / up_stride[d])
+                        )
+                        new_offset.append(offset)
+            region_offset.extend(new_offset)
+
+        # Second, HYPER_CROSS
+        for axis_type, curr_kernel_size, d in zip(
+            axis_types, kernel_size_list, range(dimension)
+        ):
+            new_offset = []
+            if axis_type == RegionType.HYPER_CROSS:
+                for curr_offset in range(curr_kernel_size):
+                    off_center = (
+                        int(math.floor((curr_kernel_size - 1) / 2)) if center else 0
+                    )
+                    offset = [
+                        0,
+                    ] * dimension
+                    # Exclude the coord (0, 0, ..., 0)
+                    if curr_offset == off_center:
+                        continue
+                    offset[d] = (
+                        (curr_offset - off_center)
+                        * dilation[d]
+                        * (tensor_stride[d] / up_stride[d])
+                    )
+                    new_offset.append(offset)
+            region_offset.extend(new_offset)
+
+        # Convert to CUSTOM type
+        region_type = RegionType.CUSTOM
+        region_offset = torch.IntTensor(region_offset)
+        kernel_volume = int(region_offset.size(0))
+
+    elif region_type == RegionType.CUSTOM:
+        assert (
+            region_offset.numel() > 0
+        ), "region_offset must be non empty when region_type is CUSTOM"
+        assert (
+            region_offset.size(1) == dimension
+        ), "region_offset must have the same dimension as the network"
+        kernel_volume = int(region_offset.size(0))
+        assert isinstance(
+            region_offset.dtype, torch.IntTensor
+        ), "region_offset must be a torch.IntTensor."
+    else:
+        raise NotImplementedError()
+
+    if region_offset is None:
+        region_offset = torch.IntTensor()
+
+    return region_type, region_offset, kernel_volume
+
+
+class KernelGenerator:
+    __slots__ = (
+        "cache",
+        "kernel_size",
+        "kernel_stride",
+        "kernel_dilation",
+        "region_type",
+        "region_offsets",
+        "axis_types",
+        "dimension",
+        "kernel_volume",
+        "requires_strided_coordinates",
+        "expand_coordinates",
+    )
+
+    def __init__(
+        self,
+        kernel_size=-1,
+        stride=1,
+        dilation=1,
+        is_transpose: bool = False,
+        region_type: RegionType = RegionType.HYPER_CUBE,
+        region_offsets: torch.Tensor = None,
+        expand_coordinates: bool = False,
+        axis_types=None,
+        dimension=-1,
+    ):
+        r"""
+        :attr:`region_type` (RegionType, optional): defines the kernel
+        shape. Please refer to MinkowskiEngine.Comon for details.
+
+        :attr:`region_offset` (torch.IntTensor, optional): when the
+        :attr:`region_type` is :attr:`RegionType.CUSTOM`, the convolution
+        kernel uses the provided `region_offset` to define offsets. It
+        should be a matrix of size :math:`N \times D` where :math:`N` is
+        the number of offsets and :math:`D` is the dimension of the
+        space.
+
+        :attr:`axis_types` (list of RegionType, optional): If given, it
+        uses different methods to create a kernel for each axis. e.g., when
+        it is `[RegionType.HYPER_CUBE, RegionType.HYPER_CUBE,
+        RegionType.HYPER_CROSS]`, the kernel would be rectangular for the
+        first two dimensions and cross shaped for the thrid dimension.
+        """
+        assert dimension > 0
+        assert isinstance(region_type, RegionType)
+
+        kernel_size = convert_to_int_list(kernel_size, dimension)
+        kernel_stride = convert_to_int_list(stride, dimension)
+        kernel_dilation = convert_to_int_list(dilation, dimension)
+
+        self.cache = {}
+        self.kernel_size = kernel_size
+        self.kernel_stride = kernel_stride
+        self.kernel_dilation = kernel_dilation
+        self.region_type = region_type
+        self.region_offsets = region_offsets if region_offsets else torch.IntTensor()
+        self.axis_types = axis_types
+        self.dimension = dimension
+        self.kernel_volume = get_kernel_volume(
+            region_type, kernel_size, region_offsets, axis_types, dimension
+        )
+        self.requires_strided_coordinates = reduce(
+            lambda s1, s2: s1 == 1 and s2 == 1, kernel_stride
+        )
+        self.expand_coordinates = expand_coordinates
+
+    def get_kernel(self, tensor_stride, is_transpose):
+        assert len(tensor_stride) == self.dimension
+        if tuple(tensor_stride) not in self.cache:
+            up_stride = (
+                self.stride
+                if is_transpose
+                else torch.Tensor(
+                    [
+                        1,
+                    ]
+                    * self.dimension
+                )
+            )
+
+            self.cache[tuple(tensor_stride)] = convert_region_type(
+                self.region_type,
+                tensor_stride,
+                self.kernel_size,
+                up_stride,
+                self.kernel_dilation,
+                self.region_offsets,
+                self.axis_types,
+                self.dimension,
+            )
+
+        return self.cache[tuple(tensor_stride)]
+
+    def __repr__(self):
+        return (
+            self.__class__.__name__
+            + f"(kernel_size={self.kernel_size}, kernel_stride={self.kernel_stride}, kernel_dilation={self.kernel_dilation}, "
+            + f"region_type={self.region_type}, expand_coordinates={self.expand_coordinates}, dimension={self.dimension})"
+        )
+
+
+class KernelRegion(
+    namedtuple(
+        "KernelRegion",
+        (
+            "kernel_size",
+            "kernel_stride",
+            "kernel_dilation",
+            "region_type",
+            "offset",
+            "D",
+        ),
+    )
+):
+    """adding functionality to a named tuple"""
+
+    __slots__ = ()
+
+    def __init__(
+        self,
+        kernel_size,
+        kernel_stride,
+        kernel_dilation,
+        region_type,
+        offset,
+        dimension,
+    ):
+        kernel_size = convert_to_int_list(kernel_size, dimension)
+        kernel_stride = convert_to_int_list(kernel_stride, dimension)
+        kernel_dilation = convert_to_int_list(kernel_dilation, dimension)
+        super(KernelRegion, self).__init__(
+            kernel_size, kernel_stride, kernel_dilation, region_type, offset, dimension
+        )
+
+    def __str__(self):
+        return "kernel_size:{self.kernel_size}, kernel_stride:{self.kernel_stride}, region_type:{self.region_type}"
+
+
+def save_ctx(
+    ctx,  # function object context
+    kernel_generator: KernelGenerator,
+    in_coords_key: CoordinateMapKey,
+    out_coords_key: CoordinateMapKey,
+    coordinate_manager: CoordinateManager,
+):
+    ctx.kernel_generator = kernel_generator
+    ctx.in_coordinate_map_key = in_coords_key
+    ctx.out_coordinate_map_key = out_coords_key
+    ctx.coordinate_manager = coordinate_manager
+    return ctx

MinkowskiEngine/MinkowskiEngine/MinkowskiNetwork.py

@@ -0,0 +1,57 @@
+# Copyright (c) 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.
+from abc import ABC, abstractmethod
+
+import torch.nn as nn
+
+from MinkowskiSparseTensor import SparseTensor
+
+
+class MinkowskiNetwork(nn.Module, ABC):
+    """
+    MinkowskiNetwork: an abstract class for sparse convnets.
+
+    Note: All modules that use the same coordinates must use the same net_metadata
+    """
+
+    def __init__(self, D):
+        super(MinkowskiNetwork, self).__init__()
+        self.D = D
+
+    @abstractmethod
+    def forward(self, x):
+        pass
+
+    def init(self, x):
+        """
+        Initialize coordinates if it does not exist
+        """
+        nrows = self.get_nrows(1)
+        if nrows < 0:
+            if isinstance(x, SparseTensor):
+                self.initialize_coords(x.coords_man)
+            else:
+                raise ValueError('Initialize input coordinates')
+        elif nrows != x.F.size(0):
+            raise ValueError('Input size does not match the coordinate size')

MinkowskiEngine/MinkowskiEngine/MinkowskiNonlinearity.py

@@ -0,0 +1,200 @@
+# Copyright (c) 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.
+from typing import Union
+
+import torch
+import torch.nn as nn
+
+from MinkowskiCommon import MinkowskiModuleBase
+from MinkowskiSparseTensor import SparseTensor
+from MinkowskiTensorField import TensorField
+
+
+class MinkowskiNonlinearityBase(MinkowskiModuleBase):
+    MODULE = None
+
+    def __init__(self, *args, **kwargs):
+        super(MinkowskiNonlinearityBase, self).__init__()
+        self.module = self.MODULE(*args, **kwargs)
+
+    def forward(self, input):
+        output = self.module(input.F)
+        if isinstance(input, TensorField):
+            return TensorField(
+                output,
+                coordinate_field_map_key=input.coordinate_field_map_key,
+                coordinate_manager=input.coordinate_manager,
+                quantization_mode=input.quantization_mode,
+            )
+        else:
+            return SparseTensor(
+                output,
+                coordinate_map_key=input.coordinate_map_key,
+                coordinate_manager=input.coordinate_manager,
+            )
+
+    def __repr__(self):
+        return self.__class__.__name__ + "()"
+
+
+class MinkowskiELU(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.ELU
+
+
+class MinkowskiHardshrink(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.Hardshrink
+
+
+class MinkowskiHardsigmoid(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.Hardsigmoid
+
+
+class MinkowskiHardtanh(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.Hardtanh
+
+
+class MinkowskiHardswish(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.Hardswish
+
+
+class MinkowskiLeakyReLU(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.LeakyReLU
+
+
+class MinkowskiLogSigmoid(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.LogSigmoid
+
+
+class MinkowskiPReLU(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.PReLU
+
+
+class MinkowskiReLU(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.ReLU
+
+
+class MinkowskiReLU6(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.ReLU6
+
+
+class MinkowskiRReLU(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.RReLU
+
+
+class MinkowskiSELU(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.SELU
+
+
+class MinkowskiCELU(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.CELU
+
+
+class MinkowskiGELU(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.GELU
+
+
+class MinkowskiSigmoid(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.Sigmoid
+
+
+class MinkowskiSiLU(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.SiLU
+
+
+class MinkowskiSoftplus(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.Softplus
+
+
+class MinkowskiSoftshrink(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.Softshrink
+
+
+class MinkowskiSoftsign(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.Softsign
+
+
+class MinkowskiTanh(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.Tanh
+
+
+class MinkowskiTanhshrink(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.Tanhshrink
+
+
+class MinkowskiThreshold(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.Threshold
+
+
+# Non-linear Activations (other)
+class MinkowskiSoftmin(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.Softmin
+
+
+class MinkowskiSoftmax(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.Softmax
+
+
+class MinkowskiLogSoftmax(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.LogSoftmax
+
+
+class MinkowskiAdaptiveLogSoftmaxWithLoss(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.AdaptiveLogSoftmaxWithLoss
+
+
+# Dropouts
+class MinkowskiDropout(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.Dropout
+
+
+class MinkowskiAlphaDropout(MinkowskiNonlinearityBase):
+    MODULE = torch.nn.AlphaDropout
+
+
+class MinkowskiSinusoidal(MinkowskiModuleBase):
+    def __init__(self, in_channel, out_channel):
+        MinkowskiModuleBase.__init__(self)
+        self.in_channel = in_channel
+        self.out_channel = out_channel
+        self.kernel = nn.Parameter(torch.rand(in_channel, out_channel))
+        self.bias = nn.Parameter(torch.rand(1, out_channel))
+        self.coef = nn.Parameter(torch.rand(1, out_channel))
+
+    def forward(self, input: Union[SparseTensor, TensorField]):
+
+        out_F = torch.sin(input.F.mm(self.kernel) + self.bias) * self.coef
+
+        if isinstance(input, TensorField):
+            return TensorField(
+                out_F,
+                coordinate_field_map_key=input.coordinate_field_map_key,
+                coordinate_manager=input.coordinate_manager,
+                quantization_mode=input.quantization_mode,
+            )
+        else:
+            return SparseTensor(
+                out_F,
+                coordinate_map_key=input.coordinate_map_key,
+                coordinate_manager=input.coordinate_manager,
+            )

MinkowskiEngine/MinkowskiEngine/MinkowskiNormalization.py

@@ -0,0 +1,399 @@
+# 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.
+import torch
+import torch.nn as nn
+from torch.nn import Module
+from torch.autograd import Function
+
+from MinkowskiSparseTensor import SparseTensor
+from MinkowskiTensorField import TensorField
+
+from MinkowskiPooling import MinkowskiGlobalAvgPooling
+from MinkowskiBroadcast import (
+    MinkowskiBroadcastAddition,
+    MinkowskiBroadcastMultiplication,
+)
+from MinkowskiEngineBackend._C import (
+    CoordinateMapKey,
+    BroadcastMode,
+    PoolingMode,
+)
+from MinkowskiCoordinateManager import CoordinateManager
+
+from MinkowskiCommon import (
+    MinkowskiModuleBase,
+    get_minkowski_function,
+)
+
+
+class MinkowskiBatchNorm(Module):
+    r"""A batch normalization layer for a sparse tensor.
+
+    See the pytorch :attr:`torch.nn.BatchNorm1d` for more details.
+    """
+
+    def __init__(
+        self,
+        num_features,
+        eps=1e-5,
+        momentum=0.1,
+        affine=True,
+        track_running_stats=True,
+    ):
+        super(MinkowskiBatchNorm, self).__init__()
+        self.bn = torch.nn.BatchNorm1d(
+            num_features,
+            eps=eps,
+            momentum=momentum,
+            affine=affine,
+            track_running_stats=track_running_stats,
+        )
+
+    def forward(self, input):
+        output = self.bn(input.F)
+        if isinstance(input, TensorField):
+            return TensorField(
+                output,
+                coordinate_field_map_key=input.coordinate_field_map_key,
+                coordinate_manager=input.coordinate_manager,
+                quantization_mode=input.quantization_mode,
+            )
+        else:
+            return SparseTensor(
+                output,
+                coordinate_map_key=input.coordinate_map_key,
+                coordinate_manager=input.coordinate_manager,
+            )
+
+    def __repr__(self):
+        s = "({}, eps={}, momentum={}, affine={}, track_running_stats={})".format(
+            self.bn.num_features,
+            self.bn.eps,
+            self.bn.momentum,
+            self.bn.affine,
+            self.bn.track_running_stats,
+        )
+        return self.__class__.__name__ + s
+
+
+class MinkowskiSyncBatchNorm(MinkowskiBatchNorm):
+    r"""A batch normalization layer with multi GPU synchronization."""
+
+    def __init__(
+        self,
+        num_features,
+        eps=1e-5,
+        momentum=0.1,
+        affine=True,
+        track_running_stats=True,
+        process_group=None,
+    ):
+        Module.__init__(self)
+        self.bn = torch.nn.SyncBatchNorm(
+            num_features,
+            eps=eps,
+            momentum=momentum,
+            affine=affine,
+            track_running_stats=track_running_stats,
+            process_group=process_group,
+        )
+
+    def forward(self, input):
+        output = self.bn(input.F)
+        if isinstance(input, TensorField):
+            return TensorField(
+                output,
+                coordinate_field_map_key=input.coordinate_field_map_key,
+                coordinate_manager=input.coordinate_manager,
+                quantization_mode=input.quantization_mode,
+            )
+        else:
+            return SparseTensor(
+                output,
+                coordinate_map_key=input.coordinate_map_key,
+                coordinate_manager=input.coordinate_manager,
+            )
+
+    @classmethod
+    def convert_sync_batchnorm(cls, module, process_group=None):
+        r"""Helper function to convert
+        :attr:`MinkowskiEngine.MinkowskiBatchNorm` layer in the model to
+        :attr:`MinkowskiEngine.MinkowskiSyncBatchNorm` layer.
+
+        Args:
+            module (nn.Module): containing module
+            process_group (optional): process group to scope synchronization,
+            default is the whole world
+
+        Returns:
+            The original module with the converted
+            :attr:`MinkowskiEngine.MinkowskiSyncBatchNorm` layer
+
+        Example::
+
+            >>> # Network with MinkowskiBatchNorm layer
+            >>> module = torch.nn.Sequential(
+            >>>            MinkowskiLinear(20, 100),
+            >>>            MinkowskiBatchNorm1d(100)
+            >>>          ).cuda()
+            >>> # creating process group (optional)
+            >>> # process_ids is a list of int identifying rank ids.
+            >>> process_group = torch.distributed.new_group(process_ids)
+            >>> sync_bn_module = convert_sync_batchnorm(module, process_group)
+
+        """
+        module_output = module
+        if isinstance(module, MinkowskiBatchNorm):
+            module_output = MinkowskiSyncBatchNorm(
+                module.bn.num_features,
+                module.bn.eps,
+                module.bn.momentum,
+                module.bn.affine,
+                module.bn.track_running_stats,
+                process_group,
+            )
+            if module.bn.affine:
+                with torch.no_grad():
+                    module_output.bn.weight = module.bn.weight
+                    module_output.bn.bias = module.bn.bias
+            module_output.bn.running_mean = module.bn.running_mean
+            module_output.bn.running_var = module.bn.running_var
+            module_output.bn.num_batches_tracked = module.bn.num_batches_tracked
+            if hasattr(module, "qconfig"):
+                module_output.bn.qconfig = module.bn.qconfig
+        for name, child in module.named_children():
+            module_output.add_module(
+                name, cls.convert_sync_batchnorm(child, process_group)
+            )
+        del module
+        return module_output
+
+
+class MinkowskiInstanceNormFunction(Function):
+    @staticmethod
+    def forward(
+        ctx,
+        in_feat: torch.Tensor,
+        in_coords_key: CoordinateMapKey,
+        glob_coords_key: CoordinateMapKey = None,
+        coords_manager: CoordinateManager = None,
+        gpooling_mode=PoolingMode.GLOBAL_AVG_POOLING_KERNEL,
+    ):
+        if glob_coords_key is None:
+            glob_coords_key = CoordinateMapKey(in_coords_key.get_coordinate_size())
+
+        gpool_avg_forward = get_minkowski_function("GlobalPoolingForward", in_feat)
+        broadcast_forward = get_minkowski_function("BroadcastForward", in_feat)
+
+        mean, num_nonzero = gpool_avg_forward(
+            in_feat,
+            gpooling_mode,
+            in_coords_key,
+            glob_coords_key,
+            coords_manager._manager,
+        )
+
+        # X - \mu
+        centered_feat = broadcast_forward(
+            in_feat,
+            -mean,
+            BroadcastMode.ELEMENTWISE_ADDITON,
+            in_coords_key,
+            glob_coords_key,
+            coords_manager._manager,
+        )
+
+        # Variance = 1/N \sum (X - \mu) ** 2
+        variance, num_nonzero = gpool_avg_forward(
+            centered_feat ** 2,
+            gpooling_mode,
+            in_coords_key,
+            glob_coords_key,
+            coords_manager._manager,
+        )
+
+        # norm_feat = (X - \mu) / \sigma
+        inv_std = 1 / (variance + 1e-8).sqrt()
+        norm_feat = broadcast_forward(
+            centered_feat,
+            inv_std,
+            BroadcastMode.ELEMENTWISE_MULTIPLICATION,
+            in_coords_key,
+            glob_coords_key,
+            coords_manager._manager,
+        )
+
+        ctx.saved_vars = (in_coords_key, glob_coords_key, coords_manager, gpooling_mode)
+        # For GPU tensors, must use save_for_backward.
+        ctx.save_for_backward(inv_std, norm_feat)
+        return norm_feat
+
+    @staticmethod
+    def backward(ctx, out_grad):
+        # https://kevinzakka.github.io/2016/09/14/batch_normalization/
+        in_coords_key, glob_coords_key, coords_manager, gpooling_mode = ctx.saved_vars
+
+        # To prevent the memory leakage, compute the norm again
+        inv_std, norm_feat = ctx.saved_tensors
+
+        gpool_avg_forward = get_minkowski_function("GlobalPoolingForward", out_grad)
+        broadcast_forward = get_minkowski_function("BroadcastForward", out_grad)
+
+        # 1/N \sum dout
+        mean_dout, num_nonzero = gpool_avg_forward(
+            out_grad,
+            gpooling_mode,
+            in_coords_key,
+            glob_coords_key,
+            coords_manager._manager,
+        )
+
+        # 1/N \sum (dout * out)
+        mean_dout_feat, num_nonzero = gpool_avg_forward(
+            out_grad * norm_feat,
+            gpooling_mode,
+            in_coords_key,
+            glob_coords_key,
+            coords_manager._manager,
+        )
+
+        # out * 1/N \sum (dout * out)
+        feat_mean_dout_feat = broadcast_forward(
+            norm_feat,
+            mean_dout_feat,
+            BroadcastMode.ELEMENTWISE_MULTIPLICATION,
+            in_coords_key,
+            glob_coords_key,
+            coords_manager._manager,
+        )
+
+        unnorm_din = broadcast_forward(
+            out_grad - feat_mean_dout_feat,
+            -mean_dout,
+            BroadcastMode.ELEMENTWISE_ADDITON,
+            in_coords_key,
+            glob_coords_key,
+            coords_manager._manager,
+        )
+
+        norm_din = broadcast_forward(
+            unnorm_din,
+            inv_std,
+            BroadcastMode.ELEMENTWISE_MULTIPLICATION,
+            in_coords_key,
+            glob_coords_key,
+            coords_manager._manager,
+        )
+
+        return norm_din, None, None, None, None
+
+
+class MinkowskiStableInstanceNorm(MinkowskiModuleBase):
+    def __init__(self, num_features):
+        Module.__init__(self)
+        self.num_features = num_features
+        self.eps = 1e-6
+        self.weight = nn.Parameter(torch.ones(1, num_features))
+        self.bias = nn.Parameter(torch.zeros(1, num_features))
+
+        self.mean_in = MinkowskiGlobalAvgPooling()
+        self.glob_sum = MinkowskiBroadcastAddition()
+        self.glob_sum2 = MinkowskiBroadcastAddition()
+        self.glob_mean = MinkowskiGlobalAvgPooling()
+        self.glob_times = MinkowskiBroadcastMultiplication()
+        self.reset_parameters()
+
+    def __repr__(self):
+        s = f"(nchannels={self.num_features})"
+        return self.__class__.__name__ + s
+
+    def reset_parameters(self):
+        self.weight.data.fill_(1)
+        self.bias.data.zero_()
+
+    def forward(self, x):
+        neg_mean_in = self.mean_in(
+            SparseTensor(-x.F, coords_key=x.coords_key, coords_manager=x.coords_man)
+        )
+        centered_in = self.glob_sum(x, neg_mean_in)
+        temp = SparseTensor(
+            centered_in.F ** 2,
+            coordinate_map_key=centered_in.coordinate_map_key,
+            coordinate_manager=centered_in.coordinate_manager,
+        )
+        var_in = self.glob_mean(temp)
+        instd_in = SparseTensor(
+            1 / (var_in.F + self.eps).sqrt(),
+            coordinate_map_key=var_in.coordinate_map_key,
+            coordinate_manager=var_in.coordinate_manager,
+        )
+
+        x = self.glob_times(self.glob_sum2(x, neg_mean_in), instd_in)
+        return SparseTensor(
+            x.F * self.weight + self.bias,
+            coordinate_map_key=x.coordinate_map_key,
+            coordinate_manager=x.coordinate_manager,
+        )
+
+
+class MinkowskiInstanceNorm(MinkowskiModuleBase):
+    r"""A instance normalization layer for a sparse tensor."""
+
+    def __init__(self, num_features):
+        r"""
+        Args:
+
+            num_features (int): the dimension of the input feautres.
+
+            mode (GlobalPoolingModel, optional): The internal global pooling computation mode.
+        """
+        Module.__init__(self)
+        self.num_features = num_features
+        self.weight = nn.Parameter(torch.ones(1, num_features))
+        self.bias = nn.Parameter(torch.zeros(1, num_features))
+        self.reset_parameters()
+        self.inst_norm = MinkowskiInstanceNormFunction()
+
+    def __repr__(self):
+        s = f"(nchannels={self.num_features})"
+        return self.__class__.__name__ + s
+
+    def reset_parameters(self):
+        self.weight.data.fill_(1)
+        self.bias.data.zero_()
+
+    def forward(self, input: SparseTensor):
+        assert isinstance(input, SparseTensor)
+
+        output = self.inst_norm.apply(
+            input.F, input.coordinate_map_key, None, input.coordinate_manager
+        )
+        output = output * self.weight + self.bias
+
+        return SparseTensor(
+            output,
+            coordinate_map_key=input.coordinate_map_key,
+            coordinate_manager=input.coordinate_manager,
+        )

MinkowskiEngine/MinkowskiEngine/MinkowskiOps.py

@@ -0,0 +1,497 @@
+# Copyright (c) 2021 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.
+from typing import Union
+import numpy as np
+
+import torch
+from torch.nn.modules import Module
+from MinkowskiSparseTensor import SparseTensor
+from MinkowskiTensor import (
+    COORDINATE_MANAGER_DIFFERENT_ERROR,
+    COORDINATE_KEY_DIFFERENT_ERROR,
+)
+from MinkowskiTensorField import TensorField
+from MinkowskiCommon import MinkowskiModuleBase
+from MinkowskiEngineBackend._C import CoordinateMapKey
+
+
+class MinkowskiLinear(Module):
+    def __init__(self, in_features, out_features, bias=True):
+        super(MinkowskiLinear, self).__init__()
+        self.linear = torch.nn.Linear(in_features, out_features, bias=bias)
+
+    def forward(self, input: Union[SparseTensor, TensorField]):
+        output = self.linear(input.F)
+        if isinstance(input, TensorField):
+            return TensorField(
+                output,
+                coordinate_field_map_key=input.coordinate_field_map_key,
+                coordinate_manager=input.coordinate_manager,
+                quantization_mode=input.quantization_mode,
+            )
+        else:
+            return SparseTensor(
+                output,
+                coordinate_map_key=input.coordinate_map_key,
+                coordinate_manager=input.coordinate_manager,
+            )
+
+    def __repr__(self):
+        s = "(in_features={}, out_features={}, bias={})".format(
+            self.linear.in_features,
+            self.linear.out_features,
+            self.linear.bias is not None,
+        )
+        return self.__class__.__name__ + s
+
+
+def _tuple_operator(*sparse_tensors, operator):
+    if len(sparse_tensors) == 1:
+        assert isinstance(sparse_tensors[0], (tuple, list))
+        sparse_tensors = sparse_tensors[0]
+
+    assert (
+        len(sparse_tensors) > 1
+    ), f"Invalid number of inputs. The input must be at least two len(sparse_tensors) > 1"
+
+    if isinstance(sparse_tensors[0], SparseTensor):
+        device = sparse_tensors[0].device
+        coordinate_manager = sparse_tensors[0].coordinate_manager
+        coordinate_map_key = sparse_tensors[0].coordinate_map_key
+        for s in sparse_tensors:
+            assert isinstance(
+                s, SparseTensor
+            ), "Inputs must be either SparseTensors or TensorFields."
+            assert (
+                device == s.device
+            ), f"Device must be the same. {device} != {s.device}"
+            assert (
+                coordinate_manager == s.coordinate_manager
+            ), COORDINATE_MANAGER_DIFFERENT_ERROR
+            assert coordinate_map_key == s.coordinate_map_key, (
+                COORDINATE_KEY_DIFFERENT_ERROR
+                + str(coordinate_map_key)
+                + " != "
+                + str(s.coordinate_map_key)
+            )
+        tens = []
+        for s in sparse_tensors:
+            tens.append(s.F)
+        return SparseTensor(
+            operator(tens),
+            coordinate_map_key=coordinate_map_key,
+            coordinate_manager=coordinate_manager,
+        )
+    elif isinstance(sparse_tensors[0], TensorField):
+        device = sparse_tensors[0].device
+        coordinate_manager = sparse_tensors[0].coordinate_manager
+        coordinate_field_map_key = sparse_tensors[0].coordinate_field_map_key
+        for s in sparse_tensors:
+            assert isinstance(
+                s, TensorField
+            ), "Inputs must be either SparseTensors or TensorFields."
+            assert (
+                device == s.device
+            ), f"Device must be the same. {device} != {s.device}"
+            assert (
+                coordinate_manager == s.coordinate_manager
+            ), COORDINATE_MANAGER_DIFFERENT_ERROR
+            assert coordinate_field_map_key == s.coordinate_field_map_key, (
+                COORDINATE_KEY_DIFFERENT_ERROR
+                + str(coordinate_field_map_key)
+                + " != "
+                + str(s.coordinate_field_map_key)
+            )
+        tens = []
+        for s in sparse_tensors:
+            tens.append(s.F)
+        return TensorField(
+            operator(tens),
+            coordinate_field_map_key=coordinate_field_map_key,
+            coordinate_manager=coordinate_manager,
+        )
+    else:
+        raise ValueError(
+            "Invalid data type. The input must be either a list of sparse tensors or a list of tensor fields."
+        )
+
+
+def cat(*sparse_tensors):
+    r"""Concatenate sparse tensors
+
+    Concatenate sparse tensor features. All sparse tensors must have the same
+    `coordinate_map_key` (the same coordinates). To concatenate sparse tensors
+    with different sparsity patterns, use SparseTensor binary operations, or
+    :attr:`MinkowskiEngine.MinkowskiUnion`.
+
+    Example::
+
+       >>> import MinkowskiEngine as ME
+       >>> sin = ME.SparseTensor(feats, coords)
+       >>> sin2 = ME.SparseTensor(feats2, coordinate_map_key=sin.coordinate_map_key, coordinate_mananger=sin.coordinate_manager)
+       >>> sout = UNet(sin)  # Returns an output sparse tensor on the same coordinates
+       >>> sout2 = ME.cat(sin, sin2, sout)  # Can concatenate multiple sparse tensors
+
+    """
+    return _tuple_operator(*sparse_tensors, operator=lambda xs: torch.cat(xs, dim=-1))
+
+
+def _sum(*sparse_tensors):
+    r"""Compute the sum of sparse tensor features
+
+    Sum all sparse tensor features. All sparse tensors must have the same
+    `coordinate_map_key` (the same coordinates). To sum sparse tensors with
+    different sparsity patterns, use SparseTensor binary operations, or
+    :attr:`MinkowskiEngine.MinkowskiUnion`.
+
+    Example::
+
+       >>> import MinkowskiEngine as ME
+       >>> sin = ME.SparseTensor(feats, coords)
+       >>> sin2 = ME.SparseTensor(feats2, coordinate_map_key=sin.coordinate_map_key, coordinate_manager=sin.coordinate_manager)
+       >>> sout = UNet(sin)  # Returns an output sparse tensor on the same coordinates
+       >>> sout2 = ME.sum(sin, sin2, sout)  # Can concatenate multiple sparse tensors
+
+    """
+
+    def return_sum(xs):
+        tmp = xs[0] + xs[1]
+        for x in xs[2:]:
+            tmp += x
+        return tmp
+
+    return _tuple_operator(*sparse_tensors, operator=lambda xs: return_sum(xs))
+
+
+def mean(*sparse_tensors):
+    r"""Compute the average of sparse tensor features
+
+    Sum all sparse tensor features. All sparse tensors must have the same
+    `coordinate_map_key` (the same coordinates). To sum sparse tensors with
+    different sparsity patterns, use SparseTensor binary operations, or
+    :attr:`MinkowskiEngine.MinkowskiUnion`.
+
+    Example::
+
+       >>> import MinkowskiEngine as ME
+       >>> sin = ME.SparseTensor(feats, coords)
+       >>> sin2 = ME.SparseTensor(feats2, coordinate_map_key=sin.coordinate_map_key, coordinate_manager=sin.coordinate_manager)
+       >>> sout = UNet(sin)  # Returns an output sparse tensor on the same coordinates
+       >>> sout2 = ME.mean(sin, sin2, sout)  # Can concatenate multiple sparse tensors
+
+    """
+
+    def return_mean(xs):
+        tmp = xs[0] + xs[1]
+        for x in xs[2:]:
+            tmp += x
+        return tmp / len(xs)
+
+    return _tuple_operator(*sparse_tensors, operator=lambda xs: return_mean(xs))
+
+
+def var(*sparse_tensors):
+    r"""Compute the variance of sparse tensor features
+
+    Sum all sparse tensor features. All sparse tensors must have the same
+    `coordinate_map_key` (the same coordinates). To sum sparse tensors with
+    different sparsity patterns, use SparseTensor binary operations, or
+    :attr:`MinkowskiEngine.MinkowskiUnion`.
+
+    Example::
+
+       >>> import MinkowskiEngine as ME
+       >>> sin = ME.SparseTensor(feats, coords)
+       >>> sin2 = ME.SparseTensor(feats2, coordinate_map_key=sin.coordinate_map_key, coordinate_manager=sin.coordinate_manager)
+       >>> sout = UNet(sin)  # Returns an output sparse tensor on the same coordinates
+       >>> sout2 = ME.var(sin, sin2, sout)  # Can concatenate multiple sparse tensors
+
+    """
+
+    def return_var(xs):
+        tmp = xs[0] + xs[1]
+        for x in xs[2:]:
+            tmp += x
+        mean = tmp / len(xs)
+        var = (xs[0] - mean) ** 2
+        for x in xs[1:]:
+            var += (x - mean) ** 2
+        return var / len(xs)
+
+    return _tuple_operator(*sparse_tensors, operator=lambda xs: return_var(xs))
+
+
+def dense_coordinates(shape: Union[list, torch.Size]):
+    """
+    coordinates = dense_coordinates(tensor.shape)
+    """
+    r"""
+    Assume the input to have BxCxD1xD2x....xDN format.
+
+    If the shape of the tensor do not change, use 
+    """
+    spatial_dim = len(shape) - 2
+    assert (
+        spatial_dim > 0
+    ), "Invalid shape. Shape must be batch x channel x spatial dimensions."
+
+    # Generate coordinates
+    size = [i for i in shape]
+    B = size[0]
+    coordinates = torch.from_numpy(
+        np.stack(
+            [
+                s.reshape(-1)
+                for s in np.meshgrid(
+                    np.linspace(0, B - 1, B),
+                    *(np.linspace(0, s - 1, s) for s in size[2:]),
+                    indexing="ij",
+                )
+            ],
+            1,
+        )
+    ).int()
+    return coordinates
+
+
+def to_sparse(x: torch.Tensor, format: str = None, coordinates=None, device=None):
+    r"""Convert a batched tensor (dimension 0 is the batch dimension) to a SparseTensor
+
+    :attr:`x` (:attr:`torch.Tensor`): a batched tensor. The first dimension is the batch dimension.
+
+    :attr:`format` (:attr:`str`): Format of the tensor. It must include 'B' and 'C' indicating the batch and channel dimension respectively. The rest of the dimensions must be 'X'. .e.g. format="BCXX" if image data with BCHW format is used. If a 3D data with the channel at the last dimension, use format="BXXXC" indicating Batch X Height X Width X Depth X Channel. If not provided, the format will be "BCX...X".
+
+    :attr:`device`: Device the sparse tensor will be generated on. If not provided, the device of the input tensor will be used.
+
+    """
+    assert x.ndim > 2, "Input has 0 spatial dimension."
+    assert isinstance(x, torch.Tensor)
+    if format is None:
+        format = [
+            "X",
+        ] * x.ndim
+        format[0] = "B"
+        format[1] = "C"
+        format = "".join(format)
+    assert x.ndim == len(format), f"Invalid format: {format}. len(format) != x.ndim"
+    assert (
+        "B" in format and "B" == format[0] and format.count("B") == 1
+    ), "The input must have the batch axis and the format must include 'B' indicating the batch axis."
+    assert (
+        "C" in format and format.count("C") == 1
+    ), "The format must indicate the channel axis"
+    if device is None:
+        device = x.device
+    ch_dim = format.find("C")
+    reduced_x = torch.abs(x).sum(ch_dim)
+    bcoords = torch.where(reduced_x != 0)
+    stacked_bcoords = torch.stack(bcoords, dim=1).int()
+    indexing = [f"bcoords[{i}]" for i in range(len(bcoords))]
+    indexing.insert(ch_dim, ":")
+    features = torch.zeros(
+        (len(stacked_bcoords), x.size(ch_dim)), dtype=x.dtype, device=x.device
+    )
+    exec("features[:] = x[" + ", ".join(indexing) + "]")
+    return SparseTensor(features=features, coordinates=stacked_bcoords, device=device)
+
+
+def to_sparse_all(dense_tensor: torch.Tensor, coordinates: torch.Tensor = None):
+    r"""Converts a (differentiable) dense tensor to a sparse tensor with all coordinates.
+
+    Assume the input to have BxCxD1xD2x....xDN format.
+
+    If the shape of the tensor do not change, use `dense_coordinates` to cache the coordinates.
+    Please refer to tests/python/dense.py for usage
+
+    Example::
+
+       >>> dense_tensor = torch.rand(3, 4, 5, 6, 7, 8)  # BxCxD1xD2xD3xD4
+       >>> dense_tensor.requires_grad = True
+       >>> stensor = to_sparse(dense_tensor)
+
+    """
+    spatial_dim = dense_tensor.ndim - 2
+    assert (
+        spatial_dim > 0
+    ), "Invalid shape. Shape must be batch x channel x spatial dimensions."
+
+    if coordinates is None:
+        coordinates = dense_coordinates(dense_tensor.shape)
+
+    feat_tensor = dense_tensor.permute(0, *(2 + i for i in range(spatial_dim)), 1)
+    return SparseTensor(
+        feat_tensor.reshape(-1, dense_tensor.size(1)),
+        coordinates,
+        device=dense_tensor.device,
+    )
+
+
+class MinkowskiToSparseTensor(MinkowskiModuleBase):
+    r"""Converts a (differentiable) dense tensor or a :attr:`MinkowskiEngine.TensorField` to a :attr:`MinkowskiEngine.SparseTensor`.
+
+    For dense tensor, the input must have the BxCxD1xD2x....xDN format.
+
+    :attr:`remove_zeros` (bool): if True, removes zero valued coordinates. If
+    False, use all coordinates to populate a sparse tensor. True by default.
+
+    :attr:`coordinates` (torch.Tensor): if set, use the provided coordinates
+    only for sparse tensor generation. Will ignore `remove_zeros`.
+    
+    If the shape of the tensor do not change, use `dense_coordinates` to cache the coordinates.
+    Please refer to tests/python/dense.py for usage.
+
+    Example::
+
+       >>> # Differentiable dense torch.Tensor to sparse tensor.
+       >>> dense_tensor = torch.rand(3, 4, 11, 11, 11, 11)  # BxCxD1xD2x....xDN
+       >>> dense_tensor.requires_grad = True
+
+       >>> # Since the shape is fixed, cache the coordinates for faster inference
+       >>> coordinates = dense_coordinates(dense_tensor.shape)
+
+       >>> network = nn.Sequential(
+       >>>     # Add layers that can be applied on a regular pytorch tensor
+       >>>     nn.ReLU(),
+       >>>     MinkowskiToSparseTensor(coordinates=coordinates),
+       >>>     MinkowskiConvolution(4, 5, kernel_size=3, dimension=4),
+       >>>     MinkowskiBatchNorm(5),
+       >>>     MinkowskiReLU(),
+       >>> )
+
+       >>> for i in range(5):
+       >>>   print(f"Iteration: {i}")
+       >>>   soutput = network(dense_tensor)
+       >>>   soutput.F.sum().backward()
+       >>>   soutput.dense(shape=dense_tensor.shape)
+
+    """
+
+    def __init__(self, remove_zeros=True, coordinates: torch.Tensor = None):
+        MinkowskiModuleBase.__init__(self)
+        self.remove_zeros = remove_zeros
+        self.coordinates = coordinates
+
+    def forward(self, input: Union[TensorField, torch.Tensor]):
+        if isinstance(input, TensorField):
+            return input.sparse()
+        elif isinstance(input, torch.Tensor):
+            # dense tensor to sparse tensor conversion
+            if self.remove_zeros and self.coordinates is not None:
+                return to_sparse(input)
+            else:
+                return to_sparse_all(input, self.coordinates)
+        else:
+            raise ValueError(
+                "Unsupported type. Only TensorField and torch.Tensor are supported"
+            )
+
+    def __repr__(self):
+        return self.__class__.__name__ + "()"
+
+
+class MinkowskiToDenseTensor(MinkowskiModuleBase):
+    r"""Converts a (differentiable) sparse tensor to a torch tensor.
+
+    The return type has the BxCxD1xD2x....xDN format.
+
+    Example::
+
+       >>> dense_tensor = torch.rand(3, 4, 11, 11, 11, 11)  # BxCxD1xD2x....xDN
+       >>> dense_tensor.requires_grad = True
+
+       >>> # Since the shape is fixed, cache the coordinates for faster inference
+       >>> coordinates = dense_coordinates(dense_tensor.shape)
+
+       >>> network = nn.Sequential(
+       >>>     # Add layers that can be applied on a regular pytorch tensor
+       >>>     nn.ReLU(),
+       >>>     MinkowskiToSparseTensor(coordinates=coordinates),
+       >>>     MinkowskiConvolution(4, 5, stride=2, kernel_size=3, dimension=4),
+       >>>     MinkowskiBatchNorm(5),
+       >>>     MinkowskiReLU(),
+       >>>     MinkowskiConvolutionTranspose(5, 6, stride=2, kernel_size=3, dimension=4),
+       >>>     MinkowskiToDenseTensor(
+       >>>         dense_tensor.shape
+       >>>     ),  # must have the same tensor stride.
+       >>> )
+
+       >>> for i in range(5):
+       >>>     print(f"Iteration: {i}")
+       >>>     output = network(dense_tensor) # returns a regular pytorch tensor
+       >>>     output.sum().backward()
+
+    """
+
+    def __init__(self, shape: torch.Size = None):
+        MinkowskiModuleBase.__init__(self)
+        self.shape = shape
+
+    def forward(self, input: SparseTensor):
+        # dense tensor to sparse tensor conversion
+        dense_tensor, _, _ = input.dense(shape=self.shape)
+        return dense_tensor
+
+    def __repr__(self):
+        return self.__class__.__name__ + "()"
+
+
+class MinkowskiToFeature(MinkowskiModuleBase):
+    r"""
+    Extract features from a sparse tensor and returns a pytorch tensor.
+
+    Can be used to to make a network construction simpler.
+
+    Example::
+
+       >>> net = nn.Sequential(MinkowskiConvolution(...), MinkowskiGlobalMaxPooling(...), MinkowskiToFeature(), nn.Linear(...))
+       >>> torch_tensor = net(sparse_tensor)
+
+    """
+
+    def forward(self, x: SparseTensor):
+        assert isinstance(
+            x, (SparseTensor, TensorField)
+        ), "Invalid input type for MinkowskiToFeature"
+        return x.F
+
+
+class MinkowskiStackCat(torch.nn.Sequential):
+    def forward(self, x):
+        return cat([module(x) for module in self])
+
+
+class MinkowskiStackSum(torch.nn.Sequential):
+    def forward(self, x):
+        return _sum([module(x) for module in self])
+
+
+class MinkowskiStackMean(torch.nn.Sequential):
+    def forward(self, x):
+        return mean([module(x) for module in self])
+
+
+class MinkowskiStackVar(torch.nn.Sequential):
+    def forward(self, x):
+        return var([module(x) for module in self])

MinkowskiEngine/MinkowskiEngine/MinkowskiPooling.py

@@ -0,0 +1,780 @@
+# 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.
+from typing import Union
+
+import torch
+from torch.autograd import Function
+
+import MinkowskiEngineBackend._C as _C
+from MinkowskiEngineBackend._C import CoordinateMapKey, PoolingMode
+from MinkowskiSparseTensor import SparseTensor, _get_coordinate_map_key
+from MinkowskiCoordinateManager import CoordinateManager
+from MinkowskiKernelGenerator import KernelGenerator, save_ctx
+from MinkowskiCommon import (
+    MinkowskiModuleBase,
+    get_minkowski_function,
+)
+import MinkowskiEngine as ME
+
+
+class MinkowskiLocalPoolingFunction(Function):
+    @staticmethod
+    def forward(
+        ctx,
+        input_features: torch.Tensor,
+        pooling_mode: PoolingMode,
+        kernel_generator: KernelGenerator,
+        in_coordinate_map_key: CoordinateMapKey,
+        out_coordinate_map_key: CoordinateMapKey = None,
+        coordinate_manager: CoordinateManager = None,
+    ):
+        if out_coordinate_map_key is None:
+            out_coordinate_map_key = CoordinateMapKey(
+                in_coordinate_map_key.get_coordinate_size()
+            )
+
+        input_features = input_features.contiguous()
+        ctx.input_features = input_features
+        ctx = save_ctx(
+            ctx,
+            kernel_generator,
+            in_coordinate_map_key,
+            out_coordinate_map_key,
+            coordinate_manager,
+        )
+        ctx.pooling_mode = pooling_mode
+
+        fw_fn = get_minkowski_function("LocalPoolingForward", input_features)
+        out_feat, num_nonzero = fw_fn(
+            ctx.input_features,
+            kernel_generator.kernel_size,
+            kernel_generator.kernel_stride,
+            kernel_generator.kernel_dilation,
+            kernel_generator.region_type,
+            kernel_generator.region_offsets,
+            pooling_mode,
+            ctx.in_coordinate_map_key,
+            ctx.out_coordinate_map_key,
+            ctx.coordinate_manager._manager,
+        )
+        ctx.num_nonzero = num_nonzero
+        return out_feat
+
+    @staticmethod
+    def backward(ctx, grad_out_feat):
+        grad_out_feat = grad_out_feat.contiguous()
+        bw_fn = get_minkowski_function("LocalPoolingBackward", grad_out_feat)
+        grad_in_feat = bw_fn(
+            ctx.input_features,
+            grad_out_feat,
+            ctx.num_nonzero,
+            ctx.kernel_generator.kernel_size,
+            ctx.kernel_generator.kernel_stride,
+            ctx.kernel_generator.kernel_dilation,
+            ctx.kernel_generator.region_type,
+            ctx.kernel_generator.region_offsets,
+            ctx.pooling_mode,
+            ctx.in_coordinate_map_key,
+            ctx.out_coordinate_map_key,
+            ctx.coordinate_manager._manager,
+        )
+        return (
+            grad_in_feat,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+class MinkowskiPoolingBase(MinkowskiModuleBase):
+
+    __slots__ = (
+        "is_transpose",
+        "kernel_generator",
+        "pooling_mode",
+        "dimension",
+        "pooling",
+    )
+
+    def __init__(
+        self,
+        kernel_size,
+        stride=1,
+        dilation=1,
+        kernel_generator=None,
+        is_transpose=False,
+        pooling_mode=PoolingMode.LOCAL_AVG_POOLING,
+        dimension=-1,
+    ):
+        super(MinkowskiPoolingBase, self).__init__()
+        assert (
+            dimension > 0
+        ), f"Invalid dimension. Please provide a valid dimension argument. dimension={dimension}"
+
+        if kernel_generator is None:
+            kernel_generator = KernelGenerator(
+                kernel_size=kernel_size,
+                stride=stride,
+                dilation=dilation,
+                dimension=dimension,
+            )
+
+        self.is_transpose = is_transpose
+        self.kernel_generator = kernel_generator
+        self.pooling_mode = pooling_mode
+        self.dimension = dimension
+        self.pooling = MinkowskiLocalPoolingFunction()
+
+    def forward(
+        self,
+        input: SparseTensor,
+        coordinates: Union[torch.IntTensor, CoordinateMapKey, SparseTensor] = None,
+    ):
+        r"""
+        :attr:`input` (`MinkowskiEngine.SparseTensor`): Input sparse tensor to apply a
+        convolution on.
+
+        :attr:`coordinates` ((`torch.IntTensor`, `MinkowskiEngine.CoordsKey`,
+        `MinkowskiEngine.SparseTensor`), optional): If provided, generate
+        results on the provided coordinates. None by default.
+
+        """
+        assert isinstance(input, SparseTensor)
+        assert input.D == self.dimension
+
+        # Get a new coordinate map key or extract one from the coordinates
+        out_coordinate_map_key = _get_coordinate_map_key(input, coordinates)
+        outfeat = self.pooling.apply(
+            input.F,
+            self.pooling_mode,
+            self.kernel_generator,
+            input.coordinate_map_key,
+            out_coordinate_map_key,
+            input._manager,
+        )
+
+        return SparseTensor(
+            outfeat,
+            coordinate_map_key=out_coordinate_map_key,
+            coordinate_manager=input.coordinate_manager,
+        )
+
+    def __repr__(self):
+        s = "(kernel_size={}, stride={}, dilation={})".format(
+            self.kernel_generator.kernel_size,
+            self.kernel_generator.kernel_stride,
+            self.kernel_generator.kernel_dilation,
+        )
+        return self.__class__.__name__ + s
+
+
+class MinkowskiAvgPooling(MinkowskiPoolingBase):
+    r"""Average input features within a kernel.
+
+    .. math::
+
+        \mathbf{y}_\mathbf{u} = \frac{1}{|\mathcal{N}^D(\mathbf{u},
+        \mathcal{C}^\text{in})|} \sum_{\mathbf{i} \in \mathcal{N}^D(\mathbf{u},
+        \mathcal{C}^\text{in})} \mathbf{x}_{\mathbf{u} + \mathbf{i}}
+        \; \text{for} \; \mathbf{u} \in \mathcal{C}^\text{out}
+
+    For each output :math:`\mathbf{u}` in :math:`\mathcal{C}^\text{out}`,
+    average input features.
+
+    .. note::
+
+        An average layer first computes the cardinality of the input features,
+        the number of input features for each output, and divide the sum of the
+        input features by the cardinality. For a dense tensor, the cardinality
+        is a constant, the volume of a kernel. However, for a sparse tensor, the
+        cardinality varies depending on the number of input features per output.
+        Thus, the average pooling for a sparse tensor is not equivalent to the
+        conventional average pooling layer for a dense tensor. Please refer to
+        the :attr:`MinkowskiSumPooling` for the equivalent layer.
+
+    .. note::
+
+       The engine will generate the in-out mapping corresponding to a
+       pooling function faster if the kernel sizes is equal to the stride
+       sizes, e.g. `kernel_size = [2, 1], stride = [2, 1]`.
+
+       If you use a U-network architecture, use the transposed version of
+       the same function for up-sampling. e.g. `pool =
+       MinkowskiSumPooling(kernel_size=2, stride=2, D=D)`, then use the
+       `unpool = MinkowskiPoolingTranspose(kernel_size=2, stride=2, D=D)`.
+
+    """
+
+    def __init__(
+        self,
+        kernel_size=-1,
+        stride=1,
+        dilation=1,
+        kernel_generator=None,
+        dimension=None,
+    ):
+        r"""a high-dimensional sparse average pooling layer.
+
+        Args:
+            :attr:`kernel_size` (int, optional): the size of the kernel in the
+            output tensor. If not provided, :attr:`region_offset` should be
+            :attr:`RegionType.CUSTOM` and :attr:`region_offset` should be a 2D
+            matrix with size :math:`N\times D` such that it lists all :math:`N`
+            offsets in D-dimension.
+
+            :attr:`stride` (int, or list, optional): stride size of the
+            convolution layer. If non-identity is used, the output coordinates
+            will be at least :attr:`stride` :math:`\times` :attr:`tensor_stride`
+            away. When a list is given, the length must be D; each element will
+            be used for stride size for the specific axis.
+
+            :attr:`dilation` (int, or list, optional): dilation size for the
+            convolution kernel. When a list is given, the length must be D and
+            each element is an axis specific dilation. All elements must be > 0.
+
+            :attr:`kernel_generator` (:attr:`MinkowskiEngine.KernelGenerator`,
+            optional): define custom kernel shape.
+
+            :attr:`dimension` (int): the spatial dimension of the space where
+            all the inputs and the network are defined. For example, images are
+            in a 2D space, meshes and 3D shapes are in a 3D space.
+
+        .. warning::
+
+           Custom kernel shapes are not supported when kernel_size == stride.
+
+        """
+        is_transpose = False
+        MinkowskiPoolingBase.__init__(
+            self,
+            kernel_size,
+            stride,
+            dilation,
+            kernel_generator,
+            is_transpose,
+            pooling_mode=PoolingMode.LOCAL_AVG_POOLING,
+            dimension=dimension,
+        )
+
+
+class MinkowskiSumPooling(MinkowskiPoolingBase):
+    r"""Sum all input features within a kernel.
+
+    .. math::
+
+        \mathbf{y}_\mathbf{u} = \sum_{\mathbf{i} \in \mathcal{N}^D(\mathbf{u},
+        \mathcal{C}^\text{in})} \mathbf{x}_{\mathbf{u} + \mathbf{i}}
+        \; \text{for} \; \mathbf{u} \in \mathcal{C}^\text{out}
+
+    For each output :math:`\mathbf{u}` in :math:`\mathcal{C}^\text{out}`,
+    average input features.
+
+    .. note::
+
+        An average layer first computes the cardinality of the input features,
+        the number of input features for each output, and divide the sum of the
+        input features by the cardinality. For a dense tensor, the cardinality
+        is a constant, the volume of a kernel. However, for a sparse tensor, the
+        cardinality varies depending on the number of input features per output.
+        Thus, averaging the input features with the cardinality may not be
+        equivalent to the conventional average pooling for a dense tensor.
+        This layer provides an alternative that does not divide the sum by the
+        cardinality.
+
+    .. note::
+
+       The engine will generate the in-out mapping corresponding to a
+       pooling function faster if the kernel sizes is equal to the stride
+       sizes, e.g. `kernel_size = [2, 1], stride = [2, 1]`.
+
+       If you use a U-network architecture, use the transposed version of
+       the same function for up-sampling. e.g. `pool =
+       MinkowskiSumPooling(kernel_size=2, stride=2, D=D)`, then use the
+       `unpool = MinkowskiPoolingTranspose(kernel_size=2, stride=2, D=D)`.
+
+
+    """
+
+    def __init__(
+        self, kernel_size, stride=1, dilation=1, kernel_generator=None, dimension=None
+    ):
+        r"""a high-dimensional sum pooling layer
+
+        Args:
+            :attr:`kernel_size` (int, optional): the size of the kernel in the
+            output tensor. If not provided, :attr:`region_offset` should be
+            :attr:`RegionType.CUSTOM` and :attr:`region_offset` should be a 2D
+            matrix with size :math:`N\times D` such that it lists all :math:`N`
+            offsets in D-dimension.
+
+            :attr:`stride` (int, or list, optional): stride size of the
+            convolution layer. If non-identity is used, the output coordinates
+            will be at least :attr:`stride` :math:`\times` :attr:`tensor_stride`
+            away. When a list is given, the length must be D; each element will
+            be used for stride size for the specific axis.
+
+            :attr:`dilation` (int, or list, optional): dilation size for the
+            convolution kernel. When a list is given, the length must be D and
+            each element is an axis specific dilation. All elements must be > 0.
+
+            :attr:`kernel_generator` (:attr:`MinkowskiEngine.KernelGenerator`,
+            optional): define custom kernel shape.
+
+            :attr:`dimension` (int): the spatial dimension of the space where
+            all the inputs and the network are defined. For example, images are
+            in a 2D space, meshes and 3D shapes are in a 3D space.
+
+        .. warning::
+
+           Custom kernel shapes are not supported when kernel_size == stride.
+
+        """
+        is_transpose = False
+        MinkowskiPoolingBase.__init__(
+            self,
+            kernel_size,
+            stride,
+            dilation,
+            kernel_generator,
+            is_transpose,
+            pooling_mode=PoolingMode.LOCAL_SUM_POOLING,
+            dimension=dimension,
+        )
+
+
+class MinkowskiMaxPooling(MinkowskiPoolingBase):
+    r"""A max pooling layer for a sparse tensor.
+
+    .. math::
+
+        y^c_\mathbf{u} = \max_{\mathbf{i} \in \mathcal{N}^D(\mathbf{u},
+        \mathcal{C}^\text{in})} x^c_{\mathbf{u} + \mathbf{i}} \; \text{for} \;
+        \mathbf{u} \in \mathcal{C}^\text{out}
+
+    where :math:`y^c_\mathbf{u}` is a feature at channel :math:`c` and a
+    coordinate :math:`\mathbf{u}`.
+
+    .. note::
+
+       The engine will generate the in-out mapping corresponding to a
+       pooling function faster if the kernel sizes is equal to the stride
+       sizes, e.g. `kernel_size = [2, 1], stride = [2, 1]`.
+
+       If you use a U-network architecture, use the transposed version of
+       the same function for up-sampling. e.g. `pool =
+       MinkowskiSumPooling(kernel_size=2, stride=2, D=D)`, then use the
+       `unpool = MinkowskiPoolingTranspose(kernel_size=2, stride=2, D=D)`.
+
+    """
+
+    def __init__(
+        self, kernel_size, stride=1, dilation=1, kernel_generator=None, dimension=None
+    ):
+        r"""a high-dimensional max pooling layer for sparse tensors.
+
+        Args:
+            :attr:`kernel_size` (int, optional): the size of the kernel in the
+            output tensor. If not provided, :attr:`region_offset` should be
+            :attr:`RegionType.CUSTOM` and :attr:`region_offset` should be a 2D
+            matrix with size :math:`N\times D` such that it lists all :math:`N`
+            offsets in D-dimension.
+
+            :attr:`stride` (int, or list, optional): stride size of the
+            convolution layer. If non-identity is used, the output coordinates
+            will be at least :attr:`stride` :math:`\times` :attr:`tensor_stride`
+            away. When a list is given, the length must be D; each element will
+            be used for stride size for the specific axis.
+
+            :attr:`dilation` (int, or list, optional): dilation size for the
+            convolution kernel. When a list is given, the length must be D and
+            each element is an axis specific dilation. All elements must be > 0.
+
+            :attr:`kernel_generator` (:attr:`MinkowskiEngine.KernelGenerator`,
+            optional): define custom kernel shape.
+
+            :attr:`dimension` (int): the spatial dimension of the space where
+            all the inputs and the network are defined. For example, images are
+            in a 2D space, meshes and 3D shapes are in a 3D space.
+
+        .. warning::
+
+           Custom kernel shapes are not supported when kernel_size == stride.
+
+        """
+
+        MinkowskiPoolingBase.__init__(
+            self,
+            kernel_size,
+            stride,
+            dilation,
+            kernel_generator,
+            is_transpose=False,
+            pooling_mode=PoolingMode.LOCAL_MAX_POOLING,
+            dimension=dimension,
+        )
+
+
+class MinkowskiLocalPoolingTransposeFunction(Function):
+    @staticmethod
+    def forward(
+        ctx,
+        input_features: torch.Tensor,
+        pooling_mode: PoolingMode,
+        kernel_generator: KernelGenerator,
+        in_coordinate_map_key: CoordinateMapKey,
+        out_coordinate_map_key: CoordinateMapKey = None,
+        coordinate_manager: CoordinateManager = None,
+    ):
+        if out_coordinate_map_key is None:
+            out_coordinate_map_key = CoordinateMapKey(
+                in_coordinate_map_key.get_coordinate_size()
+            )
+
+        input_features = input_features.contiguous()
+        ctx.input_features = input_features
+        ctx = save_ctx(
+            ctx,
+            kernel_generator,
+            in_coordinate_map_key,
+            out_coordinate_map_key,
+            coordinate_manager,
+        )
+        ctx.pooling_mode = pooling_mode
+
+        fw_fn = get_minkowski_function("LocalPoolingTransposeForward", input_features)
+        out_feat, num_nonzero = fw_fn(
+            ctx.input_features,
+            kernel_generator.kernel_size,
+            kernel_generator.kernel_stride,
+            kernel_generator.kernel_dilation,
+            kernel_generator.region_type,
+            kernel_generator.region_offsets,
+            kernel_generator.expand_coordinates,
+            pooling_mode,
+            ctx.in_coordinate_map_key,
+            ctx.out_coordinate_map_key,
+            ctx.coordinate_manager._manager,
+        )
+        ctx.num_nonzero = num_nonzero
+        return out_feat
+
+    @staticmethod
+    def backward(ctx, grad_out_feat):
+        grad_out_feat = grad_out_feat.contiguous()
+        bw_fn = get_minkowski_function("LocalPoolingTransposeBackward", grad_out_feat)
+        grad_in_feat = bw_fn(
+            ctx.input_features,
+            grad_out_feat,
+            ctx.num_nonzero,
+            ctx.kernel_generator.kernel_size,
+            ctx.kernel_generator.kernel_stride,
+            ctx.kernel_generator.kernel_dilation,
+            ctx.kernel_generator.region_type,
+            ctx.kernel_generator.region_offsets,
+            ctx.pooling_mode,
+            ctx.in_coordinate_map_key,
+            ctx.out_coordinate_map_key,
+            ctx.coordinate_manager._manager,
+        )
+        return (
+            grad_in_feat,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+
+
+class MinkowskiPoolingTranspose(MinkowskiPoolingBase):
+    r"""A pooling transpose layer for a sparse tensor.
+
+    Unpool the features and divide it by the number of non zero elements that
+    contributed.
+    """
+
+    def __init__(
+        self,
+        kernel_size,
+        stride,
+        dilation=1,
+        kernel_generator=None,
+        expand_coordinates=False,
+        dimension=None,
+    ):
+        r"""a high-dimensional unpooling layer for sparse tensors.
+
+        Args:
+            :attr:`kernel_size` (int, optional): the size of the kernel in the
+            output tensor. If not provided, :attr:`region_offset` should be
+            :attr:`RegionType.CUSTOM` and :attr:`region_offset` should be a 2D
+            matrix with size :math:`N\times D` such that it lists all :math:`N`
+            offsets in D-dimension.
+
+            :attr:`stride` (int, or list, optional): stride size of the
+            convolution layer. If non-identity is used, the output coordinates
+            will be at least :attr:`stride` :math:`\times` :attr:`tensor_stride`
+            away. When a list is given, the length must be D; each element will
+            be used for stride size for the specific axis.
+
+            :attr:`dilation` (int, or list, optional): dilation size for the
+            convolution kernel. When a list is given, the length must be D and
+            each element is an axis specific dilation. All elements must be > 0.
+
+            :attr:`kernel_generator` (:attr:`MinkowskiEngine.KernelGenerator`,
+            optional): define custom kernel shape.
+
+            :attr:`expand_coordinates` (bool, optional): Force generation of
+            new coordinates. When True, the output coordinates will be the
+            outer product of the kernel shape and the input coordinates.
+            `False` by default.
+
+            :attr:`dimension` (int): the spatial dimension of the space where
+            all the inputs and the network are defined. For example, images are
+            in a 2D space, meshes and 3D shapes are in a 3D space.
+
+        """
+        is_transpose = True
+        if kernel_generator is None:
+            kernel_generator = KernelGenerator(
+                kernel_size=kernel_size,
+                stride=stride,
+                dilation=dilation,
+                expand_coordinates=expand_coordinates,
+                dimension=dimension,
+            )
+
+        MinkowskiPoolingBase.__init__(
+            self,
+            kernel_size,
+            stride,
+            dilation,
+            kernel_generator,
+            is_transpose,
+            dimension=dimension,
+        )
+        self.pooling = MinkowskiLocalPoolingTransposeFunction()
+
+
+class MinkowskiGlobalPoolingFunction(Function):
+    @staticmethod
+    def forward(
+        ctx,
+        input_features: torch.Tensor,
+        pooling_mode: PoolingMode,
+        in_coordinate_map_key: CoordinateMapKey,
+        out_coordinate_map_key: CoordinateMapKey = None,
+        coordinate_manager: CoordinateManager = None,
+    ):
+        if out_coordinate_map_key is None:
+            out_coordinate_map_key = CoordinateMapKey(
+                in_coordinate_map_key.get_coordinate_size()
+            )
+        input_features = input_features.contiguous()
+
+        ctx.input_features = input_features
+        ctx.in_coords_key = in_coordinate_map_key
+        ctx.out_coords_key = out_coordinate_map_key
+        ctx.coordinate_manager = coordinate_manager
+        ctx.pooling_mode = pooling_mode
+
+        fw_fn = get_minkowski_function("GlobalPoolingForward", input_features)
+        out_feat, num_nonzero = fw_fn(
+            input_features,
+            pooling_mode,
+            ctx.in_coords_key,
+            ctx.out_coords_key,
+            ctx.coordinate_manager._manager,
+        )
+        ctx.num_nonzero = num_nonzero
+        return out_feat
+
+    @staticmethod
+    def backward(ctx, grad_out_feat):
+        grad_out_feat = grad_out_feat.contiguous()
+        bw_fn = get_minkowski_function("GlobalPoolingBackward", grad_out_feat)
+        grad_in_feat = bw_fn(
+            ctx.input_features,
+            grad_out_feat,
+            ctx.num_nonzero,
+            ctx.pooling_mode,
+            ctx.in_coords_key,
+            ctx.out_coords_key,
+            ctx.coordinate_manager._manager,
+        )
+        return grad_in_feat, None, None, None, None, None
+
+
+class MinkowskiGlobalPooling(MinkowskiModuleBase):
+    r"""Pool all input features to one output."""
+
+    def __init__(
+        self, mode: PoolingMode = PoolingMode.GLOBAL_AVG_POOLING_PYTORCH_INDEX
+    ):
+        r"""Reduces sparse coords into points at origin, i.e. reduce each point
+        cloud into a point at the origin, returning batch_size number of points
+        [[0, 0, ..., 0], [0, 0, ..., 1],, [0, 0, ..., 2]] where the last elem
+        of the coords is the batch index. The reduction function should be
+        provided as the mode.
+
+        Args:
+            :attr:`mode` (PoolingMode): Reduction function mode. E.g.
+            `PoolingMode.GLOBAL_SUM_POOLING_DEFAULT`
+
+        """
+        super(MinkowskiGlobalPooling, self).__init__()
+        assert isinstance(
+            mode, PoolingMode
+        ), f"Mode must be an instance of PoolingMode. mode={mode}"
+
+        self.pooling_mode = mode
+        self.pooling = MinkowskiGlobalPoolingFunction()
+
+    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,
+        )
+
+    def __repr__(self):
+        return self.__class__.__name__ + f"(mode={str(self.pooling_mode)})"
+
+
+class MinkowskiGlobalSumPooling(MinkowskiGlobalPooling):
+    def __init__(self, mode=PoolingMode.GLOBAL_SUM_POOLING_PYTORCH_INDEX):
+        r"""Reduces sparse coords into points at origin, i.e. reduce each point
+        cloud into a point at the origin, returning batch_size number of points
+        [[0, 0, ..., 0], [0, 0, ..., 1],, [0, 0, ..., 2]] where the last elem
+        of the coords is the batch index.
+
+        """
+        MinkowskiGlobalPooling.__init__(self, mode=mode)
+
+
+class MinkowskiGlobalAvgPooling(MinkowskiGlobalPooling):
+    def __init__(self, mode=PoolingMode.GLOBAL_AVG_POOLING_PYTORCH_INDEX):
+        r"""Reduces sparse coords into points at origin, i.e. reduce each point
+        cloud into a point at the origin, returning batch_size number of points
+        [[0, 0, ..., 0], [0, 0, ..., 1],, [0, 0, ..., 2]] where the last elem
+        of the coords is the batch index.
+
+        """
+        MinkowskiGlobalPooling.__init__(self, mode=mode)
+
+
+class MinkowskiGlobalMaxPooling(MinkowskiGlobalPooling):
+    r"""Max pool all input features to one output feature at the origin.
+
+    .. math::
+
+        \mathbf{y} = \max_{\mathbf{i} \in
+        \mathcal{C}^\text{in}} \mathbf{x}_{\mathbf{i}}
+
+    """
+
+    def __init__(self, mode=PoolingMode.GLOBAL_MAX_POOLING_PYTORCH_INDEX):
+        r"""Reduces sparse coords into points at origin, i.e. reduce each point
+        cloud into a point at the origin, returning batch_size number of points
+        [[0, 0, ..., 0], [0, 0, ..., 1],, [0, 0, ..., 2]] where the last elem
+        of the coords is the batch index.
+
+        """
+        MinkowskiGlobalPooling.__init__(self, mode=mode)
+
+    def forward(
+        self,
+        input,
+        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()
+            )
+        else:
+            in_coordinate_map_key = input.coordinate_map_key
+            out_coordinate_map_key = _get_coordinate_map_key(input, coordinates)
+        output = self.pooling.apply(
+            input.F,
+            self.pooling_mode,
+            in_coordinate_map_key,
+            out_coordinate_map_key,
+            input._manager,
+        )
+
+        return SparseTensor(
+            output,
+            coordinate_map_key=out_coordinate_map_key,
+            coordinate_manager=input.coordinate_manager,
+        )
+
+
+class MinkowskiDirectMaxPoolingFunction(Function):
+    @staticmethod
+    def forward(
+        ctx,
+        in_map: torch.Tensor,
+        out_map: torch.Tensor,
+        in_feat: torch.Tensor,
+        out_nrows: int,
+        is_sorted: bool = False,
+    ):
+        out_feat, max_mask = _C.direct_max_pool_fw(
+            in_map, out_map, in_feat, out_nrows, is_sorted
+        )
+        ctx.in_nrows = in_feat.size(0)
+        ctx.save_for_backward(max_mask)
+        return out_feat
+
+    @staticmethod
+    def backward(ctx, grad_out_feat):
+        grad_out_feat = grad_out_feat.contiguous()
+        max_mask = ctx.saved_tensors[0]
+        grad = _C.direct_max_pool_bw(grad_out_feat, max_mask, ctx.in_nrows)
+        return (
+            None,
+            None,
+            grad,
+            None,
+            None,
+        )

MinkowskiEngine/MinkowskiEngine/MinkowskiPruning.py

@@ -0,0 +1,121 @@
+# 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.
+import torch
+from torch.nn import Module
+from torch.autograd import Function
+
+from MinkowskiEngineBackend._C import CoordinateMapKey
+from MinkowskiSparseTensor import SparseTensor
+from MinkowskiCommon import (
+    MinkowskiModuleBase,
+    get_minkowski_function,
+)
+from MinkowskiCoordinateManager import CoordinateManager
+
+
+class MinkowskiPruningFunction(Function):
+    @staticmethod
+    def forward(
+        ctx,
+        in_feat: torch.Tensor,
+        mask: torch.Tensor,
+        in_coords_key: CoordinateMapKey,
+        out_coords_key: CoordinateMapKey = None,
+        coords_manager: CoordinateManager = None,
+    ):
+        ctx.in_coords_key = in_coords_key
+        ctx.out_coords_key = out_coords_key
+        ctx.coords_manager = coords_manager
+
+        in_feat = in_feat.contiguous()
+        fw_fn = get_minkowski_function("PruningForward", in_feat)
+        return fw_fn(
+            in_feat,
+            mask,
+            ctx.in_coords_key,
+            ctx.out_coords_key,
+            ctx.coords_manager._manager,
+        )
+
+    @staticmethod
+    def backward(ctx, grad_out_feat: torch.Tensor):
+        grad_out_feat = grad_out_feat.contiguous()
+        bw_fn = get_minkowski_function("PruningBackward", grad_out_feat)
+        grad_in_feat = bw_fn(
+            grad_out_feat,
+            ctx.in_coords_key,
+            ctx.out_coords_key,
+            ctx.coords_manager._manager,
+        )
+        return grad_in_feat, None, None, None, None
+
+
+class MinkowskiPruning(MinkowskiModuleBase):
+    r"""Remove specified coordinates from a :attr:`MinkowskiEngine.SparseTensor`.
+
+    """
+
+    def __init__(self):
+        super(MinkowskiPruning, self).__init__()
+        self.pruning = MinkowskiPruningFunction()
+
+    def forward(self, input: SparseTensor, mask: torch.Tensor):
+        r"""
+        Args:
+            :attr:`input` (:attr:`MinkowskiEnigne.SparseTensor`): a sparse tensor
+            to remove coordinates from.
+
+            :attr:`mask` (:attr:`torch.BoolTensor`): mask vector that specifies
+            which one to keep. Coordinates with False will be removed.
+
+        Returns:
+            A :attr:`MinkowskiEngine.SparseTensor` with C = coordinates
+            corresponding to `mask == True` F = copy of the features from `mask ==
+            True`.
+
+        Example::
+
+            >>> # Define inputs
+            >>> input = SparseTensor(feats, coords=coords)
+            >>> # Any boolean tensor can be used as the filter
+            >>> mask = torch.rand(feats.size(0)) < 0.5
+            >>> pruning = MinkowskiPruning()
+            >>> output = pruning(input, mask)
+
+        """
+        assert isinstance(input, SparseTensor)
+
+        out_coords_key = CoordinateMapKey(
+            input.coordinate_map_key.get_coordinate_size()
+        )
+        output = self.pruning.apply(
+            input.F, mask, input.coordinate_map_key, out_coords_key, input._manager
+        )
+        return SparseTensor(
+            output, coordinate_map_key=out_coords_key, coordinate_manager=input._manager
+        )
+
+    def __repr__(self):
+        return self.__class__.__name__ + "()"

MinkowskiEngine/MinkowskiEngine/MinkowskiSparseTensor.py

@@ -0,0 +1,783 @@
+# 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.
+import os
+import torch
+import warnings
+
+from MinkowskiCommon import convert_to_int_list, StrideType
+from MinkowskiEngineBackend._C import (
+    CoordinateMapKey,
+    CoordinateMapType,
+    GPUMemoryAllocatorType,
+    MinkowskiAlgorithm,
+)
+from MinkowskiTensor import (
+    SparseTensorQuantizationMode,
+    SparseTensorOperationMode,
+    Tensor,
+    sparse_tensor_operation_mode,
+    global_coordinate_manager,
+    set_global_coordinate_manager,
+)
+from MinkowskiCoordinateManager import CoordinateManager
+from sparse_matrix_functions import MinkowskiSPMMFunction, MinkowskiSPMMAverageFunction
+
+
+class SparseTensor(Tensor):
+    r"""A sparse tensor class. Can be accessed via
+    :attr:`MinkowskiEngine.SparseTensor`.
+
+    The :attr:`SparseTensor` class is the basic tensor in MinkowskiEngine. For
+    the definition of a sparse tensor, please visit `the terminology page
+    <https://nvidia.github.io/MinkowskiEngine/terminology.html#sparse-tensor>`_.
+    We use the COOrdinate (COO) format to save a sparse tensor `[1]
+    <http://groups.csail.mit.edu/commit/papers/2016/parker-thesis.pdf>`_. This
+    representation is simply a concatenation of coordinates in a matrix
+    :math:`C` and associated features :math:`F`.
+
+    .. math::
+
+       \mathbf{C} = \begin{bmatrix}
+       b_1    & x_1^1  & x_1^2  & \cdots & x_1^D  \\
+       \vdots & \vdots & \vdots & \ddots & \vdots \\
+       b_N    & x_N^1  & x_N^2  & \cdots & x_N^D
+       \end{bmatrix}, \; \mathbf{F} = \begin{bmatrix}
+       \mathbf{f}_1^T\\
+       \vdots\\
+       \mathbf{f}_N^T
+       \end{bmatrix}
+
+    where :math:`\mathbf{x}_i \in \mathcal{Z}^D` is a :math:`D`-dimensional
+    coordinate and :math:`b_i \in \mathcal{Z}_+` denotes the corresponding
+    batch index. :math:`N` is the number of non-zero elements in the sparse
+    tensor, each with the coordinate :math:`(b_i, x_i^1, x_i^1, \cdots,
+    x_i^D)`, and the associated feature :math:`\mathbf{f}_i`. Internally, we
+    handle the batch index as an additional spatial dimension.
+
+    Example::
+
+        >>> coords, feats = ME.utils.sparse_collate([coords_batch0, coords_batch1], [feats_batch0, feats_batch1])
+        >>> A = ME.SparseTensor(features=feats, coordinates=coords)
+        >>> B = ME.SparseTensor(features=feats, coordinate_map_key=A.coordiante_map_key, coordinate_manager=A.coordinate_manager)
+        >>> C = ME.SparseTensor(features=feats, coordinates=coords, quantization_mode=ME.SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE)
+        >>> D = ME.SparseTensor(features=feats, coordinates=coords, quantization_mode=ME.SparseTensorQuantizationMode.RANDOM_SUBSAMPLE)
+        >>> E = ME.SparseTensor(features=feats, coordinates=coords, tensor_stride=2)
+
+    .. warning::
+
+       To use the GPU-backend for coordinate management, the
+       :attr:`coordinates` must be a torch tensor on GPU. Applying `to(device)`
+       after :attr:`MinkowskiEngine.SparseTensor` initialization with a CPU
+       `coordinates` will waste time and computation on creating an unnecessary
+       CPU CoordinateMap since the GPU CoordinateMap will be created from
+       scratch as well.
+
+    .. warning::
+
+       Before MinkowskiEngine version 0.4, we put the batch indices on the last
+       column. Thus, direct manipulation of coordinates will be incompatible
+       with the latest versions. Instead, please use
+       :attr:`MinkowskiEngine.utils.batched_coordinates` or
+       :attr:`MinkowskiEngine.utils.sparse_collate` to create batched
+       coordinates.
+
+       Also, to access coordinates or features batch-wise, use the functions
+       :attr:`coordinates_at(batch_index : int)`, :attr:`features_at(batch_index : int)` of
+       a sparse tensor. Or to access all batch-wise coordinates and features,
+       `decomposed_coordinates`, `decomposed_features`,
+       `decomposed_coordinates_and_features` of a sparse tensor.
+
+       Example::
+
+           >>> coords, feats = ME.utils.sparse_collate([coords_batch0, coords_batch1], [feats_batch0, feats_batch1])
+           >>> A = ME.SparseTensor(features=feats, coordinates=coords)
+           >>> coords_batch0 = A.coordinates_at(batch_index=0)
+           >>> feats_batch1 = A.features_at(batch_index=1)
+           >>> list_of_coords, list_of_featurs = A.decomposed_coordinates_and_features
+
+    """
+
+    def __init__(
+        self,
+        features: torch.Tensor,
+        coordinates: torch.Tensor = None,
+        # optional coordinate related arguments
+        tensor_stride: StrideType = 1,
+        coordinate_map_key: CoordinateMapKey = None,
+        coordinate_manager: CoordinateManager = None,
+        quantization_mode: SparseTensorQuantizationMode = SparseTensorQuantizationMode.RANDOM_SUBSAMPLE,
+        # optional manager related arguments
+        allocator_type: GPUMemoryAllocatorType = None,
+        minkowski_algorithm: MinkowskiAlgorithm = None,
+        requires_grad=None,
+        device=None,
+    ):
+        r"""
+
+        Args:
+            :attr:`features` (:attr:`torch.FloatTensor`,
+            :attr:`torch.DoubleTensor`, :attr:`torch.cuda.FloatTensor`, or
+            :attr:`torch.cuda.DoubleTensor`): The features of a sparse
+            tensor.
+
+            :attr:`coordinates` (:attr:`torch.IntTensor`): The coordinates
+            associated to the features. If not provided, :attr:`coordinate_map_key`
+            must be provided.
+
+            :attr:`tensor_stride` (:attr:`int`, :attr:`list`,
+            :attr:`numpy.array`, or :attr:`tensor.Tensor`): The tensor stride
+            of the current sparse tensor. By default, it is 1.
+
+            :attr:`coordinate_map_key`
+            (:attr:`MinkowskiEngine.CoordinateMapKey`): When the coordinates
+            are already cached in the MinkowskiEngine, we could reuse the same
+            coordinate map by simply providing the coordinate map key. In most
+            case, this process is done automatically. When you provide a
+            `coordinate_map_key`, `coordinates` will be be ignored.
+
+            :attr:`coordinate_manager`
+            (:attr:`MinkowskiEngine.CoordinateManager`): The MinkowskiEngine
+            manages all coordinate maps using the `_C.CoordinateMapManager`. If
+            not provided, the MinkowskiEngine will create a new computation
+            graph. In most cases, this process is handled automatically and you
+            do not need to use this.
+
+            :attr:`quantization_mode`
+            (:attr:`MinkowskiEngine.SparseTensorQuantizationMode`): Defines how
+            continuous coordinates will be quantized to define a sparse tensor.
+            Please refer to :attr:`SparseTensorQuantizationMode` for details.
+
+            :attr:`allocator_type`
+            (:attr:`MinkowskiEngine.GPUMemoryAllocatorType`): Defines the GPU
+            memory allocator type. By default, it uses the c10 allocator.
+
+            :attr:`minkowski_algorithm`
+            (:attr:`MinkowskiEngine.MinkowskiAlgorithm`): Controls the mode the
+            minkowski engine runs, Use
+            :attr:`MinkowskiAlgorithm.MEMORY_EFFICIENT` if you want to reduce
+            the memory footprint. Or use
+            :attr:`MinkowskiAlgorithm.SPEED_OPTIMIZED` if you want to make it
+            run fasterat the cost of more memory.
+
+            :attr:`requires_grad` (:attr:`bool`): Set the requires_grad flag.
+
+            :attr:`device` (:attr:`torch.device`): Set the device the sparse
+            tensor is defined.
+
+        """
+        # Type checks
+        assert isinstance(features, torch.Tensor), "Features must be a torch.Tensor"
+        assert (
+            features.ndim == 2
+        ), f"The feature should be a matrix, The input feature is an order-{features.ndim} tensor."
+        assert isinstance(quantization_mode, SparseTensorQuantizationMode)
+        self.quantization_mode = quantization_mode
+
+        if coordinates is not None:
+            assert isinstance(coordinates, torch.Tensor)
+        if coordinate_map_key is not None:
+            assert isinstance(coordinate_map_key, CoordinateMapKey)
+            assert (
+                coordinate_manager is not None
+            ), "Must provide coordinate_manager if coordinate_map_key is provided"
+            assert (
+                coordinates is None
+            ), "Must not provide coordinates if coordinate_map_key is provided"
+        if coordinate_manager is not None:
+            assert isinstance(coordinate_manager, CoordinateManager)
+        if coordinates is None and (
+            coordinate_map_key is None or coordinate_manager is None
+        ):
+            raise ValueError(
+                "Either coordinates or (coordinate_map_key, coordinate_manager) pair must be provided."
+            )
+
+        Tensor.__init__(self)
+
+        # To device
+        if device is not None:
+            features = features.to(device)
+            if coordinates is not None:
+                # assertion check for the map key done later
+                coordinates = coordinates.to(device)
+
+        self._D = (
+            coordinates.size(1) - 1 if coordinates is not None else coordinate_manager.D
+        )
+        ##########################
+        # Setup CoordsManager
+        ##########################
+        if coordinate_manager is None:
+            # If set to share the coords man, use the global coords man
+            if (
+                sparse_tensor_operation_mode()
+                == SparseTensorOperationMode.SHARE_COORDINATE_MANAGER
+            ):
+                coordinate_manager = global_coordinate_manager()
+                if coordinate_manager is None:
+                    coordinate_manager = CoordinateManager(
+                        D=self._D,
+                        coordinate_map_type=CoordinateMapType.CUDA
+                        if coordinates.is_cuda
+                        else CoordinateMapType.CPU,
+                        allocator_type=allocator_type,
+                        minkowski_algorithm=minkowski_algorithm,
+                    )
+                    set_global_coordinate_manager(coordinate_manager)
+            else:
+                coordinate_manager = CoordinateManager(
+                    D=coordinates.size(1) - 1,
+                    coordinate_map_type=CoordinateMapType.CUDA
+                    if coordinates.is_cuda
+                    else CoordinateMapType.CPU,
+                    allocator_type=allocator_type,
+                    minkowski_algorithm=minkowski_algorithm,
+                )
+        self._manager = coordinate_manager
+
+        ##########################
+        # Initialize coords
+        ##########################
+        if coordinates is not None:
+            assert (
+                features.shape[0] == coordinates.shape[0]
+            ), "The number of rows in features and coordinates must match."
+
+            assert (
+                features.is_cuda == coordinates.is_cuda
+            ), "Features and coordinates must have the same backend."
+
+            coordinate_map_key = CoordinateMapKey(
+                convert_to_int_list(tensor_stride, self._D), ""
+            )
+            coordinates, features, coordinate_map_key = self.initialize_coordinates(
+                coordinates, features, coordinate_map_key
+            )
+        else:  # coordinate_map_key is not None:
+            assert coordinate_map_key.is_key_set(), "The coordinate key must be valid."
+
+        if requires_grad is not None:
+            features.requires_grad_(requires_grad)
+
+        self._F = features
+        self._C = coordinates
+        self.coordinate_map_key = coordinate_map_key
+        self._batch_rows = None
+
+    @property
+    def coordinate_key(self):
+        return self.coordinate_map_key
+
+    def initialize_coordinates(self, coordinates, features, coordinate_map_key):
+        if not isinstance(coordinates, (torch.IntTensor, torch.cuda.IntTensor)):
+            warnings.warn(
+                "coordinates implicitly converted to torch.IntTensor. "
+                + "To remove this warning, use `.int()` to convert the "
+                + "coords into an torch.IntTensor"
+            )
+            coordinates = torch.floor(coordinates).int()
+        (
+            coordinate_map_key,
+            (unique_index, inverse_mapping),
+        ) = self._manager.insert_and_map(coordinates, *coordinate_map_key.get_key())
+        self.unique_index = unique_index.long()
+        coordinates = coordinates[self.unique_index]
+
+        if len(inverse_mapping) == 0:
+            # When the input has the same shape as the output
+            self.inverse_mapping = torch.arange(
+                len(features),
+                dtype=inverse_mapping.dtype,
+                device=inverse_mapping.device,
+            )
+            return coordinates, features, coordinate_map_key
+
+        self.inverse_mapping = inverse_mapping
+        if self.quantization_mode == SparseTensorQuantizationMode.UNWEIGHTED_SUM:
+            spmm = MinkowskiSPMMFunction()
+            N = len(features)
+            cols = torch.arange(
+                N,
+                dtype=self.inverse_mapping.dtype,
+                device=self.inverse_mapping.device,
+            )
+            vals = torch.ones(N, dtype=features.dtype, device=features.device)
+            size = torch.Size([len(self.unique_index), len(self.inverse_mapping)])
+            features = spmm.apply(self.inverse_mapping, cols, vals, size, features)
+        elif self.quantization_mode == SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE:
+            spmm_avg = MinkowskiSPMMAverageFunction()
+            N = len(features)
+            cols = torch.arange(
+                N,
+                dtype=self.inverse_mapping.dtype,
+                device=self.inverse_mapping.device,
+            )
+            size = torch.Size([len(self.unique_index), len(self.inverse_mapping)])
+            features = spmm_avg.apply(self.inverse_mapping, cols, size, features)
+        elif self.quantization_mode == SparseTensorQuantizationMode.RANDOM_SUBSAMPLE:
+            features = features[self.unique_index]
+        else:
+            # No quantization
+            pass
+
+        return coordinates, features, coordinate_map_key
+
+    # Conversion functions
+    def sparse(self, min_coords=None, max_coords=None, contract_coords=True):
+        r"""Convert the :attr:`MinkowskiEngine.SparseTensor` to a torch sparse
+        tensor.
+
+        Args:
+            :attr:`min_coords` (torch.IntTensor, optional): The min
+            coordinates of the output sparse tensor. Must be divisible by the
+            current :attr:`tensor_stride`.
+
+            :attr:`max_coords` (torch.IntTensor, optional): The max coordinates
+            of the output sparse tensor (inclusive). Must be divisible by the
+            current :attr:`tensor_stride`.
+
+            :attr:`contract_coords` (bool, optional): Given True, the output
+            coordinates will be divided by the tensor stride to make features
+            contiguous.
+
+        Returns:
+            :attr:`spare_tensor` (torch.sparse.Tensor): the torch sparse tensor
+            representation of the self in `[Batch Dim, Spatial Dims..., Feature
+            Dim]`. The coordinate of each feature can be accessed via
+            `min_coord + tensor_stride * [the coordinate of the dense tensor]`.
+
+            :attr:`min_coords` (torch.IntTensor): the D-dimensional vector
+            defining the minimum coordinate of the output sparse tensor. If
+            :attr:`contract_coords` is True, the :attr:`min_coords` will also
+            be contracted.
+
+            :attr:`tensor_stride` (torch.IntTensor): the D-dimensional vector
+            defining the stride between tensor elements.
+
+        """
+
+        if min_coords is not None:
+            assert isinstance(min_coords, torch.IntTensor)
+            assert min_coords.numel() == self._D
+        if max_coords is not None:
+            assert isinstance(max_coords, torch.IntTensor)
+            assert min_coords.numel() == self._D
+
+        def torch_sparse_Tensor(coords, feats, size=None):
+            if size is None:
+                if feats.dtype == torch.float64:
+                    return torch.sparse.DoubleTensor(coords, feats)
+                elif feats.dtype == torch.float32:
+                    return torch.sparse.FloatTensor(coords, feats)
+                else:
+                    raise ValueError("Feature type not supported.")
+            else:
+                if feats.dtype == torch.float64:
+                    return torch.sparse.DoubleTensor(coords, feats, size)
+                elif feats.dtype == torch.float32:
+                    return torch.sparse.FloatTensor(coords, feats, size)
+                else:
+                    raise ValueError("Feature type not supported.")
+
+        # Use int tensor for all operations
+        tensor_stride = torch.IntTensor(self.tensor_stride)
+
+        # New coordinates
+        coords = self.C
+        coords, batch_indices = coords[:, 1:], coords[:, 0]
+
+        if min_coords is None:
+            min_coords, _ = coords.min(0, keepdim=True)
+        elif min_coords.ndim == 1:
+            min_coords = min_coords.unsqueeze(0)
+
+        assert (
+            min_coords % tensor_stride
+        ).sum() == 0, "The minimum coordinates must be divisible by the tensor stride."
+
+        if max_coords is not None:
+            if max_coords.ndim == 1:
+                max_coords = max_coords.unsqueeze(0)
+            assert (
+                max_coords % tensor_stride
+            ).sum() == 0, (
+                "The maximum coordinates must be divisible by the tensor stride."
+            )
+
+        coords -= min_coords
+
+        if coords.ndim == 1:
+            coords = coords.unsqueeze(1)
+        if batch_indices.ndim == 1:
+            batch_indices = batch_indices.unsqueeze(1)
+
+        # return the contracted tensor
+        if contract_coords:
+            coords = coords // tensor_stride
+            if max_coords is not None:
+                max_coords = max_coords // tensor_stride
+            min_coords = min_coords // tensor_stride
+
+        new_coords = torch.cat((batch_indices, coords), dim=1).long()
+
+        size = None
+        if max_coords is not None:
+            size = max_coords - min_coords + 1  # inclusive
+            # Squeeze to make the size one-dimensional
+            size = size.squeeze()
+
+            max_batch = max(self._manager.get_batch_indices())
+            size = torch.Size([max_batch + 1, *size, self.F.size(1)])
+
+        sparse_tensor = torch_sparse_Tensor(
+            new_coords.t().to(self.F.device), self.F, size
+        )
+        tensor_stride = torch.IntTensor(self.tensor_stride)
+        return sparse_tensor, min_coords, tensor_stride
+
+    def dense(self, shape=None, min_coordinate=None, contract_stride=True):
+        r"""Convert the :attr:`MinkowskiEngine.SparseTensor` to a torch dense
+        tensor.
+
+        Args:
+            :attr:`shape` (torch.Size, optional): The size of the output tensor.
+
+            :attr:`min_coordinate` (torch.IntTensor, optional): The min
+            coordinates of the output sparse tensor. Must be divisible by the
+            current :attr:`tensor_stride`. If 0 is given, it will use the origin for the min coordinate.
+
+            :attr:`contract_stride` (bool, optional): The output coordinates
+            will be divided by the tensor stride to make features spatially
+            contiguous. True by default.
+
+        Returns:
+            :attr:`tensor` (torch.Tensor): the torch tensor with size `[Batch
+            Dim, Feature Dim, Spatial Dim..., Spatial Dim]`. The coordinate of
+            each feature can be accessed via `min_coordinate + tensor_stride *
+            [the coordinate of the dense tensor]`.
+
+            :attr:`min_coordinate` (torch.IntTensor): the D-dimensional vector
+            defining the minimum coordinate of the output tensor.
+
+            :attr:`tensor_stride` (torch.IntTensor): the D-dimensional vector
+            defining the stride between tensor elements.
+
+        """
+        if min_coordinate is not None:
+            assert isinstance(min_coordinate, torch.IntTensor)
+            assert min_coordinate.numel() == self._D
+        if shape is not None:
+            assert isinstance(shape, torch.Size)
+            assert len(shape) == self._D + 2  # batch and channel
+            if shape[1] != self._F.size(1):
+                shape = torch.Size([shape[0], self._F.size(1), *[s for s in shape[2:]]])
+
+        # Exception handling for empty tensor
+        if self.__len__() == 0:
+            assert shape is not None, "shape is required to densify an empty tensor"
+            return (
+                torch.zeros(shape, dtype=self.dtype, device=self.device),
+                torch.zeros(self._D, dtype=torch.int32, device=self.device),
+                self.tensor_stride,
+            )
+
+        # Use int tensor for all operations
+        tensor_stride = torch.IntTensor(self.tensor_stride).to(self.device)
+
+        # New coordinates
+        batch_indices = self.C[:, 0]
+
+        if min_coordinate is None:
+            min_coordinate, _ = self.C.min(0, keepdim=True)
+            min_coordinate = min_coordinate[:, 1:]
+            if not torch.all(min_coordinate >= 0):
+                raise ValueError(
+                    f"Coordinate has a negative value: {min_coordinate}. Please provide min_coordinate argument"
+                )
+            coords = self.C[:, 1:]
+        elif isinstance(min_coordinate, int) and min_coordinate == 0:
+            coords = self.C[:, 1:]
+        else:
+            min_coordinate = min_coordinate.to(self.device)
+            if min_coordinate.ndim == 1:
+                min_coordinate = min_coordinate.unsqueeze(0)
+            coords = self.C[:, 1:] - min_coordinate
+
+        assert (
+            min_coordinate % tensor_stride
+        ).sum() == 0, "The minimum coordinates must be divisible by the tensor stride."
+
+        if coords.ndim == 1:
+            coords = coords.unsqueeze(1)
+
+        # return the contracted tensor
+        if contract_stride:
+            coords = coords // tensor_stride
+
+        nchannels = self.F.size(1)
+        if shape is None:
+            size = coords.max(0)[0] + 1
+            shape = torch.Size(
+                [batch_indices.max() + 1, nchannels, *size.cpu().numpy()]
+            )
+
+        dense_F = torch.zeros(shape, dtype=self.dtype, device=self.device)
+
+        tcoords = coords.t().long()
+        batch_indices = batch_indices.long()
+        exec(
+            "dense_F[batch_indices, :, "
+            + ", ".join([f"tcoords[{i}]" for i in range(len(tcoords))])
+            + "] = self.F"
+        )
+
+        tensor_stride = torch.IntTensor(self.tensor_stride)
+        return dense_F, min_coordinate, tensor_stride
+
+    def interpolate(self, X):
+        from MinkowskiTensorField import TensorField
+
+        assert isinstance(X, TensorField)
+        if self.coordinate_map_key in X._splat:
+            tensor_map, field_map, weights, size = X._splat[self.coordinate_map_key]
+            size = torch.Size([size[1], size[0]])  # transpose
+            features = MinkowskiSPMMFunction().apply(
+                field_map, tensor_map, weights, size, self._F
+            )
+        else:
+            features = self.features_at_coordinates(X.C)
+        return TensorField(
+            features=features,
+            coordinate_field_map_key=X.coordinate_field_map_key,
+            coordinate_manager=X.coordinate_manager,
+        )
+
+    def slice(self, X):
+        r"""
+
+        Args:
+           :attr:`X` (:attr:`MinkowskiEngine.SparseTensor`): a sparse tensor
+           that discretized the original input.
+
+        Returns:
+           :attr:`tensor_field` (:attr:`MinkowskiEngine.TensorField`): the
+           resulting tensor field contains features on the continuous
+           coordinates that generated the input X.
+
+        Example::
+
+           >>> # coords, feats from a data loader
+           >>> print(len(coords))  # 227742
+           >>> tfield = ME.TensorField(coordinates=coords, features=feats, quantization_mode=SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE)
+           >>> print(len(tfield))  # 227742
+           >>> sinput = tfield.sparse() # 161890 quantization results in fewer voxels
+           >>> soutput = MinkUNet(sinput)
+           >>> print(len(soutput))  # 161890 Output with the same resolution
+           >>> ofield = soutput.slice(tfield)
+           >>> assert isinstance(ofield, ME.TensorField)
+           >>> len(ofield) == len(coords)  # recovers the original ordering and length
+           >>> assert isinstance(ofield.F, torch.Tensor)  # .F returns the features
+        """
+        # Currently only supports unweighted slice.
+        assert X.quantization_mode in [
+            SparseTensorQuantizationMode.RANDOM_SUBSAMPLE,
+            SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE,
+        ], "slice only available for sparse tensors with quantization RANDOM_SUBSAMPLE or UNWEIGHTED_AVERAGE"
+
+        from MinkowskiTensorField import TensorField
+
+        if isinstance(X, TensorField):
+            return TensorField(
+                self.F[X.inverse_mapping(self.coordinate_map_key).long()],
+                coordinate_field_map_key=X.coordinate_field_map_key,
+                coordinate_manager=X.coordinate_manager,
+                quantization_mode=X.quantization_mode,
+            )
+        elif isinstance(X, SparseTensor):
+            inv_map = X.inverse_mapping
+            assert (
+                X.coordinate_map_key == self.coordinate_map_key
+            ), "Slice can only be applied on the same coordinates (coordinate_map_key)"
+            return TensorField(
+                self.F[inv_map],
+                coordinates=self.C[inv_map],
+                coordinate_manager=self.coordinate_manager,
+                quantization_mode=self.quantization_mode,
+            )
+        else:
+            raise ValueError(
+                "Invalid input. The input must be an instance of TensorField or SparseTensor."
+            )
+
+    def cat_slice(self, X):
+        r"""
+
+        Args:
+           :attr:`X` (:attr:`MinkowskiEngine.SparseTensor`): a sparse tensor
+           that discretized the original input.
+
+        Returns:
+           :attr:`tensor_field` (:attr:`MinkowskiEngine.TensorField`): the
+           resulting tensor field contains the concatenation of features on the
+           original continuous coordinates that generated the input X and the
+           self.
+
+        Example::
+
+           >>> # coords, feats from a data loader
+           >>> print(len(coords))  # 227742
+           >>> sinput = ME.SparseTensor(coordinates=coords, features=feats, quantization_mode=SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE)
+           >>> print(len(sinput))  # 161890 quantization results in fewer voxels
+           >>> soutput = network(sinput)
+           >>> print(len(soutput))  # 161890 Output with the same resolution
+           >>> ofield = soutput.cat_slice(sinput)
+           >>> assert soutput.F.size(1) + sinput.F.size(1) == ofield.F.size(1)  # concatenation of features
+        """
+        # Currently only supports unweighted slice.
+        assert X.quantization_mode in [
+            SparseTensorQuantizationMode.RANDOM_SUBSAMPLE,
+            SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE,
+        ], "slice only available for sparse tensors with quantization RANDOM_SUBSAMPLE or UNWEIGHTED_AVERAGE"
+
+        from MinkowskiTensorField import TensorField
+
+        inv_map = X.inverse_mapping(self.coordinate_map_key)
+        features = torch.cat((self.F[inv_map], X.F), dim=1)
+        if isinstance(X, TensorField):
+            return TensorField(
+                features,
+                coordinate_field_map_key=X.coordinate_field_map_key,
+                coordinate_manager=X.coordinate_manager,
+                quantization_mode=X.quantization_mode,
+            )
+        elif isinstance(X, SparseTensor):
+            assert (
+                X.coordinate_map_key == self.coordinate_map_key
+            ), "Slice can only be applied on the same coordinates (coordinate_map_key)"
+            return TensorField(
+                features,
+                coordinates=self.C[inv_map],
+                coordinate_manager=self.coordinate_manager,
+                quantization_mode=self.quantization_mode,
+            )
+        else:
+            raise ValueError(
+                "Invalid input. The input must be an instance of TensorField or SparseTensor."
+            )
+
+    def features_at_coordinates(self, query_coordinates: torch.Tensor):
+        r"""Extract features at the specified continuous coordinate matrix.
+
+        Args:
+           :attr:`query_coordinates` (:attr:`torch.FloatTensor`): a coordinate
+           matrix of size :math:`N \times (D + 1)` where :math:`D` is the size
+           of the spatial dimension.
+
+        Returns:
+           :attr:`queried_features` (:attr:`torch.Tensor`): a feature matrix of
+           size :math:`N \times D_F` where :math:`D_F` is the number of
+           channels in the feature. For coordinates not present in the current
+           sparse tensor, corresponding feature rows will be zeros.
+        """
+        from MinkowskiInterpolation import MinkowskiInterpolationFunction
+
+        assert (
+            self.dtype == query_coordinates.dtype
+        ), "Invalid query_coordinates dtype. use {self.dtype}"
+
+        assert (
+            query_coordinates.device == self.device
+        ), "query coordinates device ({query_coordinates.device}) does not match the sparse tensor device ({self.device})."
+        return MinkowskiInterpolationFunction().apply(
+            self._F,
+            query_coordinates,
+            self.coordinate_map_key,
+            self.coordinate_manager,
+        )[0]
+
+    def __repr__(self):
+        return (
+            self.__class__.__name__
+            + "("
+            + os.linesep
+            + "  coordinates="
+            + str(self.C)
+            + os.linesep
+            + "  features="
+            + str(self.F)
+            + os.linesep
+            + "  coordinate_map_key="
+            + str(self.coordinate_map_key)
+            + os.linesep
+            + "  coordinate_manager="
+            + str(self._manager)
+            + "  spatial dimension="
+            + str(self._D)
+            + ")"
+        )
+
+    __slots__ = (
+        "_C",
+        "_F",
+        "_D",
+        "coordinate_map_key",
+        "_manager",
+        "unique_index",
+        "inverse_mapping",
+        "quantization_mode",
+        "_batch_rows",
+    )
+
+
+def _get_coordinate_map_key(
+    input: SparseTensor,
+    coordinates: torch.Tensor = None,
+    tensor_stride: StrideType = 1,
+    expand_coordinates: bool = False,
+):
+    r"""Returns the coordinates map key."""
+    if coordinates is not None and not expand_coordinates:
+        assert isinstance(coordinates, (CoordinateMapKey, torch.Tensor, SparseTensor))
+        if isinstance(coordinates, torch.Tensor):
+            assert coordinates.ndim == 2
+            coordinate_map_key = CoordinateMapKey(
+                convert_to_int_list(tensor_stride, coordinates.size(1) - 1), ""
+            )
+
+            (
+                coordinate_map_key,
+                (unique_index, inverse_mapping),
+            ) = input._manager.insert_and_map(
+                coordinates, *coordinate_map_key.get_key()
+            )
+        elif isinstance(coordinates, SparseTensor):
+            coordinate_map_key = coordinates.coordinate_map_key
+        else:  # CoordinateMapKey type due to the previous assertion
+            coordinate_map_key = coordinates
+    else:  # coordinates is None
+        coordinate_map_key = CoordinateMapKey(
+            input.coordinate_map_key.get_coordinate_size()
+        )
+    return coordinate_map_key

MinkowskiEngine/MinkowskiEngine/MinkowskiTensor.py

@@ -0,0 +1,604 @@
+# 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.
+import os
+import torch
+import copy
+from enum import Enum
+
+from MinkowskiEngineBackend._C import CoordinateMapKey
+
+
+class SparseTensorOperationMode(Enum):
+    r"""Enum class for SparseTensor internal instantiation modes.
+
+    :attr:`SEPARATE_COORDINATE_MANAGER`: always create a new coordinate manager.
+
+    :attr:`SHARE_COORDINATE_MANAGER`: always use the globally defined coordinate
+    manager. Must clear the coordinate manager manually by
+    :attr:`MinkowskiEngine.SparseTensor.clear_global_coordinate_manager`.
+
+    """
+    SEPARATE_COORDINATE_MANAGER = 0
+    SHARE_COORDINATE_MANAGER = 1
+
+
+class SparseTensorQuantizationMode(Enum):
+    r"""
+    `RANDOM_SUBSAMPLE`: Subsample one coordinate per each quantization block randomly.
+    `UNWEIGHTED_AVERAGE`: average all features within a quantization block equally.
+    `UNWEIGHTED_SUM`: sum all features within a quantization block equally.
+    `NO_QUANTIZATION`: No quantization is applied. Should not be used for normal operation.
+    `MAX_POOL`: Voxel-wise max pooling is applied.
+    `SPLAT_LINEAR_INTERPOLATION`: Splat features using N-dimensional linear interpolation to 2^N neighbors.
+    """
+    RANDOM_SUBSAMPLE = 0
+    UNWEIGHTED_AVERAGE = 1
+    UNWEIGHTED_SUM = 2
+    NO_QUANTIZATION = 3
+    MAX_POOL = 4
+    SPLAT_LINEAR_INTERPOLATION = 5
+
+
+_sparse_tensor_operation_mode = SparseTensorOperationMode.SEPARATE_COORDINATE_MANAGER
+_global_coordinate_manager = None
+
+COORDINATE_MANAGER_DIFFERENT_ERROR = "SparseTensors must share the same coordinate manager for this operation. Please refer to the SparseTensor creation API (https://nvidia.github.io/MinkowskiEngine/sparse_tensor.html) to share the coordinate manager, or set the sparse tensor operation mode with `set_sparse_tensor_operation_mode` to share it by default."
+COORDINATE_KEY_DIFFERENT_ERROR = "SparseTensors must have the same coordinate_map_key."
+
+
+def set_sparse_tensor_operation_mode(operation_mode: SparseTensorOperationMode):
+    r"""Define the sparse tensor coordinate manager operation mode.
+
+    By default, a :attr:`MinkowskiEngine.SparseTensor.SparseTensor`
+    instantiation creates a new coordinate manager that is not shared with
+    other sparse tensors. By setting this function with
+    :attr:`MinkowskiEngine.SparseTensorOperationMode.SHARE_COORDINATE_MANAGER`, you
+    can share the coordinate manager globally with other sparse tensors.
+    However, you must explicitly clear the coordinate manger after use. Please
+    refer to :attr:`MinkowskiEngine.clear_global_coordinate_manager`.
+
+    Args:
+        :attr:`operation_mode`
+        (:attr:`MinkowskiEngine.SparseTensorOperationMode`): The operation mode
+        for the sparse tensor coordinate manager. By default
+        :attr:`MinkowskiEngine.SparseTensorOperationMode.SEPARATE_COORDINATE_MANAGER`.
+
+    Example:
+
+        >>> import MinkowskiEngine as ME
+        >>> ME.set_sparse_tensor_operation_mode(ME.SparseTensorOperationMode.SHARE_COORDINATE_MANAGER)
+        >>> ...
+        >>> a = ME.SparseTensor(...)
+        >>> b = ME.SparseTensor(...)  # coords_man shared
+        >>> ...  # one feed forward and backward
+        >>> ME.clear_global_coordinate_manager()  # Must use to clear the coordinates after one forward/backward
+
+    """
+    assert isinstance(
+        operation_mode, SparseTensorOperationMode
+    ), f"Input must be an instance of SparseTensorOperationMode not {operation_mode}"
+    global _sparse_tensor_operation_mode
+    _sparse_tensor_operation_mode = operation_mode
+
+
+def sparse_tensor_operation_mode() -> SparseTensorOperationMode:
+    r"""Return the current sparse tensor operation mode."""
+    global _sparse_tensor_operation_mode
+    return copy.deepcopy(_sparse_tensor_operation_mode)
+
+
+def global_coordinate_manager():
+    r"""Return the current global coordinate manager"""
+    global _global_coordinate_manager
+    return _global_coordinate_manager
+
+
+def set_global_coordinate_manager(coordinate_manager):
+    r"""Set the global coordinate manager.
+
+    :attr:`MinkowskiEngine.CoordinateManager` The coordinate manager which will
+    be set to the global coordinate manager.
+    """
+    global _global_coordinate_manager
+    _global_coordinate_manager = coordinate_manager
+
+
+def clear_global_coordinate_manager():
+    r"""Clear the global coordinate manager cache.
+
+    When you use the operation mode:
+    :attr:`MinkowskiEngine.SparseTensor.SparseTensorOperationMode.SHARE_COORDINATE_MANAGER`,
+    you must explicitly clear the coordinate manager after each feed forward/backward.
+    """
+    global _global_coordinate_manager
+    _global_coordinate_manager = None
+
+
+class Tensor:
+    r"""A sparse tensor class. Can be accessed via
+    :attr:`MinkowskiEngine.SparseTensor`.
+
+    The :attr:`SparseTensor` class is the basic tensor in MinkowskiEngine. For
+    the definition of a sparse tensor, please visit `the terminology page
+    <https://nvidia.github.io/MinkowskiEngine/terminology.html#sparse-tensor>`_.
+    We use the COOrdinate (COO) format to save a sparse tensor `[1]
+    <http://groups.csail.mit.edu/commit/papers/2016/parker-thesis.pdf>`_. This
+    representation is simply a concatenation of coordinates in a matrix
+    :math:`C` and associated features :math:`F`.
+
+    .. math::
+
+       \mathbf{C} = \begin{bmatrix}
+       b_1    & x_1^1  & x_1^2  & \cdots & x_1^D  \\
+       \vdots & \vdots & \vdots & \ddots & \vdots \\
+       b_N    & x_N^1  & x_N^2  & \cdots & x_N^D
+       \end{bmatrix}, \; \mathbf{F} = \begin{bmatrix}
+       \mathbf{f}_1^T\\
+       \vdots\\
+       \mathbf{f}_N^T
+       \end{bmatrix}
+
+    where :math:`\mathbf{x}_i \in \mathcal{Z}^D` is a :math:`D`-dimensional
+    coordinate and :math:`b_i \in \mathcal{Z}_+` denotes the corresponding
+    batch index. :math:`N` is the number of non-zero elements in the sparse
+    tensor, each with the coordinate :math:`(b_i, x_i^1, x_i^1, \cdots,
+    x_i^D)`, and the associated feature :math:`\mathbf{f}_i`. Internally, we
+    handle the batch index as an additional spatial dimension.
+
+    Example::
+
+        >>> coords, feats = ME.utils.sparse_collate([coords_batch0, coords_batch1], [feats_batch0, feats_batch1])
+        >>> A = ME.SparseTensor(features=feats, coordinates=coords)
+        >>> B = ME.SparseTensor(features=feats, coordinate_map_key=A.coordiante_map_key, coordinate_manager=A.coordinate_manager)
+        >>> C = ME.SparseTensor(features=feats, coordinates=coords, quantization_mode=ME.SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE)
+        >>> D = ME.SparseTensor(features=feats, coordinates=coords, tensor_stride=2)
+
+    .. warning::
+
+       To use the GPU-backend for coordinate management, the
+       :attr:`coordinates` must be a torch tensor on GPU. Applying `to(device)`
+       after a :attr:`MinkowskiEngine.SparseTensor` initialization with a CPU
+       `coordinates` will waste time and computation for creating a CPU
+       CoordinateMap since GPU CoordinateMap will be created from scratch.
+
+    .. warning::
+
+       Before MinkowskiEngine version 0.4, we put the batch indices on the last
+       column. Thus, direct manipulation of coordinates will be incompatible
+       with the latest versions. Instead, please use
+       :attr:`MinkowskiEngine.utils.batched_coordinates` or
+       :attr:`MinkowskiEngine.utils.sparse_collate` to create batched
+       coordinates.
+
+       Also, to access coordinates or features batch-wise, use the functions
+       :attr:`coordinates_at(batch_index : int)`, :attr:`features_at(batch_index : int)` of
+       a sparse tensor. Or to access all batch-wise coordinates and features,
+       `decomposed_coordinates`, `decomposed_features`,
+       `decomposed_coordinates_and_features` of a sparse tensor.
+
+       Example::
+
+           >>> coords, feats = ME.utils.sparse_collate([coords_batch0, coords_batch1], [feats_batch0, feats_batch1])
+           >>> A = ME.SparseTensor(features=feats, coordinates=coords)
+           >>> coords_batch0 = A.coordinates_at(batch_index=0)
+           >>> feats_batch1 = A.features_at(batch_index=1)
+           >>> list_of_coords, list_of_featurs = A.decomposed_coordinates_and_features
+
+    """
+
+    @property
+    def coordinate_manager(self):
+        return self._manager
+
+    @property
+    def tensor_stride(self):
+        return self.coordinate_map_key.get_tensor_stride()
+
+    @tensor_stride.setter
+    def tensor_stride(self, p):
+        r"""
+        This function is not recommended to be used directly.
+        """
+        raise SyntaxError("Direct modification of tensor_stride is not permitted")
+
+    def _get_coordinates(self):
+        return self._manager.get_coordinates(self.coordinate_map_key)
+
+    @property
+    def C(self):
+        r"""The alias of :attr:`coords`."""
+        return self.coordinates
+
+    @property
+    def coordinates(self):
+        r"""
+        The coordinates of the current sparse tensor. The coordinates are
+        represented as a :math:`N \times (D + 1)` dimensional matrix where
+        :math:`N` is the number of points in the space and :math:`D` is the
+        dimension of the space (e.g. 3 for 3D, 4 for 3D + Time). Additional
+        dimension of the column of the matrix C is for batch indices which is
+        internally treated as an additional spatial dimension to disassociate
+        different instances in a batch.
+        """
+        if self._C is None:
+            self._C = self._get_coordinates()
+        return self._C
+
+    @property
+    def coordinate_key(self):
+        raise NotImplementedError("Tensor interface does not have coordinate_key")
+
+    @C.setter
+    def C(self):
+        raise SyntaxError("Direct modification of coordinates is not permitted")
+
+    @coordinates.setter
+    def coordinates(self):
+        raise SyntaxError("Direct modification of coordinates is not permitted")
+
+    @property
+    def F(self):
+        r"""The alias of :attr:`feats`."""
+        return self._F
+
+    @property
+    def features(self):
+        r"""
+        The features of the current sparse tensor. The features are :math:`N
+        \times D_F` where :math:`N` is the number of points in the space and
+        :math:`D_F` is the dimension of each feature vector. Please refer to
+        :attr:`coords` to access the associated coordinates.
+        """
+        return self._F
+
+    @property
+    def _batchwise_row_indices(self):
+        if self._batch_rows is None:
+            _, self._batch_rows = self._manager.origin_map(self.coordinate_map_key)
+        return self._batch_rows
+
+    @property
+    def _sorted_batchwise_row_indices(self):
+        if self._sorted_batch_rows is None:
+            batch_rows = self._batchwise_row_indices
+            with torch.no_grad():
+                self._sorted_batch_rows = [t.sort()[0] for t in batch_rows]
+        return self._sorted_batch_rows
+
+    @property
+    def decomposition_permutations(self):
+        r"""Returns a list of indices per batch that where indices defines the permutation of the batch-wise decomposition.
+
+        Example::
+
+            >>> # coords, feats, labels are given. All follow the same order
+            >>> stensor = ME.SparseTensor(feats, coords)
+            >>> conv = ME.MinkowskiConvolution(in_channels=3, out_nchannel=3, kernel_size=3, dimension=3)
+            >>> list_of_featurs = stensor.decomposed_features
+            >>> list_of_permutations = stensor.decomposition_permutations
+            >>> # list_of_features == [feats[inds] for inds in list_of_permutations]
+            >>> list_of_decomposed_labels = [labels[inds] for inds in list_of_permutations]
+            >>> for curr_feats, curr_labels in zip(list_of_features, list_of_decomposed_labels):
+            >>>     loss += torch.functional.mse_loss(curr_feats, curr_labels)
+        """
+        return self._batchwise_row_indices
+
+    @property
+    def decomposed_coordinates(self):
+        r"""Returns a list of coordinates per batch.
+
+        Returns a list of torch.IntTensor :math:`C \in \mathcal{R}^{N_i
+        \times D}` coordinates per batch where :math:`N_i` is the number of non
+        zero elements in the :math:`i`th batch index in :math:`D` dimensional
+        space.
+
+        .. note::
+
+           The order of coordinates is non-deterministic within each batch. Use
+           :attr:`decomposed_coordinates_and_features` to retrieve both
+           coordinates features with the same order. To retrieve the order the
+           decomposed coordinates is generated, use
+           :attr:`decomposition_permutations`.
+
+        """
+        return [self.C[row_inds, 1:] for row_inds in self._batchwise_row_indices]
+
+    def coordinates_at(self, batch_index):
+        r"""Return coordinates at the specified batch index.
+
+        Returns a torch.IntTensor :math:`C \in \mathcal{R}^{N_i
+        \times D}` coordinates at the specified batch index where :math:`N_i`
+        is the number of non zero elements in the :math:`i`th batch index in
+        :math:`D` dimensional space.
+
+        .. note::
+
+           The order of coordinates is non-deterministic within each batch. Use
+           :attr:`decomposed_coordinates_and_features` to retrieve both
+           coordinates features with the same order. To retrieve the order the
+           decomposed coordinates is generated, use
+           :attr:`decomposition_permutations`.
+
+        """
+        return self.C[self._batchwise_row_indices[batch_index], 1:]
+
+    @property
+    def decomposed_features(self):
+        r"""Returns a list of features per batch.
+
+        Returns a list of torch.Tensor :math:`C \in \mathcal{R}^{N_i
+        \times N_F}` features per batch where :math:`N_i` is the number of non
+        zero elements in the :math:`i`th batch index in :math:`D` dimensional
+        space.
+
+        .. note::
+
+           The order of features is non-deterministic within each batch. Use
+           :attr:`decomposed_coordinates_and_features` to retrieve both
+           coordinates features with the same order. To retrieve the order the
+           decomposed features is generated, use
+           :attr:`decomposition_permutations`.
+
+        """
+        return [self._F[row_inds] for row_inds in self._batchwise_row_indices]
+
+    def features_at(self, batch_index):
+        r"""Returns a feature matrix at the specified batch index.
+
+        Returns a torch.Tensor :math:`C \in \mathcal{R}^{N
+        \times N_F}` feature matrix :math:`N` is the number of non
+        zero elements in the specified batch index and :math:`N_F` is the
+        number of channels.
+
+        .. note::
+
+           The order of features is non-deterministic within each batch. Use
+           :attr:`decomposed_coordinates_and_features` to retrieve both
+           coordinates features with the same order. To retrieve the order the
+           decomposed features is generated, use
+           :attr:`decomposition_permutations`.
+
+        """
+        return self._F[self._batchwise_row_indices[batch_index]]
+
+    def coordinates_and_features_at(self, batch_index):
+        r"""Returns a coordinate and feature matrix at the specified batch index.
+
+        Returns a coordinate and feature matrix at the specified `batch_index`.
+        The coordinate matrix is a torch.IntTensor :math:`C \in \mathcal{R}^{N
+        \times D}` where :math:`N` is the number of non zero elements in the
+        specified batch index in :math:`D` dimensional space. The feature
+        matrix is a torch.Tensor :math:`C \in \mathcal{R}^{N \times N_F}`
+        matrix :math:`N` is the number of non zero elements in the specified
+        batch index and :math:`N_F` is the number of channels.
+
+        .. note::
+
+           The order of features is non-deterministic within each batch. To
+           retrieve the order the decomposed features is generated, use
+           :attr:`decomposition_permutations`.
+
+        """
+        row_inds = self._batchwise_row_indices[batch_index]
+        return self.C[row_inds, 1:], self._F[row_inds]
+
+    @property
+    def decomposed_coordinates_and_features(self):
+        r"""Returns a list of coordinates and a list of features per batch.abs
+
+        .. note::
+
+           The order of decomposed coordinates and features is
+           non-deterministic within each batch. To retrieve the order the
+           decomposed features is generated, use
+           :attr:`decomposition_permutations`.
+
+        """
+        row_inds_list = self._batchwise_row_indices
+        return (
+            [self.C[row_inds, 1:] for row_inds in row_inds_list],
+            [self._F[row_inds] for row_inds in row_inds_list],
+        )
+
+    @property
+    def dimension(self):
+        r"""Alias of attr:`D`"""
+        return self._D
+
+    @dimension.setter
+    def dimension(self):
+        raise SyntaxError("Direct modification not permitted")
+
+    @property
+    def D(self):
+        r"""Alias of attr:`D`"""
+        return self._D
+
+    @D.setter
+    def D(self):
+        raise SyntaxError("Direct modification not permitted")
+
+    @property
+    def requires_grad(self):
+        return self._F.requires_grad
+
+    def requires_grad_(self, requires_grad: bool = True):
+        self._F.requires_grad_(requires_grad)
+
+    def float(self):
+        self._F = self._F.float()
+        return self
+
+    def double(self):
+        self._F = self._F.double()
+        return self
+
+    def __len__(self):
+        return len(self._F)
+
+    def size(self):
+        return self._F.size()
+
+    @property
+    def shape(self):
+        return self._F.shape
+
+    @property
+    def device(self):
+        return self._F.device
+
+    @property
+    def dtype(self):
+        return self._F.dtype
+
+    def detach(self):
+        self._F = self._F.detach()
+        return self
+
+    def get_device(self):
+        return self._F.get_device()
+
+    def _is_same_key(self, other):
+        assert isinstance(other, self.__class__)
+        assert self._manager == other._manager, COORDINATE_MANAGER_DIFFERENT_ERROR
+        assert (
+            self.coordinate_map_key == other.coordinate_map_key
+        ), COORDINATE_KEY_DIFFERENT_ERROR
+
+    # Operation overloading
+    def __iadd__(self, other):
+        self._is_same_key(other)
+        self._F += other.F
+        return self
+
+    def __isub__(self, other):
+        self._is_same_key(other)
+        self._F -= other.F
+        return self
+
+    def __imul__(self, other):
+        self._is_same_key(other)
+        self._F *= other.F
+        return self
+
+    def __idiv__(self, other):
+        self._is_same_key(other)
+        self._F /= other.F
+        return self
+
+    def _binary_functor(self, other, binary_fn):
+        assert isinstance(other, (self.__class__, torch.Tensor))
+        if isinstance(other, self.__class__):
+            assert self._manager == other._manager, COORDINATE_MANAGER_DIFFERENT_ERROR
+
+            if self.coordinate_map_key == other.coordinate_map_key:
+                return self.__class__(
+                    binary_fn(self._F, other.F),
+                    coordinate_map_key=self.coordinate_map_key,
+                    coordinate_manager=self._manager,
+                )
+            else:
+                # Generate union maps
+                out_key = CoordinateMapKey(
+                    self.coordinate_map_key.get_coordinate_size()
+                )
+                union_maps = self.coordinate_manager.union_map(
+                    [self.coordinate_map_key, other.coordinate_map_key], out_key
+                )
+                N_out = self.coordinate_manager.size(out_key)
+                out_F = torch.zeros(
+                    (N_out, self._F.size(1)), dtype=self.dtype, device=self.device
+                )
+                out_F[union_maps[0][1]] = self._F[union_maps[0][0]]
+                out_F[union_maps[1][1]] = binary_fn(
+                    out_F[union_maps[1][1]], other._F[union_maps[1][0]]
+                )
+                return self.__class__(
+                    out_F, coordinate_map_key=out_key, coordinate_manager=self._manager
+                )
+        else:  # when it is a torch.Tensor
+            return self.__class__(
+                binary_fn(self._F, other),
+                coordinate_map_key=self.coordinate_map_key,
+                coordinate_manager=self._manager,
+            )
+
+    def __add__(self, other):
+        r"""
+        Add its feature with the corresponding feature of the other
+        :attr:`MinkowskiEngine.SparseTensor` or a :attr:`torch.Tensor`
+        element-wise. For coordinates that exist on one sparse tensor but not
+        on the other, features of the counterpart that do not exist will be set
+        to 0.
+        """
+        return self._binary_functor(other, lambda x, y: x + y)
+
+    def __sub__(self, other):
+        r"""
+        Subtract the feature of the other :attr:`MinkowskiEngine.SparseTensor`
+        or a :attr:`torch.Tensor` from its corresponding feature element-wise.
+        For coordinates that exist on one sparse tensor but not on the other,
+        features of the counterpart that do not exist will be set to 0.
+        """
+        return self._binary_functor(other, lambda x, y: x - y)
+
+    def __mul__(self, other):
+        r"""
+        Multiply its feature of with the corresponding feature of the other
+        :attr:`MinkowskiEngine.SparseTensor` or a :attr:`torch.Tensor`
+        element-wise. For coordinates that exist on one sparse tensor but not
+        on the other, features of the counterpart that do not exist will be set
+        to 0.
+        """
+        return self._binary_functor(other, lambda x, y: x * y)
+
+    def __truediv__(self, other):
+        r"""
+        Divide its feature by the corresponding feature of the other
+        :attr:`MinkowskiEngine.SparseTensor` or a :attr:`torch.Tensor`
+        element-wise. For coordinates that exist on one sparse tensor but not
+        on the other, features of the counterpart that do not exist will be set
+        to 0.
+        """
+        return self._binary_functor(other, lambda x, y: x / y)
+
+    def __power__(self, power):
+        return self.__class__(
+            self._F ** power,
+            coordinate_map_key=self.coordinate_map_key,
+            coordinate_manager=self._manager,
+        )
+
+    __slots__ = (
+        "_C",
+        "_F",
+        "_D",
+        "coordinate_map_key",
+        "_manager",
+        "unique_index",
+        "inverse_mapping",
+        "quantization_mode",
+        "_batch_rows",
+    )

MinkowskiEngine/MinkowskiEngine/MinkowskiTensorField.py

@@ -0,0 +1,506 @@
+# Copyright (c) 2020 NVIDIA CORPORATION.
+#
+# 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.
+import os
+import numpy as np
+from collections.abc import Sequence
+from typing import Union, List, Tuple
+
+import torch
+from MinkowskiCommon import convert_to_int_list, StrideType
+from MinkowskiEngineBackend._C import (
+    GPUMemoryAllocatorType,
+    MinkowskiAlgorithm,
+    CoordinateMapKey,
+    CoordinateMapType,
+)
+from MinkowskiCoordinateManager import CoordinateManager
+from MinkowskiTensor import (
+    SparseTensorOperationMode,
+    SparseTensorQuantizationMode,
+    Tensor,
+    sparse_tensor_operation_mode,
+    global_coordinate_manager,
+    set_global_coordinate_manager,
+    COORDINATE_MANAGER_DIFFERENT_ERROR,
+    COORDINATE_KEY_DIFFERENT_ERROR,
+)
+from MinkowskiSparseTensor import SparseTensor
+from sparse_matrix_functions import MinkowskiSPMMFunction, MinkowskiSPMMAverageFunction
+from MinkowskiPooling import MinkowskiDirectMaxPoolingFunction
+
+
+def create_splat_coordinates(coordinates: torch.Tensor) -> torch.Tensor:
+    r"""Create splat coordinates. splat coordinates could have duplicate coordinates."""
+    dimension = coordinates.shape[1] - 1
+    region_offset = [
+        [
+            0,
+        ]
+        * (dimension + 1)
+    ]
+    for d in reversed(range(1, dimension + 1)):
+        new_offset = []
+        for offset in region_offset:
+            offset = offset.copy()  # Do not modify the original
+            offset[d] = 1
+            new_offset.append(offset)
+        region_offset.extend(new_offset)
+    region_offset = torch.IntTensor(region_offset).to(coordinates.device)
+    coordinates = torch.floor(coordinates).int().unsqueeze(1) + region_offset.unsqueeze(
+        0
+    )
+    return coordinates.reshape(-1, dimension + 1)
+
+
+class TensorField(Tensor):
+    def __init__(
+        self,
+        features: torch.Tensor,
+        coordinates: torch.Tensor = None,
+        # optional coordinate related arguments
+        tensor_stride: StrideType = 1,
+        coordinate_field_map_key: CoordinateMapKey = None,
+        coordinate_manager: CoordinateManager = None,
+        quantization_mode: SparseTensorQuantizationMode = SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE,
+        # optional manager related arguments
+        allocator_type: GPUMemoryAllocatorType = None,
+        minkowski_algorithm: MinkowskiAlgorithm = None,
+        requires_grad=None,
+        device=None,
+    ):
+        r"""
+
+        Args:
+            :attr:`features` (:attr:`torch.FloatTensor`,
+            :attr:`torch.DoubleTensor`, :attr:`torch.cuda.FloatTensor`, or
+            :attr:`torch.cuda.DoubleTensor`): The features of a sparse
+            tensor.
+
+            :attr:`coordinates` (:attr:`torch.IntTensor`): The coordinates
+            associated to the features. If not provided, :attr:`coordinate_map_key`
+            must be provided.
+
+            :attr:`tensor_stride` (:attr:`int`, :attr:`list`,
+            :attr:`numpy.array`, or :attr:`tensor.Tensor`): The tensor stride
+            of the current sparse tensor. By default, it is 1.
+
+            :attr:`coordinate_field_map_key`
+            (:attr:`MinkowskiEngine.CoordinateMapKey`): When the coordinates
+            are already cached in the MinkowskiEngine, we could reuse the same
+            coordinate map by simply providing the coordinate map key. In most
+            case, this process is done automatically. When you provide a
+            `coordinate_field_map_key`, `coordinates` will be be ignored.
+
+            :attr:`coordinate_manager`
+            (:attr:`MinkowskiEngine.CoordinateManager`): The MinkowskiEngine
+            manages all coordinate maps using the `_C.CoordinateMapManager`. If
+            not provided, the MinkowskiEngine will create a new computation
+            graph. In most cases, this process is handled automatically and you
+            do not need to use this.
+
+            :attr:`quantization_mode`
+            (:attr:`MinkowskiEngine.SparseTensorQuantizationMode`): Defines how
+            continuous coordinates will be quantized to define a sparse tensor.
+            Please refer to :attr:`SparseTensorQuantizationMode` for details.
+
+            :attr:`allocator_type`
+            (:attr:`MinkowskiEngine.GPUMemoryAllocatorType`): Defines the GPU
+            memory allocator type. By default, it uses the c10 allocator.
+
+            :attr:`minkowski_algorithm`
+            (:attr:`MinkowskiEngine.MinkowskiAlgorithm`): Controls the mode the
+            minkowski engine runs, Use
+            :attr:`MinkowskiAlgorithm.MEMORY_EFFICIENT` if you want to reduce
+            the memory footprint. Or use
+            :attr:`MinkowskiAlgorithm.SPEED_OPTIMIZED` if you want to make it
+            run fasterat the cost of more memory.
+
+            :attr:`requires_grad` (:attr:`bool`): Set the requires_grad flag.
+
+            :attr:`device` (:attr:`torch.device`): Set the device the sparse
+            tensor is defined.
+        """
+        # Type checks
+        assert isinstance(features, torch.Tensor), "Features must be a torch.Tensor"
+        assert (
+            features.ndim == 2
+        ), f"The feature should be a matrix, The input feature is an order-{features.ndim} tensor."
+        assert isinstance(quantization_mode, SparseTensorQuantizationMode)
+        assert quantization_mode in [
+            SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE,
+            SparseTensorQuantizationMode.UNWEIGHTED_SUM,
+            SparseTensorQuantizationMode.RANDOM_SUBSAMPLE,
+            SparseTensorQuantizationMode.MAX_POOL,
+        ], "invalid quantization mode"
+
+        self.quantization_mode = quantization_mode
+
+        if coordinates is not None:
+            assert isinstance(coordinates, torch.Tensor)
+        if coordinate_field_map_key is not None:
+            assert isinstance(coordinate_field_map_key, CoordinateMapKey)
+            assert coordinate_manager is not None, "Must provide coordinate_manager if coordinate_field_map_key is provided"
+            assert coordinates is None, "Must not provide coordinates if coordinate_field_map_key is provided"
+        if coordinate_manager is not None:
+            assert isinstance(coordinate_manager, CoordinateManager)
+        if coordinates is None and (
+            coordinate_field_map_key is None or coordinate_manager is None
+        ):
+            raise ValueError(
+                "Either coordinates or (coordinate_field_map_key, coordinate_manager) pair must be provided."
+            )
+
+        Tensor.__init__(self)
+
+        # To device
+        if device is not None:
+            features = features.to(device)
+            if coordinates is not None:
+                # assertion check for the map key done later
+                coordinates = coordinates.to(device)
+
+        self._D = (
+            coordinates.size(1) - 1 if coordinates is not None else coordinate_manager.D
+        )
+        ##########################
+        # Setup CoordsManager
+        ##########################
+        if coordinate_manager is None:
+            # If set to share the coords man, use the global coords man
+            if (
+                sparse_tensor_operation_mode()
+                == SparseTensorOperationMode.SHARE_COORDINATE_MANAGER
+            ):
+                coordinate_manager = global_coordinate_manager()
+                if coordinate_manager is None:
+                    coordinate_manager = CoordinateManager(
+                        D=self._D,
+                        coordinate_map_type=CoordinateMapType.CUDA
+                        if coordinates.is_cuda
+                        else CoordinateMapType.CPU,
+                        allocator_type=allocator_type,
+                        minkowski_algorithm=minkowski_algorithm,
+                    )
+                    set_global_coordinate_manager(coordinate_manager)
+            else:
+                coordinate_manager = CoordinateManager(
+                    D=coordinates.size(1) - 1,
+                    coordinate_map_type=CoordinateMapType.CUDA
+                    if coordinates.is_cuda
+                    else CoordinateMapType.CPU,
+                    allocator_type=allocator_type,
+                    minkowski_algorithm=minkowski_algorithm,
+                )
+        self._manager = coordinate_manager
+
+        ##########################
+        # Initialize coords
+        ##########################
+        # Coordinate Management
+        if coordinates is not None:
+            assert (
+                features.shape[0] == coordinates.shape[0]
+            ), "The number of rows in features and coordinates must match."
+
+            assert (
+                features.is_cuda == coordinates.is_cuda
+            ), "Features and coordinates must have the same backend."
+
+            coordinate_field_map_key = CoordinateMapKey(
+                convert_to_int_list(tensor_stride, self._D), ""
+            )
+            coordinate_field_map_key = self._manager.insert_field(
+                coordinates.float(), convert_to_int_list(tensor_stride, self._D), ""
+            )
+        else:
+            assert (
+                coordinate_field_map_key.is_key_set()
+            ), "The coordinate field map key must be valid."
+
+        if requires_grad is not None:
+            features.requires_grad_(requires_grad)
+
+        self._F = features
+        self._C = coordinates
+        self.coordinate_field_map_key = coordinate_field_map_key
+        self._batch_rows = None
+        self._inverse_mapping = {}
+        self._splat = {}
+
+    @property
+    def coordinate_key(self):
+        return self.coordinate_field_map_key
+
+    @property
+    def C(self):
+        r"""The alias of :attr:`coords`."""
+        return self.coordinates
+
+    @property
+    def coordinates(self):
+        r"""
+        The coordinates of the current sparse tensor. The coordinates are
+        represented as a :math:`N \times (D + 1)` dimensional matrix where
+        :math:`N` is the number of points in the space and :math:`D` is the
+        dimension of the space (e.g. 3 for 3D, 4 for 3D + Time). Additional
+        dimension of the column of the matrix C is for batch indices which is
+        internally treated as an additional spatial dimension to disassociate
+        different instances in a batch.
+        """
+        if self._C is None:
+            self._C = self._get_coordinate_field()
+        return self._C
+
+    @property
+    def _batchwise_row_indices(self):
+        if self._batch_rows is None:
+            _, self._batch_rows = self._manager.origin_field_map(
+                self.coordinate_field_map_key
+            )
+        return self._batch_rows
+
+    def _get_coordinate_field(self):
+        return self._manager.get_coordinate_field(self.coordinate_field_map_key)
+
+    def sparse(
+        self,
+        tensor_stride: Union[int, Sequence, np.array] = 1,
+        coordinate_map_key: CoordinateMapKey = None,
+        quantization_mode: SparseTensorQuantizationMode = None,
+    ):
+        r"""Converts the current sparse tensor field to a sparse tensor."""
+        if quantization_mode is None:
+            quantization_mode = self.quantization_mode
+        assert (
+            quantization_mode != SparseTensorQuantizationMode.SPLAT_LINEAR_INTERPOLATION
+        ), "Please use .splat() for splat quantization."
+
+        if coordinate_map_key is None:
+            tensor_stride = convert_to_int_list(tensor_stride, self.D)
+
+            coordinate_map_key, (
+                unique_index,
+                inverse_mapping,
+            ) = self._manager.field_to_sparse_insert_and_map(
+                self.coordinate_field_map_key,
+                tensor_stride,
+            )
+            N_rows = len(unique_index)
+        else:
+            # sparse index, field index
+            inverse_mapping, unique_index = self._manager.field_to_sparse_map(
+                self.coordinate_field_map_key,
+                coordinate_map_key,
+            )
+            N_rows = self._manager.size(coordinate_map_key)
+
+        assert N_rows > 0, f"Invalid out coordinate map key. Found {N_row} elements."
+
+        if len(inverse_mapping) == 0:
+            # When the input has the same shape as the output
+            self._inverse_mapping[coordinate_map_key] = torch.arange(
+                len(self._F),
+                dtype=inverse_mapping.dtype,
+                device=inverse_mapping.device,
+            )
+            return SparseTensor(
+                self._F,
+                coordinate_map_key=coordinate_map_key,
+                coordinate_manager=self._manager,
+            )
+
+        # Create features
+        if quantization_mode == SparseTensorQuantizationMode.UNWEIGHTED_SUM:
+            N = len(self._F)
+            cols = torch.arange(
+                N,
+                dtype=inverse_mapping.dtype,
+                device=inverse_mapping.device,
+            )
+            vals = torch.ones(N, dtype=self._F.dtype, device=self._F.device)
+            size = torch.Size([N_rows, len(inverse_mapping)])
+            features = MinkowskiSPMMFunction().apply(
+                inverse_mapping, cols, vals, size, self._F
+            )
+        elif quantization_mode == SparseTensorQuantizationMode.UNWEIGHTED_AVERAGE:
+            N = len(self._F)
+            cols = torch.arange(
+                N,
+                dtype=inverse_mapping.dtype,
+                device=inverse_mapping.device,
+            )
+            size = torch.Size([N_rows, len(inverse_mapping)])
+            features = MinkowskiSPMMAverageFunction().apply(
+                inverse_mapping, cols, size, self._F
+            )
+        elif quantization_mode == SparseTensorQuantizationMode.RANDOM_SUBSAMPLE:
+            features = self._F[unique_index]
+        elif quantization_mode == SparseTensorQuantizationMode.MAX_POOL:
+            N = len(self._F)
+            in_map = torch.arange(
+                N,
+                dtype=inverse_mapping.dtype,
+                device=inverse_mapping.device,
+            )
+            features = MinkowskiDirectMaxPoolingFunction().apply(
+                in_map, inverse_mapping, self._F, N_rows
+            )
+        else:
+            # No quantization
+            raise ValueError("Invalid quantization mode")
+
+        self._inverse_mapping[coordinate_map_key] = inverse_mapping
+
+        return SparseTensor(
+            features,
+            coordinate_map_key=coordinate_map_key,
+            coordinate_manager=self._manager,
+        )
+
+    def splat(self):
+        r"""
+        For slice, use Y.slice(X) where X is the tensor field and Y is the
+        resulting sparse tensor.
+        """
+        splat_coordinates = create_splat_coordinates(self.C)
+        (coordinate_map_key, _) = self._manager.insert_and_map(splat_coordinates)
+        N_rows = self._manager.size(coordinate_map_key)
+
+        tensor_map, field_map, weights = self._manager.interpolation_map_weight(
+            coordinate_map_key, self._C
+        )
+        # features
+        N = len(self._F)
+        assert weights.dtype == self._F.dtype
+        size = torch.Size([N_rows, N])
+        # Save the results for slice
+        self._splat[coordinate_map_key] = (tensor_map, field_map, weights, size)
+        features = MinkowskiSPMMFunction().apply(
+            tensor_map, field_map, weights, size, self._F
+        )
+        return SparseTensor(
+            features,
+            coordinate_map_key=coordinate_map_key,
+            coordinate_manager=self._manager,
+        )
+
+    def inverse_mapping(self, sparse_tensor_map_key: CoordinateMapKey):
+        if sparse_tensor_map_key not in self._inverse_mapping:
+            if not self._manager.exists_field_to_sparse(
+                self.coordinate_field_map_key, sparse_tensor_map_key
+            ):
+                sparse_keys = self.coordinate_manager.field_to_sparse_keys(
+                    self.coordinate_field_map_key
+                )
+                one_key = None
+                if len(sparse_keys) > 0:
+                    for key in sparse_keys:
+                        if np.prod(key.get_tensor_stride()) == 1:
+                            one_key = key
+                else:
+                    one_key = CoordinateMapKey(
+                        [
+                            1,
+                        ]
+                        * self.D,
+                        "",
+                    )
+
+                if one_key not in self._inverse_mapping:
+                    (
+                        _,
+                        self._inverse_mapping[one_key],
+                    ) = self._manager.get_field_to_sparse_map(
+                        self.coordinate_field_map_key, one_key
+                    )
+                _, stride_map = self.coordinate_manager.stride_map(
+                    one_key, sparse_tensor_map_key
+                )
+                field_map = self._inverse_mapping[one_key]
+                self._inverse_mapping[sparse_tensor_map_key] = stride_map[field_map]
+            else:
+                # Extract the mapping
+                (
+                    _,
+                    self._inverse_mapping[sparse_tensor_map_key],
+                ) = self._manager.get_field_to_sparse_map(
+                    self.coordinate_field_map_key, sparse_tensor_map_key
+                )
+        return self._inverse_mapping[sparse_tensor_map_key]
+
+    def _is_same_key(self, other):
+        assert isinstance(other, self.__class__)
+        assert self._manager == other._manager, COORDINATE_MANAGER_DIFFERENT_ERROR
+        assert (
+            self.coordinate_field_map_key == other.coordinate_field_map_key
+        ), COORDINATE_KEY_DIFFERENT_ERROR
+
+    def _binary_functor(self, other, binary_fn):
+        assert isinstance(other, (self.__class__, torch.Tensor))
+        if isinstance(other, self.__class__):
+            self._is_same_key(other)
+            return self.__class__(
+                binary_fn(self._F, other.F),
+                coordinate_map_key=self.coordinate_map_key,
+                coordinate_manager=self._manager,
+            )
+        else:  # when it is a torch.Tensor
+            return self.__class__(
+                binary_fn(self._F, other),
+                coordinate_field_map_key=self.coordinate_map_key,
+                coordinate_manager=self._manager,
+            )
+
+    def __repr__(self):
+        return (
+            self.__class__.__name__
+            + "("
+            + os.linesep
+            + "  coordinates="
+            + str(self.C)
+            + os.linesep
+            + "  features="
+            + str(self.F)
+            + os.linesep
+            + "  coordinate_field_map_key="
+            + str(self.coordinate_field_map_key)
+            + os.linesep
+            + "  coordinate_manager="
+            + str(self._manager)
+            + "  spatial dimension="
+            + str(self._D)
+            + ")"
+        )
+
+    __slots__ = (
+        "_C",
+        "_F",
+        "_D",
+        "coordinate_field_map_key",
+        "_manager",
+        "quantization_mode",
+        "_inverse_mapping",
+        "_batch_rows",
+        "_splat",
+    )

MinkowskiEngine/MinkowskiEngine/MinkowskiUnion.py

@@ -0,0 +1,156 @@
+# Copyright (c) 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.
+import torch
+from torch.nn import Module
+from torch.autograd import Function
+
+from MinkowskiEngineBackend._C import CoordinateMapKey
+from MinkowskiSparseTensor import SparseTensor
+from MinkowskiCoordinateManager import CoordinateManager
+
+
+class MinkowskiUnionFunction(Function):
+    @staticmethod
+    def forward(
+        ctx,
+        in_coords_keys: list,
+        out_coords_key: CoordinateMapKey,
+        coordinate_manager: CoordinateManager,
+        *in_feats,
+    ):
+        assert isinstance(
+            in_feats, (list, tuple)
+        ), "Input must be a collection of Tensors"
+        assert len(in_feats) > 1, "input must be a set with at least 2 Tensors"
+        assert len(in_feats) == len(
+            in_coords_keys
+        ), "The input features and keys must have the same length"
+
+        union_maps = coordinate_manager.union_map(in_coords_keys, out_coords_key)
+        out_feat = torch.zeros(
+            (coordinate_manager.size(out_coords_key), in_feats[0].shape[1]),
+            dtype=in_feats[0].dtype,
+            device=in_feats[0].device,
+        )
+        for in_feat, union_map in zip(in_feats, union_maps):
+            out_feat[union_map[1]] += in_feat[union_map[0]]
+        ctx.keys = (in_coords_keys, coordinate_manager)
+        ctx.save_for_backward(*union_maps)
+        return out_feat
+
+    @staticmethod
+    def backward(ctx, grad_out_feat):
+        if not grad_out_feat.is_contiguous():
+            grad_out_feat = grad_out_feat.contiguous()
+
+        union_maps = ctx.saved_tensors
+        in_coords_keys, coordinate_manager = ctx.keys
+        num_ch, dtype, device = (
+            grad_out_feat.shape[1],
+            grad_out_feat.dtype,
+            grad_out_feat.device,
+        )
+        grad_in_feats = []
+        for in_coords_key, union_map in zip(in_coords_keys, union_maps):
+            grad_in_feat = torch.zeros(
+                (coordinate_manager.size(in_coords_key), num_ch),
+                dtype=dtype,
+                device=device,
+            )
+            grad_in_feat[union_map[0]] = grad_out_feat[union_map[1]]
+            grad_in_feats.append(grad_in_feat)
+        return (None, None, None, *grad_in_feats)
+
+
+class MinkowskiUnion(Module):
+    r"""Create a union of all sparse tensors and add overlapping features.
+
+    Args:
+        None
+
+    .. warning::
+       This function is experimental and the usage can be changed in the future updates.
+
+    """
+
+    def __init__(self):
+        super(MinkowskiUnion, self).__init__()
+        self.union = MinkowskiUnionFunction()
+
+    def forward(self, *inputs):
+        r"""
+        Args:
+            A variable number of :attr:`MinkowskiEngine.SparseTensor`'s.
+
+        Returns:
+            A :attr:`MinkowskiEngine.SparseTensor` with coordinates = union of all
+            input coordinates, and features = sum of all features corresponding to the
+            coordinate.
+
+        Example::
+
+            >>> # Define inputs
+            >>> input1 = SparseTensor(
+            >>>     torch.rand(N, in_channels, dtype=torch.double), coords=coords)
+            >>> # All inputs must share the same coordinate manager
+            >>> input2 = SparseTensor(
+            >>>     torch.rand(N, in_channels, dtype=torch.double),
+            >>>     coords=coords + 1,
+            >>>     coords_manager=input1.coordinate_manager,  # Must use same coords manager
+            >>>     force_creation=True  # The tensor stride [1, 1] already exists.
+            >>> )
+            >>> union = MinkowskiUnion()
+            >>> output = union(input1, iput2)
+
+        """
+        assert isinstance(inputs, (list, tuple)), "The input must be a list or tuple"
+        for s in inputs:
+            assert isinstance(s, SparseTensor), "Inputs must be sparse tensors."
+        assert len(inputs) > 1, "input must be a set with at least 2 SparseTensors"
+        # Assert the same coordinate manager
+        ref_coordinate_manager = inputs[0].coordinate_manager
+        for s in inputs:
+            assert (
+                ref_coordinate_manager == s.coordinate_manager
+            ), "Invalid coordinate manager. All inputs must have the same coordinate manager."
+
+        in_coordinate_map_key = inputs[0].coordinate_map_key
+        coordinate_manager = inputs[0].coordinate_manager
+        out_coordinate_map_key = CoordinateMapKey(
+            in_coordinate_map_key.get_coordinate_size()
+        )
+        output = self.union.apply(
+            [input.coordinate_map_key for input in inputs],
+            out_coordinate_map_key,
+            coordinate_manager,
+            *[input.F for input in inputs],
+        )
+        return SparseTensor(
+            output,
+            coordinate_map_key=out_coordinate_map_key,
+            coordinate_manager=coordinate_manager,
+        )
+
+    def __repr__(self):
+        return self.__class__.__name__ + "()"

MinkowskiEngine/MinkowskiEngine/__init__.py

@@ -0,0 +1,228 @@
+# 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.
+__version__ = "0.5.4"
+
+import os
+import sys
+import warnings
+
+file_dir = os.path.dirname(__file__)
+sys.path.append(file_dir)
+
+# Force OMP_NUM_THREADS setup
+if os.cpu_count() > 16 and "OMP_NUM_THREADS" not in os.environ:
+    warnings.warn(
+        " ".join(
+            [
+                "The environment variable `OMP_NUM_THREADS` not set. MinkowskiEngine will automatically set `OMP_NUM_THREADS=16`.",
+                "If you want to set `OMP_NUM_THREADS` manually, please export it on the command line before running a python script.",
+                "e.g. `export OMP_NUM_THREADS=12; python your_program.py`.",
+                "It is recommended to set it below 24.",
+            ]
+        )
+    )
+    os.environ["OMP_NUM_THREADS"] = str(16)
+
+# Must be imported first to load all required shared libs
+import torch
+
+from diagnostics import print_diagnostics
+
+from MinkowskiEngineBackend._C import (
+    MinkowskiAlgorithm,
+    CoordinateMapKey,
+    GPUMemoryAllocatorType,
+    CoordinateMapType,
+    RegionType,
+    PoolingMode,
+    BroadcastMode,
+    is_cuda_available,
+    cuda_version,
+    cudart_version,
+    get_gpu_memory_info,
+)
+
+from MinkowskiKernelGenerator import (
+    KernelRegion,
+    KernelGenerator,
+    convert_region_type,
+    get_kernel_volume,
+)
+
+from MinkowskiTensor import (
+    SparseTensorOperationMode,
+    SparseTensorQuantizationMode,
+    set_sparse_tensor_operation_mode,
+    sparse_tensor_operation_mode,
+    global_coordinate_manager,
+    set_global_coordinate_manager,
+    clear_global_coordinate_manager,
+)
+
+from MinkowskiSparseTensor import SparseTensor
+
+from MinkowskiTensorField import TensorField
+
+from MinkowskiCommon import (
+    convert_to_int_tensor,
+    MinkowskiModuleBase,
+)
+
+from MinkowskiCoordinateManager import (
+    set_memory_manager_backend,
+    set_gpu_allocator,
+    CoordsManager,
+    CoordinateManager,
+)
+
+from MinkowskiConvolution import (
+    MinkowskiConvolutionFunction,
+    MinkowskiConvolution,
+    MinkowskiConvolutionTransposeFunction,
+    MinkowskiConvolutionTranspose,
+    MinkowskiGenerativeConvolutionTranspose,
+)
+
+
+from MinkowskiChannelwiseConvolution import MinkowskiChannelwiseConvolution
+
+from MinkowskiPooling import (
+    MinkowskiLocalPoolingFunction,
+    MinkowskiSumPooling,
+    MinkowskiAvgPooling,
+    MinkowskiMaxPooling,
+    MinkowskiLocalPoolingTransposeFunction,
+    MinkowskiPoolingTranspose,
+    MinkowskiGlobalPoolingFunction,
+    MinkowskiGlobalPooling,
+    MinkowskiGlobalSumPooling,
+    MinkowskiGlobalAvgPooling,
+    MinkowskiGlobalMaxPooling,
+    MinkowskiDirectMaxPoolingFunction,
+)
+
+from MinkowskiBroadcast import (
+    MinkowskiBroadcastFunction,
+    MinkowskiBroadcastAddition,
+    MinkowskiBroadcastMultiplication,
+    MinkowskiBroadcast,
+    MinkowskiBroadcastConcatenation,
+)
+
+from MinkowskiNonlinearity import (
+    MinkowskiELU,
+    MinkowskiHardshrink,
+    MinkowskiHardsigmoid,
+    MinkowskiHardtanh,
+    MinkowskiHardswish,
+    MinkowskiLeakyReLU,
+    MinkowskiLogSigmoid,
+    MinkowskiPReLU,
+    MinkowskiReLU,
+    MinkowskiReLU6,
+    MinkowskiRReLU,
+    MinkowskiSELU,
+    MinkowskiCELU,
+    MinkowskiGELU,
+    MinkowskiSigmoid,
+    MinkowskiSiLU,
+    MinkowskiSoftplus,
+    MinkowskiSoftshrink,
+    MinkowskiSoftsign,
+    MinkowskiTanh,
+    MinkowskiTanhshrink,
+    MinkowskiThreshold,
+    MinkowskiSoftmin,
+    MinkowskiSoftmax,
+    MinkowskiLogSoftmax,
+    MinkowskiAdaptiveLogSoftmaxWithLoss,
+    MinkowskiDropout,
+    MinkowskiAlphaDropout,
+    MinkowskiSinusoidal,
+)
+
+from MinkowskiNormalization import (
+    MinkowskiBatchNorm,
+    MinkowskiSyncBatchNorm,
+    MinkowskiInstanceNorm,
+    MinkowskiInstanceNormFunction,
+    MinkowskiStableInstanceNorm,
+)
+
+
+from MinkowskiPruning import MinkowskiPruning, MinkowskiPruningFunction
+
+from MinkowskiUnion import MinkowskiUnion, MinkowskiUnionFunction
+
+from MinkowskiInterpolation import (
+    MinkowskiInterpolation,
+    MinkowskiInterpolationFunction,
+)
+
+from MinkowskiNetwork import MinkowskiNetwork
+
+import MinkowskiOps
+
+from MinkowskiOps import (
+    MinkowskiLinear,
+    MinkowskiToSparseTensor,
+    MinkowskiToDenseTensor,
+    MinkowskiToFeature,
+    MinkowskiStackCat,
+    MinkowskiStackSum,
+    MinkowskiStackMean,
+    MinkowskiStackVar,
+    cat,
+    mean,
+    var,
+    to_sparse,
+    to_sparse_all,
+    dense_coordinates,
+)
+
+from MinkowskiOps import _sum as sum
+
+import MinkowskiFunctional
+
+import MinkowskiEngine.utils as utils
+
+import MinkowskiEngine.modules as modules
+
+from sparse_matrix_functions import (
+    spmm,
+    MinkowskiSPMMFunction,
+    MinkowskiSPMMAverageFunction,
+)
+
+
+if not is_cuda_available():
+    warnings.warn(
+        " ".join(
+            [
+                "The MinkowskiEngine was compiled with CPU_ONLY flag.",
+                "If you want to compile with CUDA support, make sure `torch.cuda.is_available()` is True when you install MinkowskiEngine.",
+            ]
+        )
+    )

MinkowskiEngine/MinkowskiEngine/diagnostics.py

@@ -0,0 +1,70 @@
+import sys
+import os
+import platform
+import subprocess
+
+
+def parse_nvidia_smi():
+    sp = subprocess.Popen(
+        ["nvidia-smi", "-q"], stdout=subprocess.PIPE, stderr=subprocess.PIPE
+    )
+    out_dict = dict()
+    for item in sp.communicate()[0].decode("utf-8").split("\n"):
+        if item.count(":") == 1:
+            key, val = [i.strip() for i in item.split(":")]
+            out_dict[key] = val
+    return out_dict
+
+
+def print_diagnostics():
+    print("==========System==========")
+    print(platform.platform())
+    os.system("cat /etc/lsb-release")
+    print(sys.version)
+
+    print("==========Pytorch==========")
+    try:
+        import torch
+
+        print(torch.__version__)
+        print(f"torch.cuda.is_available(): {torch.cuda.is_available()}")
+    except ImportError:
+        print("torch not installed")
+
+    print("==========NVIDIA-SMI==========")
+    os.system("which nvidia-smi")
+    for k, v in parse_nvidia_smi().items():
+        if "version" in k.lower():
+            print(k, v)
+
+    print("==========NVCC==========")
+    os.system("which nvcc")
+    os.system("nvcc --version")
+
+    print("==========CC==========")
+    CC = "c++"
+    if "CC" in os.environ or "CXX" in os.environ:
+        # distutils only checks CC not CXX
+        if "CXX" in os.environ:
+            os.environ["CC"] = os.environ["CXX"]
+            CC = os.environ["CXX"]
+        else:
+            CC = os.environ["CC"]
+        print(f"CC={CC}")
+    os.system(f"which {CC}")
+    os.system(f"{CC} --version")
+
+    print("==========MinkowskiEngine==========")
+    try:
+        import MinkowskiEngine as ME
+
+        print(ME.__version__)
+        print(f"MinkowskiEngine compiled with CUDA Support: {ME.is_cuda_available()}")
+        print(f"NVCC version MinkowskiEngine is compiled: {ME.cuda_version()}")
+        print(f"CUDART version MinkowskiEngine is compiled: {ME.cudart_version()}")
+    except ImportError:
+        print("MinkowskiEngine not installed")
+
+
+if __name__ == "__main__":
+    print_diagnostics()

MinkowskiEngine/MinkowskiEngine/modules/__init__.py

@@ -0,0 +1,23 @@
+# Copyright (c) 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.

MinkowskiEngine/MinkowskiEngine/modules/resnet_block.py

@@ -0,0 +1,121 @@
+# Copyright (c) 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.
+import torch.nn as nn
+
+import MinkowskiEngine as ME
+
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self,
+                 inplanes,
+                 planes,
+                 stride=1,
+                 dilation=1,
+                 downsample=None,
+                 bn_momentum=0.1,
+                 dimension=-1):
+        super(BasicBlock, self).__init__()
+        assert dimension > 0
+
+        self.conv1 = ME.MinkowskiConvolution(
+            inplanes, planes, kernel_size=3, stride=stride, dilation=dilation, dimension=dimension)
+        self.norm1 = ME.MinkowskiBatchNorm(planes, momentum=bn_momentum)
+        self.conv2 = ME.MinkowskiConvolution(
+            planes, planes, kernel_size=3, stride=1, dilation=dilation, dimension=dimension)
+        self.norm2 = ME.MinkowskiBatchNorm(planes, momentum=bn_momentum)
+        self.relu = ME.MinkowskiReLU(inplace=True)
+        self.downsample = downsample
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.norm1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.norm2(out)
+
+        if self.downsample is not None:
+          residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    expansion = 4
+
+    def __init__(self,
+                 inplanes,
+                 planes,
+                 stride=1,
+                 dilation=1,
+                 downsample=None,
+                 bn_momentum=0.1,
+                 dimension=-1):
+        super(Bottleneck, self).__init__()
+        assert dimension > 0
+
+        self.conv1 = ME.MinkowskiConvolution(
+            inplanes, planes, kernel_size=1, dimension=dimension)
+        self.norm1 = ME.MinkowskiBatchNorm(planes, momentum=bn_momentum)
+
+        self.conv2 = ME.MinkowskiConvolution(
+            planes, planes, kernel_size=3, stride=stride, dilation=dilation, dimension=dimension)
+        self.norm2 = ME.MinkowskiBatchNorm(planes, momentum=bn_momentum)
+
+        self.conv3 = ME.MinkowskiConvolution(
+            planes, planes * self.expansion, kernel_size=1, dimension=dimension)
+        self.norm3 = ME.MinkowskiBatchNorm(
+            planes * self.expansion, momentum=bn_momentum)
+
+        self.relu = ME.MinkowskiReLU(inplace=True)
+        self.downsample = downsample
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.norm1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.norm2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.norm3(out)
+
+        if self.downsample is not None:
+          residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out

MinkowskiEngine/MinkowskiEngine/modules/senet_block.py

@@ -0,0 +1,129 @@
+# Copyright (c) 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.
+import torch.nn as nn
+
+import MinkowskiEngine as ME
+
+from MinkowskiEngine.modules.resnet_block import BasicBlock, Bottleneck
+
+
+class SELayer(nn.Module):
+
+    def __init__(self, channel, reduction=16, D=-1):
+        # Global coords does not require coords_key
+        super(SELayer, self).__init__()
+        self.fc = nn.Sequential(
+            ME.MinkowskiLinear(channel, channel // reduction),
+            ME.MinkowskiReLU(inplace=True),
+            ME.MinkowskiLinear(channel // reduction, channel),
+            ME.MinkowskiSigmoid())
+        self.pooling = ME.MinkowskiGlobalPooling()
+        self.broadcast_mul = ME.MinkowskiBroadcastMultiplication()
+
+    def forward(self, x):
+        y = self.pooling(x)
+        y = self.fc(y)
+        return self.broadcast_mul(x, y)
+
+
+class SEBasicBlock(BasicBlock):
+
+    def __init__(self,
+                 inplanes,
+                 planes,
+                 stride=1,
+                 dilation=1,
+                 downsample=None,
+                 reduction=16,
+                 D=-1):
+        super(SEBasicBlock, self).__init__(
+            inplanes,
+            planes,
+            stride=stride,
+            dilation=dilation,
+            downsample=downsample,
+            D=D)
+        self.se = SELayer(planes, reduction=reduction, D=D)
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.norm1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.norm2(out)
+        out = self.se(out)
+
+        if self.downsample is not None:
+          residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class SEBottleneck(Bottleneck):
+
+    def __init__(self,
+                 inplanes,
+                 planes,
+                 stride=1,
+                 dilation=1,
+                 downsample=None,
+                 D=3,
+                 reduction=16):
+        super(SEBottleneck, self).__init__(
+            inplanes,
+            planes,
+            stride=stride,
+            dilation=dilation,
+            downsample=downsample,
+            D=D)
+        self.se = SELayer(planes * self.expansion, reduction=reduction, D=D)
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.norm1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.norm2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.norm3(out)
+        out = self.se(out)
+
+        if self.downsample is not None:
+          residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out

MinkowskiEngine/MinkowskiEngine/sparse_matrix_functions.py

@@ -0,0 +1,213 @@
+# Copyright (c) 2020 NVIDIA CORPORATION.
+#
+# 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.
+import torch
+from torch.autograd import Function
+
+import MinkowskiEngineBackend._C as MEB
+
+EPS = 1e-10
+
+
+def spmm(
+    rows: torch.Tensor,
+    cols: torch.Tensor,
+    vals: torch.Tensor,
+    size: torch.Size,
+    mat: torch.Tensor,
+    is_sorted: bool = False,
+    cuda_spmm_alg: int = 1,
+) -> torch.Tensor:
+
+    assert len(rows) == len(cols), "Invalid length"
+    assert len(rows) == len(vals), "Invalid length"
+    assert vals.dtype == mat.dtype, "dtype mismatch"
+    assert vals.device == mat.device, "device mismatch"
+    if mat.is_cuda:
+        assert (
+            rows.is_cuda and cols.is_cuda and vals.is_cuda
+        ), "All inputs must be on cuda"
+        rows = rows.int()
+        cols = cols.int()
+        result = MEB.coo_spmm_int32(
+            rows, cols, vals, size[0], size[1], mat, cuda_spmm_alg, is_sorted
+        )
+
+        # WARNING: TODO: not sorting the vals. Should not be used for generic SPMM
+        # coosort only supports int32
+        # return MEB.coo_spmm_int64(
+        #     rows, cols, vals, size[0], size[1], mat, cuda_spmm_alg
+        # )
+    else:
+        COO = torch.stack(
+            (rows, cols),
+            0,
+        ).long()
+        torchSparseTensor = None
+        if vals.dtype == torch.float64:
+            torchSparseTensor = torch.sparse.DoubleTensor
+        elif vals.dtype == torch.float32:
+            torchSparseTensor = torch.sparse.FloatTensor
+        else:
+            raise ValueError(f"Unsupported data type: {vals.dtype}")
+
+        sp = torchSparseTensor(COO, vals, size)
+        result = sp.matmul(mat)
+
+    return result
+
+
+def spmm_average(
+    rows: torch.Tensor,
+    cols: torch.Tensor,
+    size: torch.Size,
+    mat: torch.Tensor,
+    cuda_spmm_alg: int = 1,
+) -> (torch.Tensor, torch.Tensor, torch.Tensor):
+
+    assert len(rows) == len(cols), "Invalid length"
+    if mat.is_cuda:
+        assert rows.is_cuda and cols.is_cuda, "All inputs must be on cuda"
+        rows = rows.int()
+        cols = cols.int()
+        result, COO, vals = MEB.coo_spmm_average_int32(
+            rows, cols, size[0], size[1], mat, cuda_spmm_alg
+        )
+
+        # WARNING: TODO: not sorting the vals. Should not be used for generic SPMM
+        # coosort only supports int32
+        # return MEB.coo_spmm_int64(
+        #     rows, cols, vals, size[0], size[1], mat, cuda_spmm_alg
+        # )
+    else:
+        # fmt: off
+        rows, sort_ind = torch.sort(rows)
+        cols = cols[sort_ind]
+        COO = torch.stack((rows, cols), 0,).long()
+        # Vals
+        _, inverse_ind, counts = torch.unique(rows, return_counts=True, return_inverse=True)
+        vals = (1 / counts[inverse_ind]).to(mat.dtype)
+        # fmt: on
+        torchSparseTensor = None
+        if mat.dtype == torch.float64:
+            torchSparseTensor = torch.sparse.DoubleTensor
+        elif mat.dtype == torch.float32:
+            torchSparseTensor = torch.sparse.FloatTensor
+        else:
+            raise ValueError(f"Unsupported data type: {mat.dtype}")
+        sp = torchSparseTensor(COO, vals, size)
+        result = sp.matmul(mat)
+
+    return result, COO, vals
+
+
+class MinkowskiSPMMFunction(Function):
+    @staticmethod
+    def forward(
+        ctx,
+        rows: torch.Tensor,
+        cols: torch.Tensor,
+        vals: torch.Tensor,
+        size: torch.Size,
+        mat: torch.Tensor,
+        cuda_spmm_alg: int = 1,
+    ):
+        ctx.misc_args = size, cuda_spmm_alg
+        ctx.save_for_backward(rows, cols, vals)
+        result = spmm(
+            rows,
+            cols,
+            vals,
+            size,
+            mat,
+            is_sorted=False,
+            cuda_spmm_alg=cuda_spmm_alg,
+        )
+        return result
+
+    @staticmethod
+    def backward(ctx, grad: torch.Tensor):
+        size, cuda_spmm_alg = ctx.misc_args
+        rows, cols, vals = ctx.saved_tensors
+        new_size = torch.Size([size[1], size[0]])
+        grad = spmm(
+            cols,
+            rows,
+            vals,
+            new_size,
+            grad,
+            is_sorted=False,
+            cuda_spmm_alg=cuda_spmm_alg,
+        )
+        return (
+            None,
+            None,
+            None,
+            None,
+            grad,
+            None,
+        )
+
+
+class MinkowskiSPMMAverageFunction(Function):
+    @staticmethod
+    def forward(
+        ctx,
+        rows: torch.Tensor,
+        cols: torch.Tensor,
+        size: torch.Size,
+        mat: torch.Tensor,
+        cuda_spmm_alg: int = 1,
+    ):
+        ctx.misc_args = size, cuda_spmm_alg
+        result, COO, vals = spmm_average(
+            rows,
+            cols,
+            size,
+            mat,
+            cuda_spmm_alg=cuda_spmm_alg,
+        )
+        ctx.save_for_backward(COO, vals)
+        return result
+
+    @staticmethod
+    def backward(ctx, grad: torch.Tensor):
+        size, cuda_spmm_alg = ctx.misc_args
+        COO, vals = ctx.saved_tensors
+        new_size = torch.Size([size[1], size[0]])
+        grad = spmm(
+            COO[1],
+            COO[0],
+            vals,
+            new_size,
+            grad,
+            is_sorted=False,
+            cuda_spmm_alg=cuda_spmm_alg,
+        )
+        return (
+            None,
+            None,
+            None,
+            grad,
+            None,
+        )

MinkowskiEngine/MinkowskiEngine/utils/__init__.py

@@ -0,0 +1,28 @@
+# Copyright (c) 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.
+from .quantization import sparse_quantize, ravel_hash_vec, fnv_hash_vec, unique_coordinate_map
+from .collation import SparseCollation, batched_coordinates, sparse_collate, batch_sparse_collate
+# from .coords import get_coords_map
+from .init import kaiming_normal_
+from .summary import summary

MinkowskiEngine/MinkowskiEngine/utils/collation.py

@@ -0,0 +1,263 @@
+# Copyright (c) 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.
+import numpy as np
+import torch
+import logging
+import collections.abc
+
+
+def batched_coordinates(coords, dtype=torch.int32, device=None):
+    r"""Create a `ME.SparseTensor` coordinates from a sequence of coordinates
+
+    Given a list of either numpy or pytorch tensor coordinates, return the
+    batched coordinates suitable for `ME.SparseTensor`.
+
+    Args:
+        :attr:`coords` (a sequence of `torch.Tensor` or `numpy.ndarray`): a
+        list of coordinates.
+
+        :attr:`dtype`: torch data type of the return tensor. torch.int32 by default.
+
+    Returns:
+        :attr:`batched_coordindates` (`torch.Tensor`): a batched coordinates.
+
+    .. warning::
+
+       From v0.4, the batch index will be prepended before all coordinates.
+
+    """
+    assert isinstance(
+        coords, collections.abc.Sequence
+    ), "The coordinates must be a sequence."
+    assert np.array(
+        [cs.ndim == 2 for cs in coords]
+    ).all(), "All coordinates must be in a 2D array."
+    D = np.unique(np.array([cs.shape[1] for cs in coords]))
+    assert len(D) == 1, f"Dimension of the array mismatch. All dimensions: {D}"
+    D = D[0]
+    if device is None:
+        if isinstance(coords, torch.Tensor):
+            device = coords[0].device
+        else:
+            device = "cpu"
+    assert dtype in [
+        torch.int32,
+        torch.float32,
+    ], "Only torch.int32, torch.float32 supported for coordinates."
+
+    # Create a batched coordinates
+    N = np.array([len(cs) for cs in coords]).sum()
+    bcoords = torch.zeros((N, D + 1), dtype=dtype, device=device)  # uninitialized
+
+    s = 0
+    for b, cs in enumerate(coords):
+        if dtype == torch.int32:
+            if isinstance(cs, np.ndarray):
+                cs = torch.from_numpy(np.floor(cs))
+            elif not (
+                isinstance(cs, torch.IntTensor) or isinstance(cs, torch.LongTensor)
+            ):
+                cs = cs.floor()
+
+            cs = cs.int()
+        else:
+            if isinstance(cs, np.ndarray):
+                cs = torch.from_numpy(cs)
+
+        cn = len(cs)
+        # BATCH_FIRST:
+        bcoords[s : s + cn, 1:] = cs
+        bcoords[s : s + cn, 0] = b
+        s += cn
+    return bcoords
+
+
+def sparse_collate(coords, feats, labels=None, dtype=torch.int32, device=None):
+    r"""Create input arguments for a sparse tensor `the documentation
+    <https://nvidia.github.io/MinkowskiEngine/sparse_tensor.html>`_.
+
+    Convert a set of coordinates and features into the batch coordinates and
+    batch features.
+
+    Args:
+        :attr:`coords` (set of `torch.Tensor` or `numpy.ndarray`): a set of coordinates.
+
+        :attr:`feats` (set of `torch.Tensor` or `numpy.ndarray`): a set of features.
+
+        :attr:`labels` (set of `torch.Tensor` or `numpy.ndarray`): a set of labels
+        associated to the inputs.
+
+    """
+    use_label = False if labels is None else True
+    feats_batch, labels_batch = [], []
+    assert isinstance(
+        coords, collections.abc.Sequence
+    ), "The coordinates must be a sequence of arrays or tensors."
+    assert isinstance(
+        feats, collections.abc.Sequence
+    ), "The features must be a sequence of arrays or tensors."
+    D = np.unique(np.array([cs.shape[1] for cs in coords]))
+    assert len(D) == 1, f"Dimension of the array mismatch. All dimensions: {D}"
+    D = D[0]
+    if device is None:
+        if isinstance(coords[0], torch.Tensor):
+            device = coords[0].device
+        else:
+            device = "cpu"
+    assert dtype in [
+        torch.int32,
+        torch.float32,
+    ], "Only torch.int32, torch.float32 supported for coordinates."
+
+    if use_label:
+        assert isinstance(
+            labels, collections.abc.Sequence
+        ), "The labels must be a sequence of arrays or tensors."
+
+    N = np.array([len(cs) for cs in coords]).sum()
+    Nf = np.array([len(fs) for fs in feats]).sum()
+    assert N == Nf, f"Coordinate length {N} != Feature length {Nf}"
+
+    batch_id = 0
+    s = 0  # start index
+    bcoords = torch.zeros((N, D + 1), dtype=dtype, device=device)  # uninitialized
+    for coord, feat in zip(coords, feats):
+        if isinstance(coord, np.ndarray):
+            coord = torch.from_numpy(coord)
+        else:
+            assert isinstance(
+                coord, torch.Tensor
+            ), "Coords must be of type numpy.ndarray or torch.Tensor"
+        if dtype == torch.int32 and coord.dtype in [torch.float32, torch.float64]:
+            coord = coord.floor()
+
+        if isinstance(feat, np.ndarray):
+            feat = torch.from_numpy(feat)
+        else:
+            assert isinstance(
+                feat, torch.Tensor
+            ), "Features must be of type numpy.ndarray or torch.Tensor"
+
+        # Labels
+        if use_label:
+            label = labels[batch_id]
+            if isinstance(label, np.ndarray):
+                label = torch.from_numpy(label)
+            labels_batch.append(label)
+
+        cn = coord.shape[0]
+        # Batched coords
+        bcoords[s : s + cn, 1:] = coord
+        bcoords[s : s + cn, 0] = batch_id
+
+        # Features
+        feats_batch.append(feat)
+
+        # Post processing steps
+        batch_id += 1
+        s += cn
+
+    # Concatenate all lists
+    feats_batch = torch.cat(feats_batch, 0)
+    if use_label:
+        if isinstance(labels_batch[0], torch.Tensor):
+            labels_batch = torch.cat(labels_batch, 0)
+        return bcoords, feats_batch, labels_batch
+    else:
+        return bcoords, feats_batch
+
+
+def batch_sparse_collate(data, dtype=torch.int32, device=None):
+    r"""The wrapper function that can be used in in conjunction with
+    `torch.utils.data.DataLoader` to generate inputs for a sparse tensor.
+
+    Please refer to `the training example
+    <https://nvidia.github.io/MinkowskiEngine/demo/training.html>`_ for the
+    usage.
+
+    Args:
+        :attr:`data`: list of (coordinates, features, labels) tuples.
+
+    """
+    return sparse_collate(*list(zip(*data)), dtype=dtype, device=device)
+
+
+class SparseCollation:
+    r"""Generates collate function for coords, feats, labels.
+
+    Please refer to `the training example
+    <https://nvidia.github.io/MinkowskiEngine/demo/training.html>`_ for the
+    usage.
+
+    Args:
+        :attr:`limit_numpoints` (int): If positive integer, limits batch size
+        so that the number of input coordinates is below limit_numpoints. If 0
+        or False, concatenate all points. -1 by default.
+
+    Example::
+
+        >>> data_loader = torch.utils.data.DataLoader(
+        >>>     dataset,
+        >>>     ...,
+        >>>     collate_fn=SparseCollation())
+        >>> for d in iter(data_loader):
+        >>>     print(d)
+
+    """
+
+    def __init__(self, limit_numpoints=-1, dtype=torch.int32, device=None):
+        self.limit_numpoints = limit_numpoints
+        self.dtype = dtype
+        self.device = device
+
+    def __call__(self, list_data):
+        coords, feats, labels = list(zip(*list_data))
+        coords_batch, feats_batch, labels_batch = [], [], []
+
+        batch_num_points = 0
+        for batch_id, _ in enumerate(coords):
+            num_points = coords[batch_id].shape[0]
+            batch_num_points += num_points
+            if self.limit_numpoints > 0 and batch_num_points > self.limit_numpoints:
+                num_full_points = sum(len(c) for c in coords)
+                num_full_batch_size = len(coords)
+                logging.warning(
+                    f"\tCannot fit {num_full_points} points into"
+                    " {self.limit_numpoints} points limit. Truncating batch "
+                    f"size at {batch_id} out of {num_full_batch_size} with "
+                    f"{batch_num_points - num_points}."
+                )
+                break
+            coords_batch.append(coords[batch_id])
+            feats_batch.append(feats[batch_id])
+            labels_batch.append(labels[batch_id])
+
+        # Concatenate all lists
+        return sparse_collate(
+            coords_batch,
+            feats_batch,
+            labels_batch,
+            dtype=self.dtype,
+            device=self.device,
+        )

MinkowskiEngine/MinkowskiEngine/utils/coords.py

@@ -0,0 +1,63 @@
+# Copyright (c) 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.
+import torch
+
+from MinkowskiSparseTensor import SparseTensor
+
+
+def get_coords_map(x, y):
+    r"""Get mapping between sparse tensor 1 and sparse tensor 2.
+
+    Args:
+        :attr:`x` (:attr:`MinkowskiEngine.SparseTensor`): a sparse tensor with
+        `x.tensor_stride` <= `y.tensor_stride`.
+
+        :attr:`y` (:attr:`MinkowskiEngine.SparseTensor`): a sparse tensor with
+        `x.tensor_stride` <= `y.tensor_stride`.
+
+    Returns:
+        :attr:`x_indices` (:attr:`torch.LongTensor`): the indices of x that
+        corresponds to the returned indices of y.
+
+        :attr:`x_indices` (:attr:`torch.LongTensor`): the indices of y that
+        corresponds to the returned indices of x.
+
+    Example::
+
+        .. code-block:: python
+
+           sp_tensor = ME.SparseTensor(features, coordinates=coordinates)
+           out_sp_tensor = stride_2_conv(sp_tensor)
+
+           ins, outs = get_coords_map(sp_tensor, out_sp_tensor)
+           for i, o in zip(ins, outs):
+              print(f"{i} -> {o}")
+
+    """
+    assert isinstance(x, SparseTensor)
+    assert isinstance(y, SparseTensor)
+    assert (
+        x.coords_man == y.coords_man
+    ), "X and Y are using different CoordinateManagers. Y must be derived from X through strided conv/pool/etc."
+    return x.coords_man.get_coords_map(x.coords_key, y.coords_key)

MinkowskiEngine/MinkowskiEngine/utils/gradcheck.py

@@ -0,0 +1,57 @@
+# Copyright (c) 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.
+import torch
+
+assert torch.__version__ >= "1.7.0", "Gradcheck requires pytorch 1.7 or higher"
+
+from torch.types import _TensorOrTensors
+from typing import Callable, Union, Optional
+
+from torch.autograd.gradcheck import gradcheck as _gradcheck
+
+
+def gradcheck(
+    func: Callable[..., Union[_TensorOrTensors]],  # See Note [VarArg of Tensors]
+    inputs: _TensorOrTensors,
+    eps: float = 1e-6,
+    atol: float = 1e-5,
+    rtol: float = 1e-3,
+    raise_exception: bool = True,
+    check_sparse_nnz: bool = False,
+    nondet_tol: float = 0.0,
+    check_undefined_grad: bool = True,
+    check_grad_dtypes: bool = False,
+) -> bool:
+    return _gradcheck(
+        lambda *x: func.apply(*x),
+        inputs,
+        eps=eps,
+        atol=atol,
+        rtol=rtol,
+        raise_exception=raise_exception,
+        check_sparse_nnz=check_sparse_nnz,
+        nondet_tol=nondet_tol,
+        check_undefined_grad=check_undefined_grad,
+        check_grad_dtypes=check_grad_dtypes,
+    )

MinkowskiEngine/MinkowskiEngine/utils/init.py

@@ -0,0 +1,41 @@
+import math
+import torch
+
+
+def _calculate_fan_in_and_fan_out(tensor):
+    dimensions = tensor.dim()
+    if dimensions < 2:
+        raise ValueError(
+            "Fan in and fan out can not be computed for tensor with fewer than 2 dimensions"
+        )
+
+    if dimensions == 2:  # Linear
+        fan_in = tensor.size(1)
+        fan_out = tensor.size(0)
+    else:
+        num_input_fmaps = tensor.size(1)
+        num_output_fmaps = tensor.size(2)
+        receptive_field_size = tensor.size(0)
+        fan_in = num_input_fmaps * receptive_field_size
+        fan_out = num_output_fmaps * receptive_field_size
+
+    return fan_in, fan_out
+
+
+def _calculate_correct_fan(tensor, mode):
+    mode = mode.lower()
+    valid_modes = ['fan_in', 'fan_out']
+    if mode not in valid_modes:
+        raise ValueError("Mode {} not supported, please use one of {}".format(
+            mode, valid_modes))
+
+    fan_in, fan_out = _calculate_fan_in_and_fan_out(tensor)
+    return fan_in if mode == 'fan_in' else fan_out
+
+
+def kaiming_normal_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'):
+    fan = _calculate_correct_fan(tensor, mode)
+    gain = torch.nn.init.calculate_gain(nonlinearity, a)
+    std = gain / math.sqrt(fan)
+    with torch.no_grad():
+        return tensor.normal_(0, std)

MinkowskiEngine/MinkowskiEngine/utils/quantization.py

@@ -0,0 +1,363 @@
+# Copyright (c) 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.
+import torch
+import numpy as np
+from collections.abc import Sequence
+import MinkowskiEngineBackend._C as MEB
+from typing import Union, Tuple
+from MinkowskiCommon import convert_to_int_list
+
+
+def fnv_hash_vec(arr):
+    """
+    FNV64-1A
+    """
+    assert arr.ndim == 2
+    # Floor first for negative coordinates
+    arr = arr.copy()
+    arr = arr.astype(np.uint64, copy=False)
+    hashed_arr = np.uint64(14695981039346656037) * np.ones(
+        arr.shape[0], dtype=np.uint64
+    )
+    for j in range(arr.shape[1]):
+        hashed_arr *= np.uint64(1099511628211)
+        hashed_arr = np.bitwise_xor(hashed_arr, arr[:, j])
+    return hashed_arr
+
+
+def ravel_hash_vec(arr):
+    """
+    Ravel the coordinates after subtracting the min coordinates.
+    """
+    assert arr.ndim == 2
+    arr = arr.copy()
+    arr -= arr.min(0)
+    arr = arr.astype(np.uint64, copy=False)
+    arr_max = arr.max(0).astype(np.uint64) + 1
+
+    keys = np.zeros(arr.shape[0], dtype=np.uint64)
+    # Fortran style indexing
+    for j in range(arr.shape[1] - 1):
+        keys += arr[:, j]
+        keys *= arr_max[j + 1]
+    keys += arr[:, -1]
+    return keys
+
+
+def quantize(coords):
+    r"""Returns a unique index map and an inverse index map.
+
+    Args:
+        :attr:`coords` (:attr:`numpy.ndarray` or :attr:`torch.Tensor`): a
+        matrix of size :math:`N \times D` where :math:`N` is the number of
+        points in the :math:`D` dimensional space.
+
+    Returns:
+        :attr:`unique_map` (:attr:`numpy.ndarray` or :attr:`torch.Tensor`): a
+        list of indices that defines unique coordinates.
+        :attr:`coords[unique_map]` is the unique coordinates.
+
+        :attr:`inverse_map` (:attr:`numpy.ndarray` or :attr:`torch.Tensor`): a
+        list of indices that defines the inverse map that recovers the original
+        coordinates.  :attr:`coords[unique_map[inverse_map]] == coords`
+
+    Example::
+
+       >>> unique_map, inverse_map = quantize(coords)
+       >>> unique_coords = coords[unique_map]
+       >>> print(unique_coords[inverse_map] == coords)  # True, ..., True
+       >>> print(coords[unique_map[inverse_map]] == coords)  # True, ..., True
+
+    """
+    assert isinstance(coords, np.ndarray) or isinstance(
+        coords, torch.Tensor
+    ), "Invalid coords type"
+    if isinstance(coords, np.ndarray):
+        assert (
+            coords.dtype == np.int32
+        ), f"Invalid coords type {coords.dtype} != np.int32"
+        return MEB.quantize_np(coords.astype(np.int32))
+    else:
+        # Type check done inside
+        return MEB.quantize_th(coords.int())
+
+
+def quantize_label(coords, labels, ignore_label):
+    assert isinstance(coords, np.ndarray) or isinstance(
+        coords, torch.Tensor
+    ), "Invalid coords type"
+    if isinstance(coords, np.ndarray):
+        assert isinstance(labels, np.ndarray)
+        assert (
+            coords.dtype == np.int32
+        ), f"Invalid coords type {coords.dtype} != np.int32"
+        assert (
+            labels.dtype == np.int32
+        ), f"Invalid label type {labels.dtype} != np.int32"
+        return MEB.quantize_label_np(coords, labels, ignore_label)
+    else:
+        assert isinstance(labels, torch.Tensor)
+        # Type check done inside
+        return MEB.quantize_label_th(coords, labels.int(), ignore_label)
+
+
+def _auto_floor(array):
+    assert isinstance(
+        array, (np.ndarray, torch.Tensor)
+    ), "array must be either np.array or torch.Tensor."
+
+    if isinstance(array, np.ndarray):
+        return np.floor(array)
+    else:
+        return torch.floor(array)
+
+
+def sparse_quantize(
+    coordinates,
+    features=None,
+    labels=None,
+    ignore_label=-100,
+    return_index=False,
+    return_inverse=False,
+    return_maps_only=False,
+    quantization_size=None,
+    device="cpu",
+):
+    r"""Given coordinates, and features (optionally labels), the function
+    generates quantized (voxelized) coordinates.
+
+    Args:
+        :attr:`coordinates` (:attr:`numpy.ndarray` or :attr:`torch.Tensor`): a
+        matrix of size :math:`N \times D` where :math:`N` is the number of
+        points in the :math:`D` dimensional space.
+
+        :attr:`features` (:attr:`numpy.ndarray` or :attr:`torch.Tensor`, optional): a
+        matrix of size :math:`N \times D_F` where :math:`N` is the number of
+        points and :math:`D_F` is the dimension of the features. Must have the
+        same container as `coords` (i.e. if `coords` is a torch.Tensor, `feats`
+        must also be a torch.Tensor).
+
+        :attr:`labels` (:attr:`numpy.ndarray` or :attr:`torch.IntTensor`,
+        optional): integer labels associated to eah coordinates.  Must have the
+        same container as `coords` (i.e. if `coords` is a torch.Tensor,
+        `labels` must also be a torch.Tensor). For classification where a set
+        of points are mapped to one label, do not feed the labels.
+
+        :attr:`ignore_label` (:attr:`int`, optional): the int value of the
+        IGNORE LABEL.
+        :attr:`torch.nn.CrossEntropyLoss(ignore_index=ignore_label)`
+
+        :attr:`return_index` (:attr:`bool`, optional): set True if you want the
+        indices of the quantized coordinates. False by default.
+
+        :attr:`return_inverse` (:attr:`bool`, optional): set True if you want
+        the indices that can recover the discretized original coordinates.
+        False by default. `return_index` must be True when `return_reverse` is True.
+
+        :attr:`return_maps_only` (:attr:`bool`, optional): if set, return the
+        unique_map or optionally inverse map, but not the coordinates. Can be
+        used if you don't care about final coordinates or if you use
+        device==cuda and you don't need coordinates on GPU. This returns either
+        unique_map alone or (unique_map, inverse_map) if return_inverse is set.
+
+        :attr:`quantization_size` (attr:`float`, optional): if set, will use
+        the quanziation size to define the smallest distance between
+        coordinates.
+
+        :attr:`device` (attr:`str`, optional): Either 'cpu' or 'cuda'.
+
+        Example::
+
+           >>> unique_map, inverse_map = sparse_quantize(discrete_coords, return_index=True, return_inverse=True)
+           >>> unique_coords = discrete_coords[unique_map]
+           >>> print(unique_coords[inverse_map] == discrete_coords)  # True
+
+        :attr:`quantization_size` (:attr:`float`, :attr:`list`, or
+        :attr:`numpy.ndarray`, optional): the length of the each side of the
+        hyperrectangle of of the grid cell.
+
+     Example::
+
+        >>> # Segmentation
+        >>> criterion = torch.nn.CrossEntropyLoss(ignore_index=-100)
+        >>> coords, feats, labels = MinkowskiEngine.utils.sparse_quantize(
+        >>>     coords, feats, labels, ignore_label=-100, quantization_size=0.1)
+        >>> output = net(MinkowskiEngine.SparseTensor(feats, coords))
+        >>> loss = criterion(output.F, labels.long())
+        >>>
+        >>> # Classification
+        >>> criterion = torch.nn.CrossEntropyLoss(ignore_index=-100)
+        >>> coords, feats = MinkowskiEngine.utils.sparse_quantize(coords, feats)
+        >>> output = net(MinkowskiEngine.SparseTensor(feats, coords))
+        >>> loss = criterion(output.F, labels.long())
+
+
+    """
+    assert isinstance(
+        coordinates, (np.ndarray, torch.Tensor)
+    ), "Coords must be either np.array or torch.Tensor."
+
+    use_label = labels is not None
+    use_feat = features is not None
+
+    assert (
+        coordinates.ndim == 2
+    ), "The coordinates must be a 2D matrix. The shape of the input is " + str(
+        coordinates.shape
+    )
+
+    if return_inverse:
+        assert return_index, "return_reverse must be set with return_index"
+
+    if use_feat:
+        assert features.ndim == 2
+        assert coordinates.shape[0] == features.shape[0]
+
+    if use_label:
+        assert coordinates.shape[0] == len(labels)
+
+    dimension = coordinates.shape[1]
+    # Quantize the coordinates
+    if quantization_size is not None:
+        if isinstance(quantization_size, (Sequence, np.ndarray, torch.Tensor)):
+            assert (
+                len(quantization_size) == dimension
+            ), "Quantization size and coordinates size mismatch."
+            if isinstance(coordinates, np.ndarray):
+                quantization_size = np.array([i for i in quantization_size])
+            else:
+                quantization_size = torch.Tensor([i for i in quantization_size])
+            discrete_coordinates = _auto_floor(coordinates / quantization_size)
+
+        elif np.isscalar(quantization_size):  # Assume that it is a scalar
+
+            if quantization_size == 1:
+                discrete_coordinates = _auto_floor(coordinates)
+            else:
+                discrete_coordinates = _auto_floor(coordinates / quantization_size)
+        else:
+            raise ValueError("Not supported type for quantization_size.")
+    else:
+        discrete_coordinates = _auto_floor(coordinates)
+
+    if isinstance(coordinates, np.ndarray):
+        discrete_coordinates = discrete_coordinates.astype(np.int32)
+    else:
+        discrete_coordinates = discrete_coordinates.int()
+
+    if (type(device) == str and device == "cpu") or (type(device) == torch.device and device.type == "cpu"):
+        manager = MEB.CoordinateMapManagerCPU()
+    elif (type(device) == str and "cuda" in device) or (type(device) == torch.device and device.type == "cuda"):
+        manager = MEB.CoordinateMapManagerGPU_c10()
+    else:
+        raise ValueError("Invalid device. Only `cpu`, `cuda` or torch.device supported.")
+
+    # Return values accordingly
+    if use_label:
+        if isinstance(coordinates, np.ndarray):
+            unique_map, inverse_map, colabels = MEB.quantize_label_np(
+                discrete_coordinates, labels, ignore_label
+            )
+        else:
+            assert (
+                not discrete_coordinates.is_cuda
+            ), "Quantization with label requires cpu tensors."
+            assert not labels.is_cuda, "Quantization with label requires cpu tensors."
+            unique_map, inverse_map, colabels = MEB.quantize_label_th(
+                discrete_coordinates, labels, ignore_label
+            )
+        return_args = [discrete_coordinates[unique_map]]
+        if use_feat:
+            return_args.append(features[unique_map])
+        # Labels
+        return_args.append(colabels)
+        # Additional return args
+        if return_index:
+            return_args.append(unique_map)
+        if return_inverse:
+            return_args.append(inverse_map)
+
+        if len(return_args) == 1:
+            return return_args[0]
+        else:
+            return tuple(return_args)
+    else:
+        tensor_stride = [1 for i in range(discrete_coordinates.shape[1] - 1)]
+        discrete_coordinates = (
+            discrete_coordinates.to(device)
+            if isinstance(discrete_coordinates, torch.Tensor)
+            else torch.from_numpy(discrete_coordinates).to(device)
+        )
+        _, (unique_map, inverse_map) = manager.insert_and_map(
+            discrete_coordinates, tensor_stride, ""
+        )
+        if return_maps_only:
+            if return_inverse:
+                return unique_map, inverse_map
+            else:
+                return unique_map
+
+        return_args = [discrete_coordinates[unique_map]]
+        if use_feat:
+            return_args.append(features[unique_map])
+        if return_index:
+            return_args.append(unique_map)
+        if return_inverse:
+            return_args.append(inverse_map)
+
+        if len(return_args) == 1:
+            return return_args[0]
+        else:
+            return tuple(return_args)
+
+
+def unique_coordinate_map(
+    coordinates: torch.Tensor,
+    tensor_stride: Union[int, Sequence, np.ndarray] = 1,
+) -> Tuple[torch.IntTensor, torch.IntTensor]:
+    r"""Returns the unique indices and the inverse indices of the coordinates.
+
+    :attr:`coordinates`: `torch.Tensor` (Int tensor. `CUDA` if
+    coordinate_map_type == `CoordinateMapType.GPU`) that defines the
+    coordinates.
+
+    Example::
+
+       >>> coordinates = torch.IntTensor([[0, 0], [0, 0], [0, 1], [0, 2]])
+       >>> unique_map, inverse_map = unique_coordinates_map(coordinates)
+       >>> coordinates[unique_map] # unique coordinates
+       >>> torch.all(coordinates == coordinates[unique_map][inverse_map]) # True
+
+    """
+    assert coordinates.ndim == 2, "Coordinates must be a matrix"
+    assert isinstance(coordinates, torch.Tensor)
+    if not coordinates.is_cuda:
+        manager = MEB.CoordinateMapManagerCPU()
+    else:
+        manager = MEB.CoordinateMapManagerGPU_c10()
+    tensor_stride = convert_to_int_list(tensor_stride, coordinates.shape[-1] - 1)
+    key, (unique_map, inverse_map) = manager.insert_and_map(
+        coordinates, tensor_stride, ""
+    )
+    return unique_map, inverse_map

MinkowskiEngine/MinkowskiEngine/utils/summary.py

@@ -0,0 +1,136 @@
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+
+from collections import OrderedDict
+import numpy as np
+
+import MinkowskiEngine as ME
+from MinkowskiSparseTensor import SparseTensor
+
+
+def summary(model, summary_input):
+    result, params_info = minkowski_summary_string(model, summary_input)
+    print(result)
+    return params_info
+
+
+def pruned_weight_sparsity_string(module) -> float:
+    r"""
+    returns the sparsity ratio of weights.
+    """
+    for k in dir(module):
+        if '_mask' in k:
+            return (getattr(module, k.replace('_mask', '')) == 0).float().mean().item()
+    else:
+        return 0.0;
+
+
+def size2list(size: torch.Size) -> list:
+    return [i for i in size]
+
+def get_hash_occupancy_ratio(minkowski_tensor):
+    alg = minkowski_tensor.coordinate_manager.minkowski_algorithm
+    if alg == ME.MinkowskiAlgorithm.SPEED_OPTIMIZED:
+        return 25;
+    else:
+        return 50;
+
+def minkowski_summary_string(model, summary_input):
+    summary_str = ''
+
+    def register_hook(module):
+        def hook(module, input, output):
+            class_name = str(module.__class__).split(".")[-1].split("'")[0]
+            module_idx = len(summary)
+
+            m_key = "%s-%i" % (class_name, module_idx + 1)
+            summary[m_key] = OrderedDict()
+
+            # for the weight pruned model, print the sparsity information
+            summary[m_key]['sparsity_ratio'] = pruned_weight_sparsity_string(module)
+
+            # save only the size of NNZ
+            summary[m_key]["input_shape"] = input[0].shape
+            if isinstance(output, (list, tuple)):
+                summary[m_key]["output_shape"] = [size2list(o.shape) for o in output]
+            else:
+                summary[m_key]["output_shape"] = size2list(output.shape)
+
+            params = 0
+            if hasattr(module, "weight") and hasattr(module.weight, "size"):
+                params += module.weight.numel()
+                summary[m_key]["trainable"] = module.weight.requires_grad
+            if hasattr(module, "kernel") and hasattr(module.kernel, "size"):
+                params += module.kernel.numel()
+                summary[m_key]["trainable"] = module.kernel.requires_grad
+            if hasattr(module, "bias") and hasattr(module.bias, "size"):
+                params += module.bias.numel()
+            summary[m_key]["nb_params"] = params
+
+        if (
+            not isinstance(module, nn.Sequential)
+            and not isinstance(module, nn.ModuleList)
+        ):
+            hooks.append(module.register_forward_hook(hook))
+
+    # create properties
+    summary = OrderedDict()
+    hooks = []
+
+    # register hook
+    model.apply(register_hook)
+
+    # make a forward pass
+    # print(x.shape)
+    model(summary_input)
+
+    # remove these hooks
+    for h in hooks:
+        h.remove()
+
+    summary_str += "----------------------------------------------------------------" + "\n"
+    line_new = "{:>20}  {:>25} {:>15}".format(
+        "Layer (type)", "Output Shape", "Param #")
+    summary_str += line_new + "\n"
+    summary_str += "================================================================" + "\n"
+    total_params = 0
+    total_output = 0
+    trainable_params = 0
+    for layer in summary:
+        # input_shape, output_shape, trainable, nb_params
+        line_new = "{:>20}  {:>25} {:>15}".format(
+            layer,
+            str(summary[layer]["output_shape"]),
+            "{0:,}".format(summary[layer]["nb_params"]),
+        )
+        total_params += summary[layer]["nb_params"]
+
+        total_output += np.prod(summary[layer]["output_shape"])
+        if "trainable" in summary[layer]:
+            if summary[layer]["trainable"] == True:
+                trainable_params += summary[layer]["nb_params"]
+        summary_str += line_new + "\n"
+
+    # assume 4 bytes/number (float on cuda).
+    total_input_size = (len(summary_input) * summary_input.shape[1]  # feature size
+                        + len(summary_input) * (1 + summary_input.D) * (100 / get_hash_occupancy_ratio(summary_input)) # coordinate size
+                       ) * 4. / (1024 ** 2.)
+    total_output_size = abs(2. * total_output * 4. /
+                            (1024 ** 2.))  # x2 for gradients
+    total_params_size = abs(total_params * 4. / (1024 ** 2.))
+    total_size = total_params_size + total_output_size + total_input_size
+
+    summary_str += "================================================================" + "\n"
+    summary_str += "Total params: {0:,}".format(total_params) + "\n"
+    summary_str += "Trainable params: {0:,}".format(trainable_params) + "\n"
+    summary_str += "Non-trainable params: {0:,}".format(total_params -
+                                                        trainable_params) + "\n"
+    summary_str += "----------------------------------------------------------------" + "\n"
+    summary_str += "Input size (MB): %0.2f" % total_input_size + "\n"
+    summary_str += "Forward/backward pass size (MB): %0.2f" % total_output_size + "\n"
+    summary_str += "Params size (MB): %0.2f" % total_params_size + "\n"
+    summary_str += "Estimated Total Size (MB): %0.2f" % total_size + "\n"
+    summary_str += "----------------------------------------------------------------" + "\n"
+    # return summary
+    return summary_str, (total_params, trainable_params)

MinkowskiEngine/README.md

@@ -0,0 +1,396 @@
+[pypi-image]: https://badge.fury.io/py/MinkowskiEngine.svg
+[pypi-url]: https://pypi.org/project/MinkowskiEngine/
+[pypi-download]: https://img.shields.io/pypi/dm/MinkowskiEngine
+[slack-badge]: https://img.shields.io/badge/slack-join%20chats-brightgreen
+[slack-url]: https://join.slack.com/t/minkowskiengine/shared_invite/zt-piq2x02a-31dOPocLt6bRqOGY3U_9Sw
+
+# Minkowski Engine
+
+[![PyPI Version][pypi-image]][pypi-url] [![pypi monthly download][pypi-download]][pypi-url] [![slack chat][slack-badge]][slack-url]
+
+The Minkowski Engine is an auto-differentiation library for sparse tensors. It supports all standard neural network layers such as convolution, pooling, unpooling, and broadcasting operations for sparse tensors. For more information, please visit [the documentation page](http://nvidia.github.io/MinkowskiEngine/overview.html).
+
+## News
+
+- 2021-08-11 Docker installation instruction added
+- 2021-08-06 All installation errors with pytorch 1.8 and 1.9 have been resolved.
+- 2021-04-08 Due to recent errors in [pytorch 1.8 + CUDA 11](https://github.com/NVIDIA/MinkowskiEngine/issues/330), it is recommended to use [anaconda for installation](#anaconda).
+- 2020-12-24 v0.5 is now available! The new version provides CUDA accelerations for all coordinate management functions.
+
+## Example Networks
+
+The Minkowski Engine supports various functions that can be built on a sparse tensor. We list a few popular network architectures and applications here. To run the examples, please install the package and run the command in the package root directory.
+
+| Examples              | Networks and Commands                                                                                                                                                           |
+|:---------------------:|:-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------:|
+| Semantic Segmentation | <img src="https://nvidia.github.io/MinkowskiEngine/_images/segmentation_3d_net.png"> <br /> <img src="https://nvidia.github.io/MinkowskiEngine/_images/segmentation.png" width="256"> <br /> `python -m examples.indoor` |
+| Classification        | ![](https://nvidia.github.io/MinkowskiEngine/_images/classification_3d_net.png) <br /> `python -m examples.classification_modelnet40`                                                          |
+| Reconstruction        | <img src="https://nvidia.github.io/MinkowskiEngine/_images/generative_3d_net.png"> <br /> <img src="https://nvidia.github.io/MinkowskiEngine/_images/generative_3d_results.gif" width="256"> <br /> `python -m examples.reconstruction` |
+| Completion            | <img src="https://nvidia.github.io/MinkowskiEngine/_images/completion_3d_net.png"> <br /> `python -m examples.completion`                                                       |
+| Detection             | <img src="https://nvidia.github.io/MinkowskiEngine/_images/detection_3d_net.png">                                                                                               |
+
+
+## Sparse Tensor Networks: Neural Networks for Spatially Sparse Tensors
+
+Compressing a neural network to speedup inference and minimize memory footprint has been studied widely. One of the popular techniques for model compression is pruning the weights in convnets, is also known as [*sparse convolutional networks*](https://www.cv-foundation.org/openaccess/content_cvpr_2015/papers/Liu_Sparse_Convolutional_Neural_2015_CVPR_paper.pdf). Such parameter-space sparsity used for model compression compresses networks that operate on dense tensors and all intermediate activations of these networks are also dense tensors.
+
+However, in this work, we focus on [*spatially* sparse data](https://arxiv.org/abs/1409.6070), in particular, spatially sparse high-dimensional inputs and 3D data and convolution on the surface of 3D objects, first proposed in [Siggraph'17](https://wang-ps.github.io/O-CNN.html). We can also represent these data as sparse tensors, and these sparse tensors are commonplace in high-dimensional problems such as 3D perception, registration, and statistical data. We define neural networks specialized for these inputs as *sparse tensor networks*  and these sparse tensor networks process and generate sparse tensors as outputs. To construct a sparse tensor network, we build all standard neural network layers such as MLPs, non-linearities, convolution, normalizations, pooling operations as the same way we define them on a dense tensor and implemented in the Minkowski Engine.
+
+We visualized a sparse tensor network operation on a sparse tensor, convolution, below. The convolution layer on a sparse tensor works similarly to that on a dense tensor. However, on a sparse tensor, we compute convolution outputs on a few specified points which we can control in the [generalized convolution](https://nvidia.github.io/MinkowskiEngine/sparse_tensor_network.html). For more information, please visit [the documentation page on sparse tensor networks](https://nvidia.github.io/MinkowskiEngine/sparse_tensor_network.html) and [the terminology page](https://nvidia.github.io/MinkowskiEngine/terminology.html).
+
+| Dense Tensor                                                                | Sparse Tensor                                                                |
+|:---------------------------------------------------------------------------:|:----------------------------------------------------------------------------:|
+| <img src="https://nvidia.github.io/MinkowskiEngine/_images/conv_dense.gif"> | <img src="https://nvidia.github.io/MinkowskiEngine/_images/conv_sparse.gif"> |
+
+--------------------------------------------------------------------------------
+
+## Features
+
+- Unlimited high-dimensional sparse tensor support
+- All standard neural network layers (Convolution, Pooling, Broadcast, etc.)
+- Dynamic computation graph
+- Custom kernel shapes
+- Multi-GPU training
+- Multi-threaded kernel map
+- Multi-threaded compilation
+- Highly-optimized GPU kernels
+
+
+## Requirements
+
+- Ubuntu >= 14.04
+- CUDA >= 10.1.243 and **the same CUDA version used for pytorch** (e.g. if you use conda cudatoolkit=11.1, use CUDA=11.1 for MinkowskiEngine compilation)
+- pytorch >= 1.7 To specify CUDA version, please use conda for installation. You must match the CUDA version pytorch uses and CUDA version used for Minkowski Engine installation. `conda install -y -c nvidia -c pytorch pytorch=1.8.1 cudatoolkit=10.2`)
+- python >= 3.6
+- ninja (for installation)
+- GCC >= 7.4.0
+
+
+## Installation
+
+You can install the Minkowski Engine with `pip`, with anaconda, or on the system directly. If you experience issues installing the package, please checkout the [the installation wiki page](https://github.com/NVIDIA/MinkowskiEngine/wiki/Installation).
+If you cannot find a relevant problem, please report the issue on [the github issue page](https://github.com/NVIDIA/MinkowskiEngine/issues).
+
+- [PIP](https://github.com/NVIDIA/MinkowskiEngine#pip) installation
+- [Conda](https://github.com/NVIDIA/MinkowskiEngine#anaconda) installation
+- [Python](https://github.com/NVIDIA/MinkowskiEngine#system-python) installation
+- [Docker](https://github.com/NVIDIA/MinkowskiEngine#docker) installation
+
+
+### Pip
+
+The MinkowskiEngine is distributed via [PyPI MinkowskiEngine][pypi-url] which can be installed simply with `pip`.
+First, install pytorch following the [instruction](https://pytorch.org). Next, install `openblas`.
+
+```
+sudo apt install build-essential python3-dev libopenblas-dev
+pip install torch ninja
+pip install -U MinkowskiEngine --install-option="--blas=openblas" -v --no-deps
+
+# For pip installation from the latest source
+# pip install -U git+https://github.com/NVIDIA/MinkowskiEngine --no-deps
+```
+
+If you want to specify arguments for the setup script, please refer to the following command.
+
+```
+# Uncomment some options if things don't work
+# export CXX=c++; # set this if you want to use a different C++ compiler
+# export CUDA_HOME=/usr/local/cuda-11.1; # or select the correct cuda version on your system.
+pip install -U git+https://github.com/NVIDIA/MinkowskiEngine -v --no-deps \
+#                           \ # uncomment the following line if you want to force cuda installation
+#                           --install-option="--force_cuda" \
+#                           \ # uncomment the following line if you want to force no cuda installation. force_cuda supercedes cpu_only
+#                           --install-option="--cpu_only" \
+#                           \ # uncomment the following line to override to openblas, atlas, mkl, blas
+#                           --install-option="--blas=openblas" \
+```
+
+### Anaconda
+
+MinkowskiEngine supports both CUDA 10.2 and cuda 11.1, which work for most of latest pytorch versions.
+#### CUDA 10.2
+
+We recommend `python>=3.6` for installation.
+First, follow [the anaconda documentation](https://docs.anaconda.com/anaconda/install/) to install anaconda on your computer.
+
+```
+sudo apt install g++-7  # For CUDA 10.2, must use GCC < 8
+# Make sure `g++-7 --version` is at least 7.4.0
+conda create -n py3-mink python=3.8
+conda activate py3-mink
+
+conda install openblas-devel -c anaconda
+conda install pytorch=1.9.0 torchvision cudatoolkit=10.2 -c pytorch -c nvidia
+
+# Install MinkowskiEngine
+export CXX=g++-7
+# Uncomment the following line to specify the cuda home. Make sure `$CUDA_HOME/nvcc --version` is 10.2
+# export CUDA_HOME=/usr/local/cuda-10.2
+pip install -U git+https://github.com/NVIDIA/MinkowskiEngine -v --no-deps --install-option="--blas_include_dirs=${CONDA_PREFIX}/include" --install-option="--blas=openblas"
+
+# Or if you want local MinkowskiEngine
+git clone https://github.com/NVIDIA/MinkowskiEngine.git
+cd MinkowskiEngine
+export CXX=g++-7
+python setup.py install --blas_include_dirs=${CONDA_PREFIX}/include --blas=openblas
+```
+
+#### CUDA 11.X
+
+We recommend `python>=3.6` for installation.
+First, follow [the anaconda documentation](https://docs.anaconda.com/anaconda/install/) to install anaconda on your computer.
+
+```
+conda create -n py3-mink python=3.8
+conda activate py3-mink
+
+conda install openblas-devel -c anaconda
+conda install pytorch=1.9.0 torchvision cudatoolkit=11.1 -c pytorch -c nvidia
+
+# Install MinkowskiEngine
+
+# Uncomment the following line to specify the cuda home. Make sure `$CUDA_HOME/nvcc --version` is 11.X
+# export CUDA_HOME=/usr/local/cuda-11.1
+pip install -U git+https://github.com/NVIDIA/MinkowskiEngine -v --no-deps --install-option="--blas_include_dirs=${CONDA_PREFIX}/include" --install-option="--blas=openblas"
+
+# Or if you want local MinkowskiEngine
+git clone https://github.com/NVIDIA/MinkowskiEngine.git
+cd MinkowskiEngine
+python setup.py install --blas_include_dirs=${CONDA_PREFIX}/include --blas=openblas
+```
+
+### System Python
+
+Like the anaconda installation, make sure that you install pytorch with the same CUDA version that `nvcc` uses.
+
+```
+# install system requirements
+sudo apt install build-essential python3-dev libopenblas-dev
+
+# Skip if you already have pip installed on your python3
+curl https://bootstrap.pypa.io/get-pip.py | python3
+
+# Get pip and install python requirements
+python3 -m pip install torch numpy ninja
+
+git clone https://github.com/NVIDIA/MinkowskiEngine.git
+
+cd MinkowskiEngine
+
+python setup.py install
+# To specify blas, CXX, CUDA_HOME and force CUDA installation, use the following command
+# export CXX=c++; export CUDA_HOME=/usr/local/cuda-11.1; python setup.py install --blas=openblas --force_cuda
+```
+
+### Docker
+
+```
+git clone https://github.com/NVIDIA/MinkowskiEngine
+cd MinkowskiEngine
+docker build -t minkowski_engine docker
+```
+
+Once the docker is built, check it loads MinkowskiEngine correctly.
+
+```
+docker run MinkowskiEngine python3 -c "import MinkowskiEngine; print(MinkowskiEngine.__version__)"
+```
+
+## CPU only build and BLAS configuration (MKL)
+
+The Minkowski Engine supports CPU only build on other platforms that do not have NVidia GPUs. Please refer to [quick start](https://nvidia.github.io/MinkowskiEngine/quick_start.html) for more details.
+
+
+## Quick Start
+
+To use the Minkowski Engine, you first would need to import the engine.
+Then, you would need to define the network. If the data you have is not
+quantized, you would need to voxelize or quantize the (spatial) data into a
+sparse tensor.  Fortunately, the Minkowski Engine provides the quantization
+function (`MinkowskiEngine.utils.sparse_quantize`).
+
+
+### Creating a Network
+
+```python
+import torch.nn as nn
+import MinkowskiEngine as ME
+
+class ExampleNetwork(ME.MinkowskiNetwork):
+
+    def __init__(self, in_feat, out_feat, D):
+        super(ExampleNetwork, self).__init__(D)
+        self.conv1 = nn.Sequential(
+            ME.MinkowskiConvolution(
+                in_channels=in_feat,
+                out_channels=64,
+                kernel_size=3,
+                stride=2,
+                dilation=1,
+                bias=False,
+                dimension=D),
+            ME.MinkowskiBatchNorm(64),
+            ME.MinkowskiReLU())
+        self.conv2 = nn.Sequential(
+            ME.MinkowskiConvolution(
+                in_channels=64,
+                out_channels=128,
+                kernel_size=3,
+                stride=2,
+                dimension=D),
+            ME.MinkowskiBatchNorm(128),
+            ME.MinkowskiReLU())
+        self.pooling = ME.MinkowskiGlobalPooling()
+        self.linear = ME.MinkowskiLinear(128, out_feat)
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.conv2(out)
+        out = self.pooling(out)
+        return self.linear(out)
+```
+
+### Forward and backward using the custom network
+
+```python
+    # loss and network
+    criterion = nn.CrossEntropyLoss()
+    net = ExampleNetwork(in_feat=3, out_feat=5, D=2)
+    print(net)
+
+    # a data loader must return a tuple of coords, features, and labels.
+    coords, feat, label = data_loader()
+    input = ME.SparseTensor(feat, coordinates=coords)
+    # Forward
+    output = net(input)
+
+    # Loss
+    loss = criterion(output.F, label)
+```
+
+## Discussion and Documentation
+
+For discussion and questions, please use `minkowskiengine@googlegroups.com`.
+For API and general usage, please refer to the [MinkowskiEngine documentation
+page](http://nvidia.github.io/MinkowskiEngine/) for more detail.
+
+For issues not listed on the API and feature requests, feel free to submit
+an issue on the [github issue
+page](https://github.com/NVIDIA/MinkowskiEngine/issues).
+
+
+## Known Issues
+
+### Specifying CUDA architecture list
+
+In some cases, you need to explicitly specify which compute capability your GPU uses. The default list might not contain your architecture.
+
+```bash
+export TORCH_CUDA_ARCH_LIST="5.2 6.0 6.1 7.0 7.5 8.0 8.6+PTX"; python setup.py install --force_cuda
+```
+
+### Unhandled Out-Of-Memory thrust::system exception
+
+There is [a known issue](https://github.com/NVIDIA/thrust/issues/1448) in thrust with CUDA 10 that leads to an unhandled thrust exception. Please refer to the [issue](https://github.com/NVIDIA/MinkowskiEngine/issues/357) for detail.
+
+### Too much GPU memory usage or Frequent Out of Memory
+
+There are a few causes for this error.
+
+1. Out of memory during a long running training
+
+MinkowskiEngine is a specialized library that can handle different number of points or different number of non-zero elements at every iteration during training, which is common in point cloud data.
+However, pytorch is implemented assuming that the number of point, or size of the activations do not change at every iteration. Thus, the GPU memory caching used by pytorch can result in unnecessarily large memory consumption.
+
+Specifically, pytorch caches chunks of memory spaces to speed up allocation used in every tensor creation. If it fails to find the memory space, it splits an existing cached memory or allocate new space if there's no cached memory large enough for the requested size. Thus, every time we use different number of point (number of non-zero elements) with pytorch, it either split existing cache or reserve new memory. If the cache is too fragmented and allocated all GPU space, it will raise out of memory error.
+
+**To prevent this, you must clear the cache at regular interval with `torch.cuda.empty_cache()`.**
+
+### CUDA 11.1 Installation
+
+```
+wget https://developer.download.nvidia.com/compute/cuda/11.1.1/local_installers/cuda_11.1.1_455.32.00_linux.run
+sudo sh cuda_11.1.1_455.32.00_linux.run --toolkit --silent --override
+
+# Install MinkowskiEngine with CUDA 11.1
+export CUDA_HOME=/usr/local/cuda-11.1; pip install MinkowskiEngine -v --no-deps
+```
+
+### Running the MinkowskiEngine on nodes with a large number of CPUs
+
+The MinkowskiEngine uses OpenMP to parallelize the kernel map generation. However, when the number of threads used for parallelization is too large (e.g. OMP_NUM_THREADS=80), the efficiency drops rapidly as all threads simply wait for multithread locks to be released.
+In such cases, set the number of threads used for OpenMP. Usually, any number below 24 would be fine, but search for the optimal setup on your system.
+
+```
+export OMP_NUM_THREADS=<number of threads to use>; python <your_program.py>
+```
+
+## Citing Minkowski Engine
+
+If you use the Minkowski Engine, please cite:
+
+- [4D Spatio-Temporal ConvNets: Minkowski Convolutional Neural Networks, CVPR'19](https://arxiv.org/abs/1904.08755), [[pdf]](https://arxiv.org/pdf/1904.08755.pdf)
+
+```
+@inproceedings{choy20194d,
+  title={4D Spatio-Temporal ConvNets: Minkowski Convolutional Neural Networks},
+  author={Choy, Christopher and Gwak, JunYoung and Savarese, Silvio},
+  booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
+  pages={3075--3084},
+  year={2019}
+}
+```
+
+For multi-threaded kernel map generation, please cite:
+
+```
+@inproceedings{choy2019fully,
+  title={Fully Convolutional Geometric Features},
+  author={Choy, Christopher and Park, Jaesik and Koltun, Vladlen},
+  booktitle={Proceedings of the IEEE International Conference on Computer Vision},
+  pages={8958--8966},
+  year={2019}
+}
+```
+
+For strided pooling layers for high-dimensional convolutions, please cite:
+
+```
+@inproceedings{choy2020high,
+  title={High-dimensional Convolutional Networks for Geometric Pattern Recognition},
+  author={Choy, Christopher and Lee, Junha and Ranftl, Rene and Park, Jaesik and Koltun, Vladlen},
+  booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
+  year={2020}
+}
+```
+
+For generative transposed convolution, please cite:
+
+```
+@inproceedings{gwak2020gsdn,
+  title={Generative Sparse Detection Networks for 3D Single-shot Object Detection},
+  author={Gwak, JunYoung and Choy, Christopher B and Savarese, Silvio},
+  booktitle={European conference on computer vision},
+  year={2020}
+}
+```
+
+
+## Unittest
+
+For unittests and gradcheck, use torch >= 1.7
+
+## Projects using Minkowski Engine
+
+Please feel free to update [the wiki page](https://github.com/NVIDIA/MinkowskiEngine/wiki/Usage) to add your projects!
+
+- [Projects using MinkowskiEngine](https://github.com/NVIDIA/MinkowskiEngine/wiki/Usage)
+
+- Segmentation: [3D and 4D Spatio-Temporal Semantic Segmentation, CVPR'19](https://github.com/chrischoy/SpatioTemporalSegmentation)
+- Representation Learning: [Fully Convolutional Geometric Features, ICCV'19](https://github.com/chrischoy/FCGF)
+- 3D Registration: [Learning multiview 3D point cloud registration, CVPR'20](https://arxiv.org/abs/2001.05119)
+- 3D Registration: [Deep Global Registration, CVPR'20](https://arxiv.org/abs/2004.11540)
+- Pattern Recognition: [High-Dimensional Convolutional Networks for Geometric Pattern Recognition, CVPR'20](https://arxiv.org/abs/2005.08144)
+- Detection: [Generative Sparse Detection Networks for 3D Single-shot Object Detection, ECCV'20](https://arxiv.org/abs/2006.12356)
+- Image matching: [Sparse Neighbourhood Consensus Networks, ECCV'20](https://www.di.ens.fr/willow/research/sparse-ncnet/)

MinkowskiEngine/docker/Dockerfile

@@ -0,0 +1,30 @@
+# Use use previous versions, modify these variables
+# ARG PYTORCH="1.9.0"
+# ARG CUDA="11.1"
+
+ARG PYTORCH="1.12.0"
+ARG CUDA="11.3"
+ARG CUDNN="8"
+
+FROM pytorch/pytorch:${PYTORCH}-cuda${CUDA}-cudnn${CUDNN}-devel
+
+##############################################
+# You should modify this to match your GPU compute capability
+# ENV TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0+PTX"
+ENV TORCH_CUDA_ARCH_LIST="6.0 6.1 6.2 7.0 7.2 7.5 8.0 8.6"
+##############################################
+
+ENV TORCH_NVCC_FLAGS="-Xfatbin -compress-all"
+
+# Install dependencies
+RUN apt-get update
+RUN apt-get install -y git ninja-build cmake build-essential libopenblas-dev \
+    xterm xauth openssh-server tmux wget mate-desktop-environment-core
+
+RUN apt-get clean
+RUN rm -rf /var/lib/apt/lists/*
+
+# For faster build, use more jobs.
+ENV MAX_JOBS=4
+RUN git clone --recursive "https://github.com/NVIDIA/MinkowskiEngine"
+RUN cd MinkowskiEngine; python setup.py install --force_cuda --blas=openblas

MinkowskiEngine/docs/.gitignore

@@ -0,0 +1,3 @@
+source
+_build
+_static

MinkowskiEngine/docs/Makefile

@@ -0,0 +1,19 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line.
+SPHINXOPTS    =
+SPHINXBUILD   = sphinx-build
+SOURCEDIR     = .
+BUILDDIR      = _build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)