ocnn 2.2.2__tar.gz → 2.2.4__tar.gz

{ocnn-2.2.2/ocnn.egg-info → ocnn-2.2.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ocnn
-Version: 2.2.2
+Version: 2.2.4
 Summary: Octree-based Sparse Convolutional Neural Networks
 Home-page: https://github.com/octree-nn/ocnn-pytorch
 Author: Peng-Shuai Wang
@@ -23,6 +23,7 @@ Requires-Dist: packaging
 
 [![Documentation Status](https://readthedocs.org/projects/ocnn-pytorch/badge/?version=latest)](https://ocnn-pytorch.readthedocs.io/en/latest/?badge=latest)
 [![Downloads](https://static.pepy.tech/badge/ocnn)](https://pepy.tech/project/ocnn)
+[![Downloads](https://static.pepy.tech/badge/ocnn/month)](https://pepy.tech/project/ocnn)
 [![PyPI](https://img.shields.io/pypi/v/ocnn)](https://pypi.org/project/ocnn/)
 
 This repository contains the **pure PyTorch**-based implementation of
@@ -43,14 +44,14 @@ The key difference is that our O-CNN uses the `octree` to index the sparse
 voxels, while these 3 works use the `Hash Table`.
 
 Our O-CNN is published in SIGGRAPH 2017, H-CNN is published in TVCG 2018,
-SparseConvNet is published in CVPR 2018, and MinkowskiNet is published in 
+SparseConvNet is published in CVPR 2018, and MinkowskiNet is published in
 CVPR 2019. Actually, our O-CNN was submitted to SIGGRAPH in the end of 2016 and
 was officially accepted in March, 2017. The camera-ready version of our O-CNN was
 submitted to SIGGRAPH in April, 2017. We just did not post our paper on Arxiv
 during the review process of SIGGRAPH. Therefore, **the idea of constraining CNN
 computation into sparse non-emtpry voxels is first proposed by our O-CNN**.
 Currently, this type of 3D convolution is known as Sparse Convolution in the
-research community. 
+research community.
 
 ## Key benefits of ocnn-pytorch
 
@@ -65,3 +66,15 @@ research community.
   training settings, MinkowskiNet 0.4.3 takes 60 hours and MinkowskiNet 0.5.4
   takes 30 hours.
 
+## Citation
+
+  ```bibtex
+  @article {Wang-2017-ocnn,
+    title    = {{O-CNN}: Octree-based Convolutional Neural Networksfor {3D} Shape Analysis},
+    author   = {Wang, Peng-Shuai and Liu, Yang and Guo, Yu-Xiao and Sun, Chun-Yu and Tong, Xin},
+    journal  = {ACM Transactions on Graphics (SIGGRAPH)},
+    volume   = {36},
+    number   = {4},
+    year     = {2017},
+  }
+  ```

{ocnn-2.2.2 → ocnn-2.2.4}/README.md

@@ -4,6 +4,7 @@
 
 [![Documentation Status](https://readthedocs.org/projects/ocnn-pytorch/badge/?version=latest)](https://ocnn-pytorch.readthedocs.io/en/latest/?badge=latest)
 [![Downloads](https://static.pepy.tech/badge/ocnn)](https://pepy.tech/project/ocnn)
+[![Downloads](https://static.pepy.tech/badge/ocnn/month)](https://pepy.tech/project/ocnn)
 [![PyPI](https://img.shields.io/pypi/v/ocnn)](https://pypi.org/project/ocnn/)
 
 This repository contains the **pure PyTorch**-based implementation of
@@ -24,14 +25,14 @@ The key difference is that our O-CNN uses the `octree` to index the sparse
 voxels, while these 3 works use the `Hash Table`.
 
 Our O-CNN is published in SIGGRAPH 2017, H-CNN is published in TVCG 2018,
-SparseConvNet is published in CVPR 2018, and MinkowskiNet is published in 
+SparseConvNet is published in CVPR 2018, and MinkowskiNet is published in
 CVPR 2019. Actually, our O-CNN was submitted to SIGGRAPH in the end of 2016 and
 was officially accepted in March, 2017. The camera-ready version of our O-CNN was
 submitted to SIGGRAPH in April, 2017. We just did not post our paper on Arxiv
 during the review process of SIGGRAPH. Therefore, **the idea of constraining CNN
 computation into sparse non-emtpry voxels is first proposed by our O-CNN**.
 Currently, this type of 3D convolution is known as Sparse Convolution in the
-research community. 
+research community.
 
 ## Key benefits of ocnn-pytorch
 
@@ -46,3 +47,15 @@ research community.
   training settings, MinkowskiNet 0.4.3 takes 60 hours and MinkowskiNet 0.5.4
   takes 30 hours.
 
+## Citation
+
+  ```bibtex
+  @article {Wang-2017-ocnn,
+    title    = {{O-CNN}: Octree-based Convolutional Neural Networksfor {3D} Shape Analysis},
+    author   = {Wang, Peng-Shuai and Liu, Yang and Guo, Yu-Xiao and Sun, Chun-Yu and Tong, Xin},
+    journal  = {ACM Transactions on Graphics (SIGGRAPH)},
+    volume   = {36},
+    number   = {4},
+    year     = {2017},
+  }
+  ```

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/__init__.py

@@ -12,7 +12,7 @@ from . import models
 from . import dataset
 from . import utils
 
-__version__ = '2.2.2'
+__version__ = '2.2.4'
 
 __all__ = [
     'octree',

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/models/autoencoder.py

@@ -7,7 +7,6 @@
 
 import torch
 import torch.nn
-from typing import Optional
 
 import ocnn
 from ocnn.octree import Octree

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/models/lenet.py

@@ -26,7 +26,7 @@ class LeNet(torch.nn.Module):
     self.convs = torch.nn.ModuleList([ocnn.modules.OctreeConvBnRelu(
         channels[i], channels[i+1], nempty=nempty) for i in range(stages)])
     self.pools = torch.nn.ModuleList([ocnn.nn.OctreeMaxPool(
-        nempty) for i in range(stages)])
+        nempty) for _ in range(stages)])
     self.octree2voxel = ocnn.nn.Octree2Voxel(self.nempty)
     self.header = torch.nn.Sequential(
         torch.nn.Dropout(p=0.5),                     # drop1

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/models/ounet.py

@@ -23,12 +23,13 @@ class OUNet(AutoEncoder):
     self.proj = None  # remove this module used in AutoEncoder
 
   def encoder(self, octree):
-    r''' The encoder network for extracting heirarchy features. 
+    r''' The encoder network for extracting heirarchy features.
     '''
 
     convs = dict()
     depth, full_depth = self.depth, self.full_depth
-    data = self.get_input_feature(octree)
+    data = octree.get_input_feature(self.feature, nempty=False)
+    assert data.size(1) == self.channel_in
     convs[depth] = self.conv1(data, octree, depth)
     for i, d in enumerate(range(depth, full_depth-1, -1)):
       convs[d] = self.encoder_blks[i](convs[d], octree, d)

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/models/resnet.py

@@ -31,7 +31,7 @@ class ResNet(torch.nn.Module):
         channels[i], channels[i+1], resblock_num, nempty=nempty)
         for i in range(stages-1)])
     self.pools = torch.nn.ModuleList([ocnn.nn.OctreeMaxPool(
-        nempty) for i in range(stages-1)])
+        nempty) for _ in range(stages-1)])
     self.global_pool = ocnn.nn.OctreeGlobalPool(nempty)
     # self.header = torch.nn.Linear(channels[-1], out_channels, bias=True)
     self.header = torch.nn.Sequential(

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/models/segnet.py

@@ -29,7 +29,7 @@ class SegNet(torch.nn.Module):
         [ocnn.modules.OctreeConvBnRelu(channels[i], channels[i+1], nempty=nempty)
          for i in range(stages)])
     self.pools = torch.nn.ModuleList(
-        [ocnn.nn.OctreeMaxPool(nempty, return_indices) for i in range(stages)])
+        [ocnn.nn.OctreeMaxPool(nempty, return_indices) for _ in range(stages)])
 
     self.bottleneck = ocnn.modules.OctreeConvBnRelu(channels[-1], channels[-1])
 
@@ -38,7 +38,7 @@ class SegNet(torch.nn.Module):
         [ocnn.modules.OctreeConvBnRelu(channels[i], channels[i+1], nempty=nempty)
          for i in range(0, stages)])
     self.unpools = torch.nn.ModuleList(
-        [ocnn.nn.OctreeMaxUnpool(nempty) for i in range(stages)])
+        [ocnn.nn.OctreeMaxUnpool(nempty) for _ in range(stages)])
 
     self.octree_interp = ocnn.nn.OctreeInterp(interp, nempty)
     self.header = torch.nn.Sequential(

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/modules/__init__.py

@@ -7,14 +7,20 @@
 
 from .modules import (InputFeature,
                       OctreeConvBn, OctreeConvBnRelu, OctreeDeconvBnRelu,
-                      Conv1x1, Conv1x1Bn, Conv1x1BnRelu, FcBnRelu,)
-from .resblocks import OctreeResBlock, OctreeResBlock2, OctreeResBlocks
+                      Conv1x1, Conv1x1Bn, Conv1x1BnRelu, FcBnRelu,
+                      OctreeConvGn, OctreeConvGnRelu, OctreeDeconvGnRelu,
+                      Conv1x1, Conv1x1Gn, Conv1x1GnRelu)
+from .resblocks import (OctreeResBlock, OctreeResBlock2, OctreeResBlockGn,
+                        OctreeResBlocks,)
 
 __all__ = [
     'InputFeature',
     'OctreeConvBn', 'OctreeConvBnRelu', 'OctreeDeconvBnRelu',
     'Conv1x1', 'Conv1x1Bn', 'Conv1x1BnRelu', 'FcBnRelu',
-    'OctreeResBlock', 'OctreeResBlock2', 'OctreeResBlocks',
+    'OctreeConvGn', 'OctreeConvGnRelu', 'OctreeDeconvGnRelu',
+    'Conv1x1', 'Conv1x1Gn', 'Conv1x1GnRelu',
+    'OctreeResBlock', 'OctreeResBlock2', 'OctreeResBlockGn',
+    'OctreeResBlocks',
 ]
 
 classes = __all__

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/modules/modules.py

@@ -9,7 +9,7 @@ import torch
 import torch.utils.checkpoint
 from typing import List
 
-from ocnn.nn import OctreeConv, OctreeDeconv
+from ocnn.nn import OctreeConv, OctreeDeconv, OctreeGroupNorm
 from ocnn.octree import Octree
 
 
@@ -40,7 +40,7 @@ class OctreeConvBn(torch.nn.Module):
     super().__init__()
     self.conv = OctreeConv(
         in_channels, out_channels, kernel_size, stride, nempty)
-    self.bn = torch.nn.BatchNorm1d(out_channels) #, bn_eps, bn_momentum)
+    self.bn = torch.nn.BatchNorm1d(out_channels)
 
   def forward(self, data: torch.Tensor, octree: Octree, depth: int):
     r''''''
@@ -62,7 +62,7 @@ class OctreeConvBnRelu(torch.nn.Module):
     super().__init__()
     self.conv = OctreeConv(
         in_channels, out_channels, kernel_size, stride, nempty)
-    self.bn = torch.nn.BatchNorm1d(out_channels) #, bn_eps, bn_momentum)
+    self.bn = torch.nn.BatchNorm1d(out_channels)
     self.relu = torch.nn.ReLU(inplace=True)
 
   def forward(self, data: torch.Tensor, octree: Octree, depth: int):
@@ -86,7 +86,7 @@ class OctreeDeconvBnRelu(torch.nn.Module):
     super().__init__()
     self.deconv = OctreeDeconv(
         in_channels, out_channels, kernel_size, stride, nempty)
-    self.bn = torch.nn.BatchNorm1d(out_channels) #, bn_eps, bn_momentum)
+    self.bn = torch.nn.BatchNorm1d(out_channels)
     self.relu = torch.nn.ReLU(inplace=True)
 
   def forward(self, data: torch.Tensor, octree: Octree, depth: int):
@@ -123,7 +123,7 @@ class Conv1x1Bn(torch.nn.Module):
   def __init__(self, in_channels: int, out_channels: int):
     super().__init__()
     self.conv = Conv1x1(in_channels, out_channels, use_bias=False)
-    self.bn = torch.nn.BatchNorm1d(out_channels) #, bn_eps, bn_momentum)
+    self.bn = torch.nn.BatchNorm1d(out_channels)
 
   def forward(self, data: torch.Tensor):
     r''''''
@@ -140,7 +140,7 @@ class Conv1x1BnRelu(torch.nn.Module):
   def __init__(self, in_channels: int, out_channels: int):
     super().__init__()
     self.conv = Conv1x1(in_channels, out_channels, use_bias=False)
-    self.bn = torch.nn.BatchNorm1d(out_channels) #, bn_eps, bn_momentum)
+    self.bn = torch.nn.BatchNorm1d(out_channels)
     self.relu = torch.nn.ReLU(inplace=True)
 
   def forward(self, data: torch.Tensor):
@@ -160,7 +160,7 @@ class FcBnRelu(torch.nn.Module):
     super().__init__()
     self.flatten = torch.nn.Flatten(start_dim=1)
     self.fc = torch.nn.Linear(in_channels, out_channels, bias=False)
-    self.bn = torch.nn.BatchNorm1d(out_channels) #, bn_eps, bn_momentum)
+    self.bn = torch.nn.BatchNorm1d(out_channels)
     self.relu = torch.nn.ReLU(inplace=True)
 
   def forward(self, data):
@@ -173,6 +173,116 @@ class FcBnRelu(torch.nn.Module):
     return out
 
 
+class OctreeConvGn(torch.nn.Module):
+  r''' A sequence of :class:`OctreeConv` and :obj:`OctreeGroupNorm`.
+
+  Please refer to :class:`ocnn.nn.OctreeConv` for details on the parameters.
+  '''
+
+  def __init__(self, in_channels: int, out_channels: int, group: int,
+               kernel_size: List[int] = [3], stride: int = 1,
+               nempty: bool = False):
+    super().__init__()
+    self.conv = OctreeConv(
+        in_channels, out_channels, kernel_size, stride, nempty)
+    self.gn = OctreeGroupNorm(out_channels, group=group, nempty=nempty)
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+
+    out = self.conv(data, octree, depth)
+    out = self.gn(out, octree, depth)
+    return out
+
+
+class OctreeConvGnRelu(torch.nn.Module):
+  r''' A sequence of :class:`OctreeConv`, :obj:`OctreeGroupNorm`, and :obj:`Relu`.
+
+  Please refer to :class:`ocnn.nn.OctreeConv` for details on the parameters.
+  '''
+
+  def __init__(self, in_channels: int, out_channels: int, group: int,
+               kernel_size: List[int] = [3], stride: int = 1,
+               nempty: bool = False):
+    super().__init__()
+    self.stride = stride
+    self.conv = OctreeConv(
+        in_channels, out_channels, kernel_size, stride, nempty)
+    self.gn = OctreeGroupNorm(out_channels, group=group, nempty=nempty)
+    self.relu = torch.nn.ReLU(inplace=True)
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+
+    out = self.conv(data, octree, depth)
+    out = self.gn(out, octree, depth if self.stride == 1 else depth - 1)
+    out = self.relu(out)
+    return out
+
+
+class OctreeDeconvGnRelu(torch.nn.Module):
+  r''' A sequence of :class:`OctreeDeconv`, :obj:`OctreeGroupNorm`, and :obj:`Relu`.
+
+  Please refer to :class:`ocnn.nn.OctreeConv` for details on the parameters.
+  '''
+
+  def __init__(self, in_channels: int, out_channels: int, group: int,
+               kernel_size: List[int] = [3], stride: int = 1,
+               nempty: bool = False):
+    super().__init__()
+    self.stride = stride
+    self.deconv = OctreeDeconv(
+        in_channels, out_channels, kernel_size, stride, nempty)
+    self.gn = OctreeGroupNorm(out_channels, group=group, nempty=nempty)
+    self.relu = torch.nn.ReLU(inplace=True)
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+
+    out = self.deconv(data, octree, depth)
+    out = self.gn(out, octree, depth if self.stride == 1 else depth + 1)
+    out = self.relu(out)
+    return out
+
+
+class Conv1x1Gn(torch.nn.Module):
+  r''' A sequence of :class:`Conv1x1`, :class:`OctreeGroupNorm`.
+  '''
+
+  def __init__(self, in_channels: int, out_channels: int, group: int,
+               nempty: bool = False):
+    super().__init__()
+    self.conv = Conv1x1(in_channels, out_channels, use_bias=False)
+    self.gn = OctreeGroupNorm(out_channels, group=group, nempty=nempty)
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+
+    out = self.conv(data)
+    out = self.gn(out, octree, depth)
+    return out
+
+
+class Conv1x1GnRelu(torch.nn.Module):
+  r''' A sequence of :class:`Conv1x1`, :class:`OctreeGroupNorm` and :class:`Relu`.
+  '''
+
+  def __init__(self, in_channels: int, out_channels: int, group: int,
+               nempty: bool = False):
+    super().__init__()
+    self.conv = Conv1x1(in_channels, out_channels, use_bias=False)
+    self.gn = OctreeGroupNorm(out_channels, group=group, nempty=nempty)
+    self.relu = torch.nn.ReLU(inplace=True)
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+
+    out = self.conv(data)
+    out = self.gn(out, octree, depth)
+    out = self.relu(out)
+    return out
+
+
 class InputFeature(torch.nn.Module):
   r''' Returns the initial input feature stored in octree.
 

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/modules/resblocks.py

@@ -10,7 +10,9 @@ import torch.utils.checkpoint
 
 from ocnn.octree import Octree
 from ocnn.nn import OctreeMaxPool
-from ocnn.modules import Conv1x1BnRelu, OctreeConvBnRelu, Conv1x1Bn, OctreeConvBn
+from ocnn.modules import (Conv1x1BnRelu, OctreeConvBnRelu, Conv1x1Bn,
+                          OctreeConvBn, OctreeConvGnRelu, Conv1x1Gn,
+                          OctreeConvGn,)
 
 
 class OctreeResBlock(torch.nn.Module):
@@ -97,6 +99,38 @@ class OctreeResBlock2(torch.nn.Module):
     return out
 
 
+class OctreeResBlockGn(torch.nn.Module):
+
+  def __init__(self, in_channels: int, out_channels: int, stride: int = 1,
+               bottleneck: int = 4, nempty: bool = False, group: int = 32):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.stride = stride
+    channelb = int(out_channels / bottleneck)
+
+    if self.stride == 2:
+      self.maxpool = OctreeMaxPool(self.depth)
+    self.conv3x3a = OctreeConvGnRelu(in_channels, channelb, group, nempty=nempty)
+    self.conv3x3b = OctreeConvGn(channelb, out_channels, group, nempty=nempty)
+    if self.in_channels != self.out_channels:
+      self.conv1x1 = Conv1x1Gn(in_channels, out_channels, group)
+    self.relu = torch.nn.ReLU(inplace=True)
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+
+    if self.stride == 2:
+      data = self.maxpool(data, octree, depth)
+      depth = depth - 1
+    conv1 = self.conv3x3a(data, octree, depth)
+    conv2 = self.conv3x3b(conv1, octree, depth)
+    if self.in_channels != self.out_channels:
+      data = self.conv1x1(data, octree, depth)
+    out = self.relu(conv2 + data)
+    return out
+
+
 class OctreeResBlocks(torch.nn.Module):
   r''' A sequence of :attr:`resblk_num` ResNet blocks.
   '''
@@ -108,9 +142,9 @@ class OctreeResBlocks(torch.nn.Module):
     self.use_checkpoint = use_checkpoint
     channels = [in_channels] + [out_channels] * resblk_num
 
-    self.resblks = torch.nn.ModuleList(
-        [resblk(channels[i], channels[i+1], 1, bottleneck, nempty)
-         for i in range(self.resblk_num)])
+    self.resblks = torch.nn.ModuleList([resblk(
+        channels[i], channels[i+1], 1, bottleneck, nempty)
+        for i in range(self.resblk_num)])
 
   def forward(self, data: torch.Tensor, octree: Octree, depth: int):
     r''''''
@@ -118,7 +152,7 @@ class OctreeResBlocks(torch.nn.Module):
     for i in range(self.resblk_num):
       if self.use_checkpoint:
         data = torch.utils.checkpoint.checkpoint(
-            self.resblks[i], data, octree, depth)
+            self.resblks[i], data, octree, depth, use_reentrant=False)
       else:
         data = self.resblks[i](data, octree, depth)
     return data

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/nn/__init__.py

@@ -15,8 +15,10 @@ from .octree_pool import (octree_max_pool, OctreeMaxPool,
                           octree_global_pool, OctreeGlobalPool,
                           octree_avg_pool, OctreeAvgPool,)
 from .octree_conv import OctreeConv, OctreeDeconv
+from .octree_gconv import OctreeGroupConv
 from .octree_dwconv import OctreeDWConv
-from .octree_norm import OctreeBatchNorm, OctreeGroupNorm, OctreeInstanceNorm
+from .octree_norm import (OctreeBatchNorm, OctreeGroupNorm,
+                          OctreeInstanceNorm, OctreeNorm)
 from .octree_drop import OctreeDropPath
 from .octree_align import search_value, octree_align
 
@@ -32,9 +34,9 @@ __all__ = [
     'OctreeMaxPool', 'OctreeMaxUnpool',
     'OctreeGlobalPool', 'OctreeAvgPool',
     'OctreeConv', 'OctreeDeconv',
-    'OctreeDWConv',
+    'OctreeGroupConv', 'OctreeDWConv',
     'OctreeInterp', 'OctreeUpsample',
-    'OctreeInstanceNorm', 'OctreeBatchNorm', 'OctreeGroupNorm',
+    'OctreeInstanceNorm', 'OctreeBatchNorm', 'OctreeGroupNorm', 'OctreeNorm',
     'OctreeDropPath',
     'search_value', 'octree_align',
 ]

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/nn/octree_conv.py

@@ -109,6 +109,12 @@ class OctreeConvBase:
     r''' Peforms the forward pass of octree-based convolution.
     '''
 
+    # Type check
+    if data.dtype != out.dtype:
+      data = data.to(out.dtype)
+    if weights.dtype != out.dtype:
+      weights = weights.to(out.dtype)
+
     # Initialize the buffer
     buffer = data.new_empty(self.buffer_shape)
 
@@ -139,6 +145,12 @@ class OctreeConvBase:
     r''' Performs the backward pass of octree-based convolution.
     '''
 
+    # Type check
+    if grad.dtype != out.dtype:
+      grad = grad.to(out.dtype)
+    if weights.dtype != out.dtype:
+      weights = weights.to(out.dtype)
+
     # Loop over each sub-matrix
     for i in range(self.buffer_n):
       start = i * self.buffer_h
@@ -165,6 +177,12 @@ class OctreeConvBase:
     r''' Computes the gradient of the weight matrix.
     '''
 
+    # Type check
+    if data.dtype != out.dtype:
+      data = data.to(out.dtype)
+    if grad.dtype != out.dtype:
+      grad = grad.to(out.dtype)
+
     # Record the shape of out
     out_shape = out.shape
     out = out.flatten(0, 1)

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/nn/octree_dwconv.py

@@ -32,6 +32,12 @@ class OctreeDWConvBase(OctreeConvBase):
     r''' Peforms the forward pass of octree-based convolution.
     '''
 
+    # Type check
+    if data.dtype != out.dtype:
+      data = data.to(out.dtype)
+    if weights.dtype != out.dtype:
+      weights = weights.to(out.dtype)
+
     # Initialize the buffer
     buffer = data.new_empty(self.buffer_shape)
 
@@ -62,6 +68,12 @@ class OctreeDWConvBase(OctreeConvBase):
     r''' Performs the backward pass of octree-based convolution. 
     '''
 
+    # Type check
+    if grad.dtype != out.dtype:
+      grad = grad.to(out.dtype)
+    if weights.dtype != out.dtype:
+      weights = weights.to(out.dtype)
+
     # Loop over each sub-matrix
     for i in range(self.buffer_n):
       start = i * self.buffer_h
@@ -88,6 +100,12 @@ class OctreeDWConvBase(OctreeConvBase):
     r''' Computes the gradient of the weight matrix.
     '''
 
+    # Type check
+    if data.dtype != out.dtype:
+      data = data.to(out.dtype)
+    if grad.dtype != out.dtype:
+      grad = grad.to(out.dtype)
+
     # Record the shape of out
     out_shape = out.shape
     out = out.flatten(0, 1)

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/nn/octree_norm.py

@@ -7,6 +7,7 @@
 
 import torch
 import torch.nn
+from typing import Optional
 
 from ocnn.octree import Octree
 from ocnn.utils import scatter_add
@@ -19,15 +20,19 @@ class OctreeGroupNorm(torch.nn.Module):
   r''' An group normalization layer for the octree.
   '''
 
-  def __init__(self, in_channels: int, group: int, nempty: bool = False):
+  def __init__(self, in_channels: int, group: int, nempty: bool = False,
+               min_group_channels: int = 4):
     super().__init__()
     self.eps = 1e-5
     self.nempty = nempty
     self.group = group
     self.in_channels = in_channels
+    self.min_group_channels = min_group_channels
+    if self.min_group_channels * self.group > in_channels:
+      self.group = in_channels // self.min_group_channels
 
-    assert in_channels % group == 0
-    self.channels_per_group = in_channels // group
+    assert in_channels % self.group == 0
+    self.channels_per_group = in_channels // self.group
 
     self.weights = torch.nn.Parameter(torch.Tensor(1, in_channels))
     self.bias = torch.nn.Parameter(torch.Tensor(1, in_channels))
@@ -71,8 +76,8 @@ class OctreeGroupNorm(torch.nn.Module):
     return tensor
 
   def extra_repr(self) -> str:
-    return ('in_channels={}, group={}, nempty={}').format(
-            self.in_channels, self.group, self.nempty)  # noqa
+    return ('in_channels={}, group={}, nempty={}, min_group_channels={}').format(
+        self.in_channels, self.group, self.nempty, self.min_group_channels)
 
 
 class OctreeInstanceNorm(OctreeGroupNorm):
@@ -80,7 +85,42 @@ class OctreeInstanceNorm(OctreeGroupNorm):
   '''
 
   def __init__(self, in_channels: int, nempty: bool = False):
-    super().__init__(in_channels=in_channels, group=in_channels, nempty=nempty)
+    super().__init__(in_channels=in_channels, group=in_channels, nempty=nempty,
+                     min_group_channels=1)  # NOTE: group=in_channels
 
   def extra_repr(self) -> str:
     return ('in_channels={}, nempty={}').format(self.in_channels, self.nempty)
+
+
+class OctreeNorm(torch.nn.Module):
+  r''' A normalization layer for the octree. It encapsulates octree-based batch,
+  group and instance normalization.
+  '''
+
+  def __init__(self, in_channels: int, norm_type: str = 'batch_norm',
+               group: int = 32, min_group_channels: int = 4):
+    super().__init__()
+    self.in_channels = in_channels
+    self.norm_type = norm_type
+    self.group = group
+    self.min_group_channels = min_group_channels
+
+    if self.norm_type == 'batch_norm':
+      self.norm = torch.nn.BatchNorm1d(in_channels)
+    elif self.norm_type == 'group_norm':
+      self.norm = OctreeGroupNorm(in_channels, group, min_group_channels)
+    elif self.norm_type == 'instance_norm':
+      self.norm = OctreeInstanceNorm(in_channels)
+    else:
+      raise ValueError
+
+  def forward(self, x: torch.Tensor, octree: Optional[Octree] = None,
+              depth: Optional[int] = None):
+    if self.norm_type == 'batch_norm':
+      output = self.norm(x)
+    elif (self.norm_type == 'group_norm' or
+          self.norm_type == 'instance_norm'):
+      output = self.norm(x, octree, depth)
+    else:
+      raise ValueError
+    return output

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/octree/octree.py

@@ -63,14 +63,14 @@ class Octree:
 
     # octree node numbers in each octree layers.
     # TODO: decide whether to settle them to 'gpu' or not?
-    self.nnum = torch.zeros(num, dtype=torch.int32)
-    self.nnum_nempty = torch.zeros(num, dtype=torch.int32)
+    self.nnum = torch.zeros(num, dtype=torch.long)
+    self.nnum_nempty = torch.zeros(num, dtype=torch.long)
 
     # the following properties are valid after `merge_octrees`.
     # TODO: make them valid after `octree_grow`, `octree_split` and `build_octree`
     batch_size = self.batch_size
-    self.batch_nnum = torch.zeros(num, batch_size, dtype=torch.int32)
-    self.batch_nnum_nempty = torch.zeros(num, batch_size, dtype=torch.int32)
+    self.batch_nnum = torch.zeros(num, batch_size, dtype=torch.long)
+    self.batch_nnum_nempty = torch.zeros(num, batch_size, dtype=torch.long)
 
     # construct the look up tables for neighborhood searching
     device = self.device
@@ -274,7 +274,7 @@ class Octree:
       children[0] = 0
 
     # update octree
-    self.children[depth] = children
+    self.children[depth] = children.int()
     self.nnum_nempty[depth] = nnum_nempty
 
   def octree_grow(self, depth: int, update_neigh: bool = True):
@@ -290,7 +290,7 @@ class Octree:
     # node number
     nnum = self.nnum_nempty[depth-1] * 8
     self.nnum[depth] = nnum
-    self.nnum_nempty[depth] = nnum
+    self.nnum_nempty[depth] = nnum  # initialize self.nnum_nempty
 
     # update keys
     key = self.key(depth-1, nempty=True)
@@ -326,7 +326,7 @@ class Octree:
       xyz = xyz.view(-1, 3)                 # (N*27, 3)
       neigh = xyz2key(xyz[:, 0], xyz[:, 1], xyz[:, 2], depth=depth)
 
-      bs = torch.arange(self.batch_size, dtype=torch.int32, device=device)
+      bs = torch.arange(self.batch_size, dtype=torch.long, device=device)
       neigh = neigh + bs.unsqueeze(1) * nnum  # (N*27,) + (B, 1) -> (B, N*27)
 
       bound = 1 << depth
@@ -383,9 +383,10 @@ class Octree:
     # I choose `torch.bucketize` here because it has fewer dimension checks,
     # resulting in slightly better performance according to the docs of
     # pytorch-1.9.1, since `key` is always 1-D sorted sequence.
+    # https://pytorch.org/docs/1.9.1/generated/torch.searchsorted.html
     idx = torch.bucketize(query, key)
 
-    valid = idx < key.shape[0]      # valid if NOT out-of-bound
+    valid = idx < key.shape[0]      # valid if in-bound
     found = key[idx[valid]] == query[valid]
     valid[valid.clone()] = found    # valid if found
     idx[valid.logical_not()] = -1   # set to -1 if invalid
@@ -498,7 +499,7 @@ class Octree:
     # normalize xyz to [-1, 1] since the average points are in range [0, 2^d]
     if rescale:
       scale = 2 ** (1 - depth)
-      xyz = self.points[depth] * scale - 1.0
+      xyz = xyz * scale - 1.0
 
     # construct Points
     out = Points(xyz, self.normals[depth], self.features[depth],

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/octree/points.py

@@ -56,11 +56,16 @@ class Points:
       assert self.features.dim() == 2
       assert self.features.size(0) == self.points.size(0)
     if self.labels is not None:
-      assert self.labels.dim() == 2
+      assert self.labels.dim() == 2 or self.labels.dim() == 1
       assert self.labels.size(0) == self.points.size(0)
+      if self.labels.dim() == 1:
+        self.labels = self.labels.unsqueeze(1)
     if self.batch_id is not None:
-      assert self.batch_id.dim() == 2 and self.batch_id.size(1) == 1
+      assert self.batch_id.dim() == 2 or self.batch_id.dim() == 1
       assert self.batch_id.size(0) == self.points.size(0)
+      assert self.batch_id.size(1) == 1
+      if self.batch_id.dim() == 1:
+        self.batch_id = self.batch_id.unsqueeze(1)
 
   @property
   def npt(self):

{ocnn-2.2.2 → ocnn-2.2.4}/ocnn/utils.py

@@ -173,7 +173,7 @@ def resize_with_last_val(list_in: list, num: int = 3):
   '''
 
   assert (type(list_in) is list and len(list_in) < num + 1)
-  for i in range(len(list_in), num):
+  for _ in range(len(list_in), num):
     list_in.append(list_in[-1])
   return list_in
 
@@ -186,15 +186,18 @@ def list2str(list_in: list):
   return ''.join(out)
 
 
-def build_example_octree(depth: int = 5, full_depth: int = 2):
-  r''' Builds an example octree on CPU from 3 points.
+def build_example_octree(depth: int = 5, full_depth: int = 2, pt_num: int = 3):
+  r''' Builds an example octree on CPU from at most 3 points.
   '''
   # initialize the point cloud
   points = torch.Tensor([[-1, -1, -1], [0, 0, -1], [0.0625, 0.0625, -1]])
   normals = torch.Tensor([[1, 0, 0], [-1, 0, 0], [0, 1, 0]])
   features = torch.Tensor([[1, -1], [2, -2], [3, -3]])
   labels = torch.Tensor([[0], [2], [2]])
-  point_cloud = ocnn.octree.Points(points, normals, features, labels)
+
+  assert pt_num <= 3 and pt_num > 0
+  point_cloud = ocnn.octree.Points(
+      points[:pt_num], normals[:pt_num], features[:pt_num], labels[:pt_num])
 
   # build octree
   octree = ocnn.octree.Octree(depth, full_depth)

{ocnn-2.2.2 → ocnn-2.2.4/ocnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ocnn
-Version: 2.2.2
+Version: 2.2.4
 Summary: Octree-based Sparse Convolutional Neural Networks
 Home-page: https://github.com/octree-nn/ocnn-pytorch
 Author: Peng-Shuai Wang
@@ -23,6 +23,7 @@ Requires-Dist: packaging
 
 [![Documentation Status](https://readthedocs.org/projects/ocnn-pytorch/badge/?version=latest)](https://ocnn-pytorch.readthedocs.io/en/latest/?badge=latest)
 [![Downloads](https://static.pepy.tech/badge/ocnn)](https://pepy.tech/project/ocnn)
+[![Downloads](https://static.pepy.tech/badge/ocnn/month)](https://pepy.tech/project/ocnn)
 [![PyPI](https://img.shields.io/pypi/v/ocnn)](https://pypi.org/project/ocnn/)
 
 This repository contains the **pure PyTorch**-based implementation of
@@ -43,14 +44,14 @@ The key difference is that our O-CNN uses the `octree` to index the sparse
 voxels, while these 3 works use the `Hash Table`.
 
 Our O-CNN is published in SIGGRAPH 2017, H-CNN is published in TVCG 2018,
-SparseConvNet is published in CVPR 2018, and MinkowskiNet is published in 
+SparseConvNet is published in CVPR 2018, and MinkowskiNet is published in
 CVPR 2019. Actually, our O-CNN was submitted to SIGGRAPH in the end of 2016 and
 was officially accepted in March, 2017. The camera-ready version of our O-CNN was
 submitted to SIGGRAPH in April, 2017. We just did not post our paper on Arxiv
 during the review process of SIGGRAPH. Therefore, **the idea of constraining CNN
 computation into sparse non-emtpry voxels is first proposed by our O-CNN**.
 Currently, this type of 3D convolution is known as Sparse Convolution in the
-research community. 
+research community.
 
 ## Key benefits of ocnn-pytorch
 
@@ -65,3 +66,15 @@ research community.
   training settings, MinkowskiNet 0.4.3 takes 60 hours and MinkowskiNet 0.5.4
   takes 30 hours.
 
+## Citation
+
+  ```bibtex
+  @article {Wang-2017-ocnn,
+    title    = {{O-CNN}: Octree-based Convolutional Neural Networksfor {3D} Shape Analysis},
+    author   = {Wang, Peng-Shuai and Liu, Yang and Guo, Yu-Xiao and Sun, Chun-Yu and Tong, Xin},
+    journal  = {ACM Transactions on Graphics (SIGGRAPH)},
+    volume   = {36},
+    number   = {4},
+    year     = {2017},
+  }
+  ```

{ocnn-2.2.2 → ocnn-2.2.4}/setup.py

@@ -7,7 +7,7 @@
 
 from setuptools import setup, find_packages
 
-__version__ = '2.2.2'
+__version__ = '2.2.4'
 
 with open("README.md", "r", encoding="utf-8") as fid:
   long_description = fid.read()