def conv_block(input_channel, depth, channel_in):
conv_bn_re = []
conv_bn_re.append(
ocnn.OctreeConvBnRelu(depth, input_channel, 2**(9 - depth)))
pool = []
pool.append(ocnn.OctreeMaxPool(depth))
for i in range(depth - 1, 2, -1):
conv_bn_re.append(ocnn.OctreeConvBnRelu(i, 2**(9 - i - 1), 2**(9 - i)))
pool.append(ocnn.OctreeMaxPool(i))
return ModuleList(conv_bn_re), ModuleList(pool)
def test_forward_and_backward_max_pool(self):
depth, channel, height = 5, 2, 16
octree = ocnn.octree_batch(
ocnn.octree_samples(['octree_1', 'octree_1']))
data = np.array([[1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8],
[8.1, 7.2, 6.3, 5.4, 4.5, 3.6, 2.7, 1.8]],
dtype=np.float32)
data = np.concatenate([data, data], axis=1)
data = np.reshape(data, (1, channel, height, 1))
out_gt = np.array([[8.8, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
[8.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]],
dtype=np.float32)
out_gt = np.concatenate([out_gt, out_gt], axis=1)
out_gt = np.reshape(out_gt, (1, channel, height, 1))
grad_gt = np.array([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.1],
[8.1, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]],
dtype=np.float32)
grad_gt = np.concatenate([grad_gt, grad_gt], axis=1)
grad_gt = np.reshape(grad_gt, (1, channel, height, 1))
mask_gt = np.array([[[[7], [15]], [[0], [8]]]], dtype=np.int32)
# forward
octree = octree.to('cuda')
data_in = torch.from_numpy(data).to('cuda').requires_grad_()
outputs, mask_out = ocnn.OctreeMaxPool(depth)(data_in, octree)
# backward
pesudo_grad = torch.from_numpy(data).to('cuda')
outputs.backward(pesudo_grad)
# test
self.assertTrue(
np.array_equal(mask_out.cpu().detach().numpy(), mask_gt))
self.assertTrue(np.array_equal(outputs.cpu().detach().numpy(), out_gt))
self.assertTrue(np.array_equal(data_in.grad.cpu().numpy(), grad_gt))
def __init__(self,
depth,
channel_in,
channel_out,
stride=1,
bottleneck=1,
nempty=False):
super().__init__()
self.channel_in = channel_in
self.channel_out = channel_out
self.stride = stride
self.depth = depth
channelb = int(channel_out / bottleneck)
if self.stride == 2:
self.maxpool = ocnn.OctreeMaxPool(self.depth)
self.depth = self.depth - 1
self.conv3x3a = OctreeConvBnRelu(self.depth,
channel_in,
channelb,
nempty=nempty)
self.conv3x3b = OctreeConvBn(self.depth,
channelb,
channel_out,
nempty=nempty)
if self.channel_in != self.channel_out:
self.conv1x1 = OctreeConv1x1Bn(channel_in, channel_out)
self.relu = torch.nn.ReLU(inplace=True)
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, depth, channel_in, channel_out, stride=1):
super(OctreeResBlock2, self).__init__()
self.channel_in = channel_in
self.channel_out = channel_out
self.stride = stride
self.depth = depth
if self.stride == 2:
self.maxpool = ocnn.OctreeMaxPool(self.depth)
self.depth = self.depth - 1
self.conv3x3a = OctreeConvBnRelu(self.depth, channel_in, channel_out)
self.conv3x3b = OctreeConvBn(self.depth, channel_out, channel_out)
if self.channel_in != self.channel_out:
self.conv1x1 = OctreeConv1x1Bn(channel_in, channel_out)
self.relu = torch.nn.ReLU(inplace=True)
def __init__(self, depth, channel_in, channel_out, stride=1, bottleneck=4):
super(OctreeResBlock, self).__init__()
self.channel_in = channel_in
self.channel_out = channel_out
self.stride = stride
channelb = int(channel_out / bottleneck)
self.depth = depth
if self.stride == 2:
self.maxpool = ocnn.OctreeMaxPool(self.depth)
self.depth = self.depth - 1
self.conv1x1a = OctreeConv1x1BnRelu(channel_in, channelb)
self.conv3x3 = OctreeConvBnRelu(self.depth, channelb, channelb)
self.conv1x1b = OctreeConv1x1Bn(channelb, channel_out)
if self.channel_in != self.channel_out:
self.conv1x1c = OctreeConv1x1Bn(channel_in, channel_out)
self.relu = torch.nn.ReLU(inplace=True)
def __init__(self, depth, channel_in, nout, resblk_num):
super(ResNet, self).__init__()
self.depth, self.channel_in = depth, channel_in
channels = [2**max(11 - i, 2) for i in range(depth + 1)]
channels.append(channels[depth])
self.conv1 = ocnn.OctreeConvBnRelu(depth, channel_in, channels[depth])
self.resblocks = torch.nn.ModuleList([
ocnn.OctreeResBlocks(d, channels[d + 1], channels[d], resblk_num)
for d in range(depth, 2, -1)
])
self.pools = torch.nn.ModuleList(
[ocnn.OctreeMaxPool(d) for d in range(depth, 2, -1)])
self.header = torch.nn.Sequential(
ocnn.FullOctreeGlobalPool(depth=2), # global pool
# torch.nn.Dropout(p=0.5), # drop
torch.nn.Linear(channels[3], nout)) # fc