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)