ocnn 2.2.8__py3-none-any.whl → 2.3.0__py3-none-any.whl

ocnn/modules/resblocks.py

@@ -1,158 +1,158 @@
-# --------------------------------------------------------
-# Octree-based Sparse Convolutional Neural Networks
-# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Peng-Shuai Wang
-# --------------------------------------------------------
-
-import torch
-import torch.utils.checkpoint
-
-from ocnn.octree import Octree
-from ocnn.nn import OctreeMaxPool
-from ocnn.modules import (Conv1x1BnRelu, OctreeConvBnRelu, Conv1x1Bn,
-                          OctreeConvBn, OctreeConvGnRelu, Conv1x1Gn,
-                          OctreeConvGn,)
-
-
-class OctreeResBlock(torch.nn.Module):
-  r''' Octree-based ResNet block in a bottleneck style. The block is composed of
-  a series of :obj:`Conv1x1`, :obj:`Conv3x3`, and :obj:`Conv1x1`.
-
-  Args:
-    in_channels (int): Number of input channels.
-    out_channels (int): Number of output channels.
-    stride (int): The stride of the block (:obj:`1` or :obj:`2`).
-    bottleneck (int): The input and output channels of the :obj:`Conv3x3` is
-        equal to the input channel divided by :attr:`bottleneck`.
-    nempty (bool): If True, only performs the convolution on non-empty
-        octree nodes.
-  '''
-
-  def __init__(self, in_channels: int, out_channels: int, stride: int = 1,
-               bottleneck: int = 4, nempty: bool = False):
-    super().__init__()
-    self.in_channels = in_channels
-    self.out_channels = out_channels
-    self.bottleneck = bottleneck
-    self.stride = stride
-    channelb = int(out_channels / bottleneck)
-
-    if self.stride == 2:
-      self.max_pool = OctreeMaxPool(nempty)
-    self.conv1x1a = Conv1x1BnRelu(in_channels, channelb)
-    self.conv3x3 = OctreeConvBnRelu(channelb, channelb, nempty=nempty)
-    self.conv1x1b = Conv1x1Bn(channelb, out_channels)
-    if self.in_channels != self.out_channels:
-      self.conv1x1c = Conv1x1Bn(in_channels, out_channels)
-    self.relu = torch.nn.ReLU(inplace=True)
-
-  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
-    r''''''
-
-    if self.stride == 2:
-      data = self.max_pool(data, octree, depth)
-      depth = depth - 1
-    conv1 = self.conv1x1a(data)
-    conv2 = self.conv3x3(conv1, octree, depth)
-    conv3 = self.conv1x1b(conv2)
-    if self.in_channels != self.out_channels:
-      data = self.conv1x1c(data)
-    out = self.relu(conv3 + data)
-    return out
-
-
-class OctreeResBlock2(torch.nn.Module):
-  r''' Basic Octree-based ResNet block. The block is composed of
-  a series of :obj:`Conv3x3` and :obj:`Conv3x3`.
-
-  Refer to :class:`OctreeResBlock` for the details of arguments.
-  '''
-
-  def __init__(self, in_channels, out_channels, stride=1, bottleneck=1,
-               nempty=False):
-    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 = OctreeConvBnRelu(in_channels, channelb, nempty=nempty)
-    self.conv3x3b = OctreeConvBn(channelb, out_channels, nempty=nempty)
-    if self.in_channels != self.out_channels:
-      self.conv1x1 = Conv1x1Bn(in_channels, out_channels)
-    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)
-    out = self.relu(conv2 + data)
-    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.
-  '''
-
-  def __init__(self, in_channels, out_channels, resblk_num, bottleneck=4,
-               nempty=False, resblk=OctreeResBlock, use_checkpoint=False):
-    super().__init__()
-    self.resblk_num = resblk_num
-    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)])
-
-  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
-    r''''''
-
-    for i in range(self.resblk_num):
-      if self.use_checkpoint:
-        data = torch.utils.checkpoint.checkpoint(
-            self.resblks[i], data, octree, depth, use_reentrant=False)
-      else:
-        data = self.resblks[i](data, octree, depth)
-    return data
+# --------------------------------------------------------
+# Octree-based Sparse Convolutional Neural Networks
+# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Peng-Shuai Wang
+# --------------------------------------------------------
+
+import torch
+import torch.utils.checkpoint
+
+from ocnn.octree import Octree
+from ocnn.nn import OctreeMaxPool
+from ocnn.modules import (Conv1x1BnRelu, OctreeConvBnRelu, Conv1x1Bn,
+                          OctreeConvBn, OctreeConvGnRelu, Conv1x1Gn,
+                          OctreeConvGn,)
+
+
+class OctreeResBlock(torch.nn.Module):
+  r''' Octree-based ResNet block in a bottleneck style. The block is composed of
+  a series of :obj:`Conv1x1`, :obj:`Conv3x3`, and :obj:`Conv1x1`.
+
+  Args:
+    in_channels (int): Number of input channels.
+    out_channels (int): Number of output channels.
+    stride (int): The stride of the block (:obj:`1` or :obj:`2`).
+    bottleneck (int): The input and output channels of the :obj:`Conv3x3` is
+        equal to the input channel divided by :attr:`bottleneck`.
+    nempty (bool): If True, only performs the convolution on non-empty
+        octree nodes.
+  '''
+
+  def __init__(self, in_channels: int, out_channels: int, stride: int = 1,
+               bottleneck: int = 4, nempty: bool = False):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.bottleneck = bottleneck
+    self.stride = stride
+    channelb = int(out_channels / bottleneck)
+
+    if self.stride == 2:
+      self.max_pool = OctreeMaxPool(nempty)
+    self.conv1x1a = Conv1x1BnRelu(in_channels, channelb)
+    self.conv3x3 = OctreeConvBnRelu(channelb, channelb, nempty=nempty)
+    self.conv1x1b = Conv1x1Bn(channelb, out_channels)
+    if self.in_channels != self.out_channels:
+      self.conv1x1c = Conv1x1Bn(in_channels, out_channels)
+    self.relu = torch.nn.ReLU(inplace=True)
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+
+    if self.stride == 2:
+      data = self.max_pool(data, octree, depth)
+      depth = depth - 1
+    conv1 = self.conv1x1a(data)
+    conv2 = self.conv3x3(conv1, octree, depth)
+    conv3 = self.conv1x1b(conv2)
+    if self.in_channels != self.out_channels:
+      data = self.conv1x1c(data)
+    out = self.relu(conv3 + data)
+    return out
+
+
+class OctreeResBlock2(torch.nn.Module):
+  r''' Basic Octree-based ResNet block. The block is composed of
+  a series of :obj:`Conv3x3` and :obj:`Conv3x3`.
+
+  Refer to :class:`OctreeResBlock` for the details of arguments.
+  '''
+
+  def __init__(self, in_channels, out_channels, stride=1, bottleneck=1,
+               nempty=False):
+    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 = OctreeConvBnRelu(in_channels, channelb, nempty=nempty)
+    self.conv3x3b = OctreeConvBn(channelb, out_channels, nempty=nempty)
+    if self.in_channels != self.out_channels:
+      self.conv1x1 = Conv1x1Bn(in_channels, out_channels)
+    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)
+    out = self.relu(conv2 + data)
+    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.
+  '''
+
+  def __init__(self, in_channels, out_channels, resblk_num, bottleneck=4,
+               nempty=False, resblk=OctreeResBlock, use_checkpoint=False):
+    super().__init__()
+    self.resblk_num = resblk_num
+    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)])
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+
+    for i in range(self.resblk_num):
+      if self.use_checkpoint:
+        data = torch.utils.checkpoint.checkpoint(
+            self.resblks[i], data, octree, depth, use_reentrant=False)
+      else:
+        data = self.resblks[i](data, octree, depth)
+    return data

ocnn/nn/__init__.py

@@ -1,44 +1,45 @@
-# --------------------------------------------------------
-# Octree-based Sparse Convolutional Neural Networks
-# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Peng-Shuai Wang
-# --------------------------------------------------------
-
-from .octree2vox import octree2voxel, Octree2Voxel
-from .octree2col import octree2col, col2octree
-from .octree_pad import octree_pad, octree_depad
-from .octree_interp import (octree_nearest_pts, octree_linear_pts,
-                            OctreeInterp, OctreeUpsample)
-from .octree_pool import (octree_max_pool, OctreeMaxPool,
-                          octree_max_unpool, OctreeMaxUnpool,
-                          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, OctreeNorm)
-from .octree_drop import OctreeDropPath
-from .octree_align import search_value, octree_align
-
-
-__all__ = [
-    'octree2voxel',
-    'octree2col', 'col2octree',
-    'octree_pad', 'octree_depad',
-    'octree_nearest_pts', 'octree_linear_pts',
-    'octree_max_pool', 'octree_max_unpool',
-    'octree_global_pool', 'octree_avg_pool',
-    'Octree2Voxel',
-    'OctreeMaxPool', 'OctreeMaxUnpool',
-    'OctreeGlobalPool', 'OctreeAvgPool',
-    'OctreeConv', 'OctreeDeconv',
-    'OctreeGroupConv', 'OctreeDWConv',
-    'OctreeInterp', 'OctreeUpsample',
-    'OctreeInstanceNorm', 'OctreeBatchNorm', 'OctreeGroupNorm', 'OctreeNorm',
-    'OctreeDropPath',
-    'search_value', 'octree_align',
-]
-
-classes = __all__
+# --------------------------------------------------------
+# Octree-based Sparse Convolutional Neural Networks
+# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Peng-Shuai Wang
+# --------------------------------------------------------
+
+from .octree2vox import octree2voxel, Octree2Voxel
+from .octree2col import octree2col, col2octree
+from .octree_pad import octree_pad, octree_depad
+from .octree_interp import (octree_nearest_pts, octree_linear_pts,
+                            OctreeInterp, OctreeUpsample)
+from .octree_pool import (octree_max_pool, OctreeMaxPool,
+                          octree_max_unpool, OctreeMaxUnpool,
+                          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, OctreeNorm)
+from .octree_drop import OctreeDropPath
+from .octree_align import search_value, octree_align
+from .octree_conv_t import OctreeConvTriton, OctreeConvT, convert_conv_triton
+
+__all__ = [
+    'octree2voxel',
+    'octree2col', 'col2octree',
+    'octree_pad', 'octree_depad',
+    'octree_nearest_pts', 'octree_linear_pts',
+    'octree_max_pool', 'octree_max_unpool',
+    'octree_global_pool', 'octree_avg_pool',
+    'Octree2Voxel',
+    'OctreeMaxPool', 'OctreeMaxUnpool',
+    'OctreeGlobalPool', 'OctreeAvgPool',
+    'OctreeConv', 'OctreeDeconv',
+    'OctreeGroupConv', 'OctreeDWConv',
+    'OctreeInterp', 'OctreeUpsample',
+    'OctreeInstanceNorm', 'OctreeBatchNorm', 'OctreeGroupNorm', 'OctreeNorm',
+    'OctreeDropPath',
+    'search_value', 'octree_align',
+    'OctreeConvTriton', 'OctreeConvT', 'convert_conv_triton',
+]
+
+classes = __all__

ocnn/nn/kernels/__init__.py

@@ -0,0 +1,14 @@
+from .conv_fwd_implicit_gemm_splitk import conv_fwd_implicit_gemm_splitk
+from .conv_bwd_implicit_gemm_splitk import conv_bwd_implicit_gemm_splitk
+from .conv_bwd_implicit_gemm import conv_bwd_implicit_gemm
+from .conv_fwd_implicit_gemm import conv_fwd_implicit_gemm
+
+__all__ = [
+    'conv_fwd_implicit_gemm_splitk',
+    'conv_bwd_implicit_gemm_splitk',
+    'conv_bwd_implicit_gemm',
+    'conv_fwd_implicit_gemm',
+]
+
+from .autotuner import load_autotune_cache
+load_autotune_cache()