def __init__(self, depth, channel_in, nout):
super(ourNet, self).__init__()
self.depth = depth
self.channel_in = channel_in
channels = [2**max(9 - i, 2) for i in range(depth + 1)]
channels.append(channel_in)
self.conv_block, self.pool_block = conv_block(channel_in, depth,
channel_in)
self.octree2voxel = ocnn.FullOctree2Voxel(2)
self.drop = torch.nn.Dropout(p=0.4)
self.drop2 = torch.nn.Dropout(p=0.6)
self.flat = torch.nn.Flatten(start_dim=1)
self.fc = torch.nn.utils.weight_norm(
torch.nn.Linear(channels[3] * 64, channels[2], bias=False))
self.bn = torch.nn.BatchNorm1d(channels[2])
self.relu = torch.nn.ReLU(inplace=True)
self.fcRd = torch.nn.utils.weight_norm(
torch.nn.Linear(channels[2], channels[2], bias=False))
self.fc2 = torch.nn.utils.weight_norm(
torch.nn.Linear(channels[2], nout))
self.mlp1 = MLP([channels[2], 32, 64, 128])
self.mlp2 = MLP([128, 128])
# concatenate
self.mlp3 = MLP([256, 128, channels[2]])
def __init__(self, depth, channel_in, nout):
super(LeNet, self).__init__()
self.depth, self.channel_in = depth, channel_in
channels = [2 ** max(9 - i, 2) for i in range(depth + 1)]
channels.append(channel_in)
octree_conv, octree_pool = ocnn.OctreeConvBnRelu, ocnn.OctreeMaxPool
self.convs = torch.nn.ModuleList(
[octree_conv(d, channels[d + 1], channels[d]) for d in range(depth, 2, -1)])
self.pools = torch.nn.ModuleList(
[octree_pool(d) for d in range(depth, 2, -1)])
self.octree2voxel = ocnn.FullOctree2Voxel(2)
self.header = torch.nn.Sequential(
torch.nn.Dropout(p=0.5), # drop1
ocnn.FcBnRelu(channels[3] * 64, channels[2]), # fc1
torch.nn.Dropout(p=0.5), # drop2
torch.nn.Linear(channels[2], nout)) # fc2
def __init__(self, depth): super(FullOctreeGlobalPool, self).__init__(depth) self.octree2voxel = ocnn.FullOctree2Voxel(depth)
def __init__(self,
observation_space: gym.spaces.Box,
depth: int = 5,
full_depth: int = 2,
channels_in: int = 4,
channel_multiplier: int = 16,
full_depth_conv1d: bool = False,
full_depth_channels: int = 8,
features_dim: int = 128,
aux_obs_dim: int = 0,
fast_conv: bool = True,
batch_normalization: bool = True,
bn_eps: float = 0.00001,
bn_momentum: float = 0.01,
verbose: bool = False):
self._depth = depth
self._channels_in = channels_in
self._aux_obs_dim = aux_obs_dim
self._verbose = verbose
# Keyword arguments used for layers that might contain BatchNorm layers
bn_kwargs = {}
if batch_normalization:
bn_kwargs.update({'bn_eps': bn_eps, 'bn_momentum': bn_momentum})
# Determine number of stacked octrees based on observation space shape
self.n_stacks = observation_space.shape[0]
# Chain up parent constructor
super(OctreeCnnFeaturesExtractor,
self).__init__(observation_space,
self.n_stacks * (features_dim + aux_obs_dim))
# Channels ordered as [channels_in, depth, depth-1, ..., full_depth]
# I.e [channels_in, channel_multiplier*1, channel_multiplier*2, channel_multiplier*4, channel_multiplier*8,...]
channels = [
channel_multiplier * (2**i) for i in range(depth - full_depth)
]
channels.insert(0, channels_in)
# Create all Octree convolution and pooling layers in depth-descending order [depth, depth-1, ..., full_depth]
# Input to the first conv layer is the input Octree at depth=depth
# Output from the last pool layer is feature map at depth=full_depth
if fast_conv:
if batch_normalization:
OctreeConv = OctreeConvFastBnRelu
else:
OctreeConv = OctreeConvFastRelu
else:
if batch_normalization:
OctreeConv = OctreeConvBnRelu
else:
OctreeConv = OctreeConvRelu
OctreePool = ocnn.OctreeMaxPool
self.convs = torch.nn.ModuleList([
OctreeConv(depth - i, channels[i], channels[i + 1], **bn_kwargs)
for i in range(depth - full_depth)
])
self.pools = torch.nn.ModuleList(
[OctreePool(depth - i) for i in range(depth - full_depth)])
# Last convolution at depth=full_depth, which is not follewed by pooling
# This layer is used to significantly reduce the channels, decresing number of parameters required in the next linear layer(s)
self._full_depth_conv1d = full_depth_conv1d
if self._full_depth_conv1d:
# Use 1D convolution (Conv1D instead of linear is used here to preserve spatial information)
if batch_normalization:
OctreeConv1D = OctreeConv1x1BnRelu
else:
OctreeConv1D = OctreeConv1x1Relu
self.full_depth_conv = OctreeConv1D(channels[-1],
full_depth_channels,
**bn_kwargs)
else:
# Use 3D convolution (same as previous modules)
self.full_depth_conv = OctreeConv(full_depth, channels[-1],
full_depth_channels, **bn_kwargs)
# Layer that converts octree at depth=full_depth into a full voxel grid (zero padding) such that it has a fixed size
self.octree2voxel = ocnn.FullOctree2Voxel(full_depth)
full_depth_voxel_count = 2**(3 * full_depth)
# Layer that flattens the voxel grid into a feature vector
self.flatten = torch.nn.Flatten()
flatten_dim = full_depth_channels * full_depth_voxel_count
# Last linear layer of the extractor, applied to all (flattened) voxels at full depth
self.linear = LinearRelu(flatten_dim, features_dim)
# One linear layer for auxiliary observations
if self._aux_obs_dim != 0:
self.aux_obs_linear = LinearRelu(self._aux_obs_dim,
self._aux_obs_dim)
number_of_learnable_parameters = sum(p.numel()
for p in self.parameters()
if p.requires_grad)
print("Initialised OctreeCnnFeaturesExtractor with "
f"{number_of_learnable_parameters} parameters")
if verbose:
print(self)