ocnn 2.2.1__py3-none-any.whl → 2.2.2__py3-none-any.whl

ocnn/models/ounet.py

@@ -1,94 +1,94 @@
-# --------------------------------------------------------
-# 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 ocnn
-import torch
-import torch.nn
-
-from ocnn.octree import Octree
-from ocnn.models.autoencoder import AutoEncoder
-
-
-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
-    self.proj = None  # remove this module used in AutoEncoder
-
-  def encoder(self, octree):
-    r''' The encoder network for extracting heirarchy features. 
-    '''
-
-    convs = dict()
-    depth, full_depth = self.depth, self.full_depth
-    data = self.get_input_feature(octree)
-    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)
-      if d > full_depth:
-        convs[d-1] = self.downsample[i](convs[d], octree, d)
-    return convs
-
-  def decoder(self, convs: dict, octree_in: Octree, octree_out: Octree,
-              update_octree: bool = False):
-    r''' The decoder network for decode the octree.
-    '''
-
-    logits = dict()
-    deconv = convs[self.full_depth]
-    depth, full_depth = self.depth, self.full_depth
-    for i, d in enumerate(range(full_depth, depth + 1)):
-      if d > full_depth:
-        deconv = self.upsample[i-1](deconv, octree_out, d-1)
-        skip = ocnn.nn.octree_align(convs[d], octree_in, octree_out, d)
-        deconv = deconv + skip  # output-guided skip connections
-      deconv = self.decoder_blks[i](deconv, octree_out, d)
-
-      # predict the splitting label
-      logit = self.predict[i](deconv)
-      logits[d] = logit
-
-      # update the octree according to predicted labels
-      if update_octree:
-        split = logit.argmax(1).int()
-        octree_out.octree_split(split, d)
-        if d < depth:
-          octree_out.octree_grow(d + 1)
-
-      # predict the signal
-      if d == depth:
-        signal = self.header(deconv)
-        signal = torch.tanh(signal)
-        signal = ocnn.nn.octree_depad(signal, octree_out, depth)
-        if update_octree:
-          octree_out.features[depth] = signal
-
-    return {'logits': logits, 'signal': signal, 'octree_out': octree_out}
-
-  def init_octree(self, octree_in: Octree):
-    r''' Initialize a full octree for decoding.
-    '''
-
-    device = octree_in.device
-    batch_size = octree_in.batch_size
-    octree = Octree(self.depth, self.full_depth, batch_size, device)
-    for d in range(self.full_depth+1):
-      octree.octree_grow_full(depth=d)
-    return octree
-
-  def forward(self, octree_in, octree_out=None, update_octree: bool = False):
-    r''''''
-
-    if octree_out is None:
-      update_octree = True
-      octree_out = self.init_octree(octree_in)
-    convs = self.encoder(octree_in)
-    out = self.decoder(convs, octree_in, octree_out, update_octree)
-    return out
+# --------------------------------------------------------
+# 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 ocnn
+import torch
+import torch.nn
+
+from ocnn.octree import Octree
+from ocnn.models.autoencoder import AutoEncoder
+
+
+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
+    self.proj = None  # remove this module used in AutoEncoder
+
+  def encoder(self, octree):
+    r''' The encoder network for extracting heirarchy features. 
+    '''
+
+    convs = dict()
+    depth, full_depth = self.depth, self.full_depth
+    data = self.get_input_feature(octree)
+    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)
+      if d > full_depth:
+        convs[d-1] = self.downsample[i](convs[d], octree, d)
+    return convs
+
+  def decoder(self, convs: dict, octree_in: Octree, octree_out: Octree,
+              update_octree: bool = False):
+    r''' The decoder network for decode the octree.
+    '''
+
+    logits = dict()
+    deconv = convs[self.full_depth]
+    depth, full_depth = self.depth, self.full_depth
+    for i, d in enumerate(range(full_depth, depth + 1)):
+      if d > full_depth:
+        deconv = self.upsample[i-1](deconv, octree_out, d-1)
+        skip = ocnn.nn.octree_align(convs[d], octree_in, octree_out, d)
+        deconv = deconv + skip  # output-guided skip connections
+      deconv = self.decoder_blks[i](deconv, octree_out, d)
+
+      # predict the splitting label
+      logit = self.predict[i](deconv)
+      logits[d] = logit
+
+      # update the octree according to predicted labels
+      if update_octree:
+        split = logit.argmax(1).int()
+        octree_out.octree_split(split, d)
+        if d < depth:
+          octree_out.octree_grow(d + 1)
+
+      # predict the signal
+      if d == depth:
+        signal = self.header(deconv)
+        signal = torch.tanh(signal)
+        signal = ocnn.nn.octree_depad(signal, octree_out, depth)
+        if update_octree:
+          octree_out.features[depth] = signal
+
+    return {'logits': logits, 'signal': signal, 'octree_out': octree_out}
+
+  def init_octree(self, octree_in: Octree):
+    r''' Initialize a full octree for decoding.
+    '''
+
+    device = octree_in.device
+    batch_size = octree_in.batch_size
+    octree = Octree(self.depth, self.full_depth, batch_size, device)
+    for d in range(self.full_depth+1):
+      octree.octree_grow_full(depth=d)
+    return octree
+
+  def forward(self, octree_in, octree_out=None, update_octree: bool = False):
+    r''''''
+
+    if octree_out is None:
+      update_octree = True
+      octree_out = self.init_octree(octree_in)
+    convs = self.encoder(octree_in)
+    out = self.decoder(convs, octree_in, octree_out, update_octree)
+    return out

ocnn/models/resnet.py

@@ -1,53 +1,53 @@
-# --------------------------------------------------------
-# 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 ocnn
-from ocnn.octree import Octree
-
-
-class ResNet(torch.nn.Module):
-  r''' Octree-based ResNet for classification.
-  '''
-
-  def __init__(self, in_channels: int, out_channels: int, resblock_num: int,
-               stages: int, nempty: bool = False):
-    super().__init__()
-    self.in_channels = in_channels
-    self.out_channels = out_channels
-    self.resblk_num = resblock_num
-    self.stages = stages
-    self.nempty = nempty
-    channels = [2 ** max(i+9-stages, 2) for i in range(stages)]
-
-    self.conv1 = ocnn.modules.OctreeConvBnRelu(
-        in_channels, channels[0], nempty=nempty)
-    self.pool1 = ocnn.nn.OctreeMaxPool(nempty)
-    self.resblocks = torch.nn.ModuleList([ocnn.modules.OctreeResBlocks(
-        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)])
-    self.global_pool = ocnn.nn.OctreeGlobalPool(nempty)
-    # self.header = torch.nn.Linear(channels[-1], out_channels, bias=True)
-    self.header = torch.nn.Sequential(
-        ocnn.modules.FcBnRelu(channels[-1], 512),
-        torch.nn.Dropout(p=0.5),
-        torch.nn.Linear(512, out_channels))
-
-  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
-    r''''''
-
-    data = self.conv1(data, octree, depth)
-    data = self.pool1(data, octree, depth)
-    for i in range(self.stages-1):
-      d = depth - i - 1
-      data = self.resblocks[i](data, octree, d)
-      data = self.pools[i](data, octree, d)
-    data = self.global_pool(data, octree, depth-self.stages)
-    data = self.header(data)
-    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 ocnn
+from ocnn.octree import Octree
+
+
+class ResNet(torch.nn.Module):
+  r''' Octree-based ResNet for classification.
+  '''
+
+  def __init__(self, in_channels: int, out_channels: int, resblock_num: int,
+               stages: int, nempty: bool = False):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.resblk_num = resblock_num
+    self.stages = stages
+    self.nempty = nempty
+    channels = [2 ** max(i+9-stages, 2) for i in range(stages)]
+
+    self.conv1 = ocnn.modules.OctreeConvBnRelu(
+        in_channels, channels[0], nempty=nempty)
+    self.pool1 = ocnn.nn.OctreeMaxPool(nempty)
+    self.resblocks = torch.nn.ModuleList([ocnn.modules.OctreeResBlocks(
+        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)])
+    self.global_pool = ocnn.nn.OctreeGlobalPool(nempty)
+    # self.header = torch.nn.Linear(channels[-1], out_channels, bias=True)
+    self.header = torch.nn.Sequential(
+        ocnn.modules.FcBnRelu(channels[-1], 512),
+        torch.nn.Dropout(p=0.5),
+        torch.nn.Linear(512, out_channels))
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+
+    data = self.conv1(data, octree, depth)
+    data = self.pool1(data, octree, depth)
+    for i in range(self.stages-1):
+      d = depth - i - 1
+      data = self.resblocks[i](data, octree, d)
+      data = self.pools[i](data, octree, d)
+    data = self.global_pool(data, octree, depth-self.stages)
+    data = self.header(data)
+    return data

ocnn/models/segnet.py

@@ -1,72 +1,72 @@
-# --------------------------------------------------------
-# 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 ocnn
-from ocnn.octree import Octree
-
-
-class SegNet(torch.nn.Module):
-  r''' Octree-based SegNet for segmentation.
-  '''
-
-  def __init__(self, in_channels: int, out_channels: int, stages: int,
-               interp: str = 'linear', nempty: bool = False, **kwargs):
-    super().__init__()
-    self.in_channels = in_channels
-    self.out_channels = out_channels
-    self.stages = stages
-    self.nempty = nempty
-    return_indices = True
-
-    channels_stages = [2 ** max(i+8-stages, 2) for i in range(stages)]
-    channels = [in_channels] + channels_stages
-    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, return_indices) for i in range(stages)])
-
-    self.bottleneck = ocnn.modules.OctreeConvBnRelu(channels[-1], channels[-1])
-
-    channels = channels_stages[::-1] + [channels_stages[0]]
-    self.deconvs = torch.nn.ModuleList(
-        [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)])
-
-    self.octree_interp = ocnn.nn.OctreeInterp(interp, nempty)
-    self.header = torch.nn.Sequential(
-        ocnn.modules.Conv1x1BnRelu(channels[-1], 64),
-        ocnn.modules.Conv1x1(64, out_channels, use_bias=True))
-
-  def forward(self, data: torch.Tensor, octree: Octree, depth: int,
-              query_pts: torch.Tensor):
-    r''''''
-
-    # encoder
-    indices = dict()
-    for i in range(self.stages):
-      d = depth - i
-      data = self.convs[i](data, octree, d)
-      data, indices[d] = self.pools[i](data, octree, d)
-
-    # bottleneck
-    data = self.bottleneck(data, octree, depth-self.stages)
-
-    # decoder
-    for i in range(self.stages):
-      d = depth - self.stages + i
-      data = self.unpools[i](data, indices[d + 1], octree, d)
-      data = self.deconvs[i](data, octree, d + 1)
-
-    # header
-    feature = self.octree_interp(data, octree, depth, query_pts)
-    logits = self.header(feature)
-
-    return logits
+# --------------------------------------------------------
+# 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 ocnn
+from ocnn.octree import Octree
+
+
+class SegNet(torch.nn.Module):
+  r''' Octree-based SegNet for segmentation.
+  '''
+
+  def __init__(self, in_channels: int, out_channels: int, stages: int,
+               interp: str = 'linear', nempty: bool = False, **kwargs):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.stages = stages
+    self.nempty = nempty
+    return_indices = True
+
+    channels_stages = [2 ** max(i+8-stages, 2) for i in range(stages)]
+    channels = [in_channels] + channels_stages
+    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, return_indices) for i in range(stages)])
+
+    self.bottleneck = ocnn.modules.OctreeConvBnRelu(channels[-1], channels[-1])
+
+    channels = channels_stages[::-1] + [channels_stages[0]]
+    self.deconvs = torch.nn.ModuleList(
+        [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)])
+
+    self.octree_interp = ocnn.nn.OctreeInterp(interp, nempty)
+    self.header = torch.nn.Sequential(
+        ocnn.modules.Conv1x1BnRelu(channels[-1], 64),
+        ocnn.modules.Conv1x1(64, out_channels, use_bias=True))
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int,
+              query_pts: torch.Tensor):
+    r''''''
+
+    # encoder
+    indices = dict()
+    for i in range(self.stages):
+      d = depth - i
+      data = self.convs[i](data, octree, d)
+      data, indices[d] = self.pools[i](data, octree, d)
+
+    # bottleneck
+    data = self.bottleneck(data, octree, depth-self.stages)
+
+    # decoder
+    for i in range(self.stages):
+      d = depth - self.stages + i
+      data = self.unpools[i](data, indices[d + 1], octree, d)
+      data = self.deconvs[i](data, octree, d + 1)
+
+    # header
+    feature = self.octree_interp(data, octree, depth, query_pts)
+    logits = self.header(feature)
+
+    return logits

ocnn/models/unet.py

@@ -1,105 +1,105 @@
-# --------------------------------------------------------
-# 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.nn
-from typing import Dict
-
-import ocnn
-from ocnn.octree import Octree
-
-
-class UNet(torch.nn.Module):
-  r''' Octree-based UNet for segmentation.
-  '''
-
-  def __init__(self, in_channels: int, out_channels: int, interp: str = 'linear',
-               nempty: bool = False, **kwargs):
-    super(UNet, self).__init__()
-    self.in_channels = in_channels
-    self.out_channels = out_channels
-    self.nempty = nempty
-    self.config_network()
-    self.encoder_stages = len(self.encoder_blocks)
-    self.decoder_stages = len(self.decoder_blocks)
-
-    # encoder
-    self.conv1 = ocnn.modules.OctreeConvBnRelu(
-        in_channels, self.encoder_channel[0], nempty=nempty)
-    self.downsample = torch.nn.ModuleList([ocnn.modules.OctreeConvBnRelu(
-        self.encoder_channel[i], self.encoder_channel[i+1], kernel_size=[2],
-        stride=2, nempty=nempty) for i in range(self.encoder_stages)])
-    self.encoder = torch.nn.ModuleList([ocnn.modules.OctreeResBlocks(
-        self.encoder_channel[i+1], self.encoder_channel[i + 1],
-        self.encoder_blocks[i], self.bottleneck, nempty, self.resblk)
-        for i in range(self.encoder_stages)])
-
-    # decoder
-    channel = [self.decoder_channel[i+1] + self.encoder_channel[-i-2]
-               for i in range(self.decoder_stages)]
-    self.upsample = torch.nn.ModuleList([ocnn.modules.OctreeDeconvBnRelu(
-        self.decoder_channel[i], self.decoder_channel[i+1], kernel_size=[2],
-        stride=2, nempty=nempty) for i in range(self.decoder_stages)])
-    self.decoder = torch.nn.ModuleList([ocnn.modules.OctreeResBlocks(
-        channel[i], self.decoder_channel[i+1],
-        self.decoder_blocks[i], self.bottleneck, nempty, self.resblk)
-        for i in range(self.decoder_stages)])
-
-    # header
-    # channel = self.decoder_channel[self.decoder_stages]
-    self.octree_interp = ocnn.nn.OctreeInterp(interp, nempty)
-    self.header = torch.nn.Sequential(
-        ocnn.modules.Conv1x1BnRelu(self.decoder_channel[-1], self.head_channel),
-        ocnn.modules.Conv1x1(self.head_channel, self.out_channels, use_bias=True))
-
-  def config_network(self):
-    r''' Configure the network channels and Resblock numbers.
-    '''
-
-    self.encoder_channel = [32, 32, 64, 128, 256]
-    self.decoder_channel = [256, 256, 128, 96, 96]
-    self.encoder_blocks = [2, 3, 4, 6]
-    self.decoder_blocks = [2, 2, 2, 2]
-    self.head_channel = 64
-    self.bottleneck = 1
-    self.resblk = ocnn.modules.OctreeResBlock2
-
-  def unet_encoder(self, data: torch.Tensor, octree: Octree, depth: int):
-    r''' The encoder of the U-Net. 
-    '''
-
-    convd = dict()
-    convd[depth] = self.conv1(data, octree, depth)
-    for i in range(self.encoder_stages):
-      d = depth - i
-      conv = self.downsample[i](convd[d], octree, d)
-      convd[d-1] = self.encoder[i](conv, octree, d-1)
-    return convd
-
-  def unet_decoder(self, convd: Dict[int, torch.Tensor], octree: Octree, depth: int):
-    r''' The decoder of the U-Net. 
-    '''
-
-    deconv = convd[depth]
-    for i in range(self.decoder_stages):
-      d = depth + i
-      deconv = self.upsample[i](deconv, octree, d)
-      deconv = torch.cat([convd[d+1], deconv], dim=1)  # skip connections
-      deconv = self.decoder[i](deconv, octree, d+1)
-    return deconv
-
-  def forward(self, data: torch.Tensor, octree: Octree, depth: int,
-              query_pts: torch.Tensor):
-    r''''''
-
-    convd = self.unet_encoder(data, octree, depth)
-    deconv = self.unet_decoder(convd, octree, depth - self.encoder_stages)
-
-    interp_depth = depth - self.encoder_stages + self.decoder_stages
-    feature = self.octree_interp(deconv, octree, interp_depth, query_pts)
-    logits = self.header(feature)
-    return logits
+# --------------------------------------------------------
+# 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.nn
+from typing import Dict
+
+import ocnn
+from ocnn.octree import Octree
+
+
+class UNet(torch.nn.Module):
+  r''' Octree-based UNet for segmentation.
+  '''
+
+  def __init__(self, in_channels: int, out_channels: int, interp: str = 'linear',
+               nempty: bool = False, **kwargs):
+    super(UNet, self).__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.nempty = nempty
+    self.config_network()
+    self.encoder_stages = len(self.encoder_blocks)
+    self.decoder_stages = len(self.decoder_blocks)
+
+    # encoder
+    self.conv1 = ocnn.modules.OctreeConvBnRelu(
+        in_channels, self.encoder_channel[0], nempty=nempty)
+    self.downsample = torch.nn.ModuleList([ocnn.modules.OctreeConvBnRelu(
+        self.encoder_channel[i], self.encoder_channel[i+1], kernel_size=[2],
+        stride=2, nempty=nempty) for i in range(self.encoder_stages)])
+    self.encoder = torch.nn.ModuleList([ocnn.modules.OctreeResBlocks(
+        self.encoder_channel[i+1], self.encoder_channel[i + 1],
+        self.encoder_blocks[i], self.bottleneck, nempty, self.resblk)
+        for i in range(self.encoder_stages)])
+
+    # decoder
+    channel = [self.decoder_channel[i+1] + self.encoder_channel[-i-2]
+               for i in range(self.decoder_stages)]
+    self.upsample = torch.nn.ModuleList([ocnn.modules.OctreeDeconvBnRelu(
+        self.decoder_channel[i], self.decoder_channel[i+1], kernel_size=[2],
+        stride=2, nempty=nempty) for i in range(self.decoder_stages)])
+    self.decoder = torch.nn.ModuleList([ocnn.modules.OctreeResBlocks(
+        channel[i], self.decoder_channel[i+1],
+        self.decoder_blocks[i], self.bottleneck, nempty, self.resblk)
+        for i in range(self.decoder_stages)])
+
+    # header
+    # channel = self.decoder_channel[self.decoder_stages]
+    self.octree_interp = ocnn.nn.OctreeInterp(interp, nempty)
+    self.header = torch.nn.Sequential(
+        ocnn.modules.Conv1x1BnRelu(self.decoder_channel[-1], self.head_channel),
+        ocnn.modules.Conv1x1(self.head_channel, self.out_channels, use_bias=True))
+
+  def config_network(self):
+    r''' Configure the network channels and Resblock numbers.
+    '''
+
+    self.encoder_channel = [32, 32, 64, 128, 256]
+    self.decoder_channel = [256, 256, 128, 96, 96]
+    self.encoder_blocks = [2, 3, 4, 6]
+    self.decoder_blocks = [2, 2, 2, 2]
+    self.head_channel = 64
+    self.bottleneck = 1
+    self.resblk = ocnn.modules.OctreeResBlock2
+
+  def unet_encoder(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''' The encoder of the U-Net. 
+    '''
+
+    convd = dict()
+    convd[depth] = self.conv1(data, octree, depth)
+    for i in range(self.encoder_stages):
+      d = depth - i
+      conv = self.downsample[i](convd[d], octree, d)
+      convd[d-1] = self.encoder[i](conv, octree, d-1)
+    return convd
+
+  def unet_decoder(self, convd: Dict[int, torch.Tensor], octree: Octree, depth: int):
+    r''' The decoder of the U-Net. 
+    '''
+
+    deconv = convd[depth]
+    for i in range(self.decoder_stages):
+      d = depth + i
+      deconv = self.upsample[i](deconv, octree, d)
+      deconv = torch.cat([convd[d+1], deconv], dim=1)  # skip connections
+      deconv = self.decoder[i](deconv, octree, d+1)
+    return deconv
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int,
+              query_pts: torch.Tensor):
+    r''''''
+
+    convd = self.unet_encoder(data, octree, depth)
+    deconv = self.unet_decoder(convd, octree, depth - self.encoder_stages)
+
+    interp_depth = depth - self.encoder_stages + self.decoder_stages
+    feature = self.octree_interp(deconv, octree, interp_depth, query_pts)
+    logits = self.header(feature)
+    return logits