def __init__(self, inplanes, planes, stride=1, dilation=1, downsample=None, bn_momentum=0.1, leakiness=0, dimension=-1): super(BasicBlock, self).__init__() assert dimension > 0 self.conv1 = ME.MinkowskiConvolution(inplanes, planes, kernel_size=3, stride=stride, dilation=dilation, dimension=dimension) self.norm1 = ME.MinkowskiBatchNorm(planes, momentum=bn_momentum) self.conv2 = ME.MinkowskiConvolution(planes, planes, kernel_size=3, stride=1, dilation=dilation, dimension=dimension) self.norm2 = ME.MinkowskiBatchNorm(planes, momentum=bn_momentum) self.relu = ME.MinkowskiLeakyReLU(negative_slope=leakiness) self.downsample = downsample
def get_conv_block(self, in_channel, out_channel, kernel_size, stride): return nn.Sequential( ME.MinkowskiConvolution( in_channel, out_channel, kernel_size=kernel_size, stride=stride, dimension=self.D, ), ME.MinkowskiBatchNorm(out_channel), ME.MinkowskiLeakyReLU(), )
def get_mlp_block(self, in_channel, out_channel): return nn.Sequential( ME.MinkowskiLinear(in_channel, out_channel, bias=False), ME.MinkowskiBatchNorm(out_channel), ME.MinkowskiLeakyReLU(), )
def network_initialization(self, in_channels, out_channels, D): # Output of the first conv concated to conv6 self.inplanes = self.INIT_DIM self.conv0p1s1 = ME.MinkowskiConvolution( in_channels, self.inplanes, kernel_size=5, dimension=D) self.bn0 = ME.MinkowskiBatchNorm(self.inplanes) self.conv1p1s2 = ME.MinkowskiConvolution( self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn1 = ME.MinkowskiBatchNorm(self.inplanes) self.block1 = self._make_layer(self.BLOCK, self.PLANES[0], self.LAYERS[0], leakiness=self.leakiness) self.conv2p2s2 = ME.MinkowskiConvolution( self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn2 = ME.MinkowskiBatchNorm(self.inplanes) self.block2 = self._make_layer(self.BLOCK, self.PLANES[1], self.LAYERS[1], leakiness=self.leakiness) self.conv3p4s2 = ME.MinkowskiConvolution( self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn3 = ME.MinkowskiBatchNorm(self.inplanes) self.block3 = self._make_layer(self.BLOCK, self.PLANES[2], self.LAYERS[2], leakiness=self.leakiness) self.conv4p8s2 = ME.MinkowskiConvolution( self.inplanes, self.inplanes, kernel_size=2, stride=2, dimension=D) self.bn4 = ME.MinkowskiBatchNorm(self.inplanes) self.block4 = self._make_layer(self.BLOCK, self.PLANES[3], self.LAYERS[3], leakiness=self.leakiness) self.convtr4p16s2 = ME.MinkowskiConvolutionTranspose( self.inplanes, self.PLANES[4], kernel_size=2, stride=2, dimension=D) self.bntr4 = ME.MinkowskiBatchNorm(self.PLANES[4]) self.inplanes = self.PLANES[4] if self.skipconnections: self.inplanes += self.PLANES[2] * self.BLOCK.expansion self.block5 = self._make_layer(self.BLOCK, self.PLANES[4], self.LAYERS[4], leakiness=self.leakiness) self.convtr5p8s2 = ME.MinkowskiConvolutionTranspose( self.inplanes, self.PLANES[5], kernel_size=2, stride=2, dimension=D) self.bntr5 = ME.MinkowskiBatchNorm(self.PLANES[5]) self.inplanes = self.PLANES[5] #+ self.PLANES[1] * self.BLOCK.expansion if self.skipconnections: self.inplanes += self.PLANES[1] * self.BLOCK.expansion self.block6 = self._make_layer(self.BLOCK, self.PLANES[5], self.LAYERS[5], leakiness=self.leakiness) self.convtr6p4s2 = ME.MinkowskiConvolutionTranspose( self.inplanes, self.PLANES[6], kernel_size=2, stride=2, dimension=D) self.bntr6 = ME.MinkowskiBatchNorm(self.PLANES[6]) self.inplanes = self.PLANES[6] #+ self.PLANES[0] * self.BLOCK.expansion if self.skipconnections: self.inplanes += self.PLANES[0] * self.BLOCK.expansion self.block7 = self._make_layer(self.BLOCK, self.PLANES[6], self.LAYERS[6], leakiness=self.leakiness) self.convtr7p2s2 = ME.MinkowskiConvolutionTranspose( self.inplanes, self.PLANES[7], kernel_size=2, stride=2, dimension=D) self.bntr7 = ME.MinkowskiBatchNorm(self.PLANES[7]) self.inplanes = self.PLANES[7] #+ self.INIT_DIM if self.skipconnections: self.inplanes += self.INIT_DIM self.block8 = self._make_layer(self.BLOCK, self.PLANES[7], self.LAYERS[7], leakiness=self.leakiness) self.final = ME.MinkowskiConvolution( self.PLANES[7], out_channels, kernel_size=1, has_bias=True, dimension=D) self.relu = ME.MinkowskiLeakyReLU(negative_slope=self.leakiness) self.dropout = ME.MinkowskiDropout(self.dropout_p)