def __init__(self, depth, channel_in, nout, interp='linear'):
super(SegNet, self).__init__()
self.depth, self.channel_in = depth, channel_in
channels = [2**max(10 - i, 2) for i in range(depth + 1)]
channels.append(channel_in)
channels[2] = channels[3]
self.channels = channels
self.convs = torch.nn.ModuleList([
ocnn.OctreeConvBnRelu(d, channels[d + 1], channels[d])
for d in range(depth, 2, -1)
])
self.pools = torch.nn.ModuleList([
ocnn.OctreeMaxPool(d, return_indices=True)
for d in range(depth, 2, -1)
])
self.deconvs = torch.nn.ModuleList([
ocnn.OctreeConvBnRelu(d, channels[d], channels[d + 1])
for d in range(2, depth)
])
self.unpools = torch.nn.ModuleList(
[ocnn.OctreeMaxUnpool(d) for d in range(2, depth)])
self.deconv = ocnn.OctreeConvBnRelu(depth, channels[depth],
channels[depth])
self.octree_interp = ocnn.OctreeInterp(self.depth,
interp,
nempty=False)
self.header = torch.nn.Sequential(
ocnn.OctreeConv1x1BnRelu(channels[depth], 64), # fc1
ocnn.OctreeConv1x1(64, nout, use_bias=True)) # fc2
def __init__(self,
channel_in,
channel_out,
use_bias=True,
bn_eps=0.00001,
bn_momentum=0.1):
super(OctreeConv1x1BnRelu, self).__init__()
self.conv1x1 = ocnn.OctreeConv1x1(channel_in, channel_out, use_bias)
self.bn = torch.nn.BatchNorm2d(channel_out, bn_eps, bn_momentum)
self.relu = torch.nn.ReLU(inplace=True)
def __init__(self, channel_in, channel_out, use_bias=True):
super(OctreeConv1x1Relu, self).__init__()
self.conv1x1 = ocnn.OctreeConv1x1(channel_in, channel_out, use_bias)
self.relu = torch.nn.ReLU(inplace=True)
def make_predict_module(self, channel_in, channel_out=2, num_hidden=64):
return torch.nn.Sequential(
ocnn.OctreeConv1x1BnRelu(channel_in, num_hidden),
ocnn.OctreeConv1x1(num_hidden, channel_out, use_bias=True))