ocnn 2.2.3__tar.gz → 2.2.5__tar.gz

{ocnn-2.2.3/ocnn.egg-info → ocnn-2.2.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ocnn
-Version: 2.2.3
+Version: 2.2.5
 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

{ocnn-2.2.3 → ocnn-2.2.5}/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

{ocnn-2.2.3 → ocnn-2.2.5}/ocnn/__init__.py

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

{ocnn-2.2.3 → ocnn-2.2.5}/ocnn/models/autoencoder.py

@@ -7,7 +7,6 @@
 
 import torch
 import torch.nn
-from typing import Optional
 
 import ocnn
 from ocnn.octree import Octree
@@ -34,8 +33,9 @@ class AutoEncoder(torch.nn.Module):
     self.full_depth = full_depth
     self.feature = feature
     self.resblk_num = 2
-    self.code_channel = 64  # dim-of-code = code_channel * 2**(3*full_depth)
     self.channels = [512, 512, 256, 256, 128, 128, 32, 32, 16, 16]
+    # dim-of-code = code_channel * 2**(3*full_depth)
+    self.code_channel = self.channels[full_depth]
 
     # encoder
     self.conv1 = ocnn.modules.OctreeConvBnRelu(

{ocnn-2.2.3 → ocnn-2.2.5}/ocnn/models/ounet.py

@@ -18,17 +18,17 @@ class OUNet(AutoEncoder):
 
   def __init__(self, channel_in: int, channel_out: int, depth: int,
                full_depth: int = 2, feature: str = 'ND'):
-    super().__init__(channel_in, channel_out, depth, full_depth, feature,
-                     code_channel=-1)  # !set code_channe=-1
+    super().__init__(channel_in, channel_out, depth, full_depth, feature)
     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.3 → ocnn-2.2.5}/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.3 → ocnn-2.2.5}/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.3 → ocnn-2.2.5}/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.3 → ocnn-2.2.5}/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.3 → ocnn-2.2.5}/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.3 → ocnn-2.2.5}/ocnn/nn/__init__.py

@@ -17,7 +17,8 @@ from .octree_pool import (octree_max_pool, OctreeMaxPool,
 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
 
@@ -35,7 +36,7 @@ __all__ = [
     'OctreeConv', 'OctreeDeconv',
     'OctreeGroupConv', 'OctreeDWConv',
     'OctreeInterp', 'OctreeUpsample',
-    'OctreeInstanceNorm', 'OctreeBatchNorm', 'OctreeGroupNorm',
+    'OctreeInstanceNorm', 'OctreeBatchNorm', 'OctreeGroupNorm', 'OctreeNorm',
     'OctreeDropPath',
     'search_value', 'octree_align',
 ]

{ocnn-2.2.3 → ocnn-2.2.5}/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.3 → ocnn-2.2.5}/ocnn/octree/octree.py

@@ -35,7 +35,9 @@ class Octree:
     and :obj:`points`, contain only non-empty nodes.
 
   .. note::
-    The point cloud must be in range :obj:`[-1, 1]`.
+    The point cloud must be strictly in range :obj:`[-1, 1]`. A good practice
+    is to normalize it into :obj:`[-0.99, 0.99]` or :obj:`[0.9, 0.9]` to retain
+    some margin.
   '''
 
   def __init__(self, depth: int, full_depth: int = 2, batch_size: int = 1,
@@ -63,14 +65,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
@@ -176,7 +178,7 @@ class Octree:
     for d in range(self.depth, self.full_depth, -1):
       # compute parent key, i.e. keys of layer (d -1)
       pkey = node_key >> 3
-      pkey, pidx, pcounts = torch.unique_consecutive(
+      pkey, pidx, _ = torch.unique_consecutive(
           pkey, return_inverse=True, return_counts=True)
 
       # augmented key
@@ -287,10 +289,27 @@ class Octree:
       update_neigh (bool): If True, construct the neighborhood indices.
     '''
 
+    # increase the octree depth if required
+    if depth > self.depth:
+      assert depth == self.depth + 1
+      self.depth = depth
+      self.keys.append(None)
+      self.children.append(None)
+      self.neighs.append(None)
+      self.features.append(None)
+      self.normals.append(None)
+      self.points.append(None)
+      zero = torch.zeros(1, dtype=torch.long)
+      self.nnum = torch.cat([self.nnum, zero])
+      self.nnum_nempty = torch.cat([self.nnum_nempty, zero])
+      zero = zero.view(1, 1)
+      self.batch_nnum = torch.cat([self.batch_nnum, zero], dim=0)
+      self.batch_nnum_nempty = torch.cat([self.batch_nnum_nempty, zero], dim=0)
+
     # 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 +345,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 +402,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

{ocnn-2.2.3 → ocnn-2.2.5/ocnn.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ocnn
-Version: 2.2.3
+Version: 2.2.5
 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

{ocnn-2.2.3 → ocnn-2.2.5}/setup.py

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