def network_initialization(self, in_channels, out_channels, D):
self.inplanes = self.INIT_DIM
self.conv1 = ME.MinkowskiConvolution(
in_channels, self.inplanes, kernel_size=5, stride=1, dimension=D)
self.bn1 = ME.MinkowskiBatchNorm(self.inplanes)
self.relu = ME.MinkowskiReLU(inplace=True)
self.pool = ME.MinkowskiSumPooling(kernel_size=2, stride=2, dimension=D)
self.layer1 = self._make_layer(
self.BLOCK, self.PLANES[0], self.LAYERS[0], stride=2)
self.layer2 = self._make_layer(
self.BLOCK, self.PLANES[1], self.LAYERS[1], stride=2)
self.layer3 = self._make_layer(
self.BLOCK, self.PLANES[2], self.LAYERS[2], stride=2)
self.layer4 = self._make_layer(
self.BLOCK, self.PLANES[3], self.LAYERS[3], stride=2)
self.glob_avg = ME.MinkowskiGlobalPooling(dimension=D)
self.classification_block = nn.Sequential(
ME.MinkowskiLinear(self.inplanes, self.inplanes, bias=False),
ME.MinkowskiBatchNorm(self.inplanes), ME.MinkowskiReLU(),
ME.MinkowskiLinear(self.inplanes, self.inplanes, bias=False),
ME.MinkowskiBatchNorm(self.inplanes))
self.final = ME.MinkowskiLinear(self.inplanes, out_channels, bias=True)
def __init__(self, channel, reduction=16, D=-1):
# Global coords does not require coords_key
super(SELayer, self).__init__()
self.fc = nn.Sequential(
ME.MinkowskiLinear(channel, channel // reduction), ME.MinkowskiReLU(inplace=True),
ME.MinkowskiLinear(channel // reduction, channel), ME.MinkowskiSigmoid())
self.pooling = ME.MinkowskiGlobalPooling(dimension=D)
self.broadcast_mul = ME.MinkowskiBroadcastMultiplication(dimension=D)
def __init__(self,
dims,
use_bn=False,
use_relu=False,
use_dropout=False,
use_bias=True):
super().__init__()
layers = []
last_dim = dims[0]
counter = 1
for dim in dims[1:]:
layers.append(ME.MinkowskiLinear(last_dim, dim, bias=use_bias))
counter += 1
if use_bn:
layers.append(
ME.MinkowskiBatchNorm(
dim,
eps=1e-5,
momentum=0.1,
))
if (counter < len(dims)) and use_relu:
layers.append(ME.MinkowskiReLU(inplace=True))
last_dim = dim
if use_dropout:
layers.append(MinkowskiDropout.Dropout())
self.clf = nn.Sequential(*layers)
def network_initialization(self, in_channels, out_channels, D):
self.inplanes = self.init_dim
self.conv1 = ME.MinkowskiConvolution(
in_channels, self.inplanes, kernel_size=5, stride=2, dimension=D)
self.bn1 = ME.MinkowskiBatchNorm(self.inplanes)
self.relu = ME.MinkowskiReLU(inplace=True)
self.pool = ME.MinkowskiAvgPooling(kernel_size=2, stride=2, dimension=D)
self.layer1 = self._make_layer(
self.block, self.planes[0], self.layers[0], stride=2)
self.layer2 = self._make_layer(
self.block, self.planes[1], self.layers[1], stride=2)
self.layer3 = self._make_layer(
self.block, self.planes[2], self.layers[2], stride=2)
self.layer4 = self._make_layer(
self.block, self.planes[3], self.layers[3], stride=2)
self.conv5 = ME.MinkowskiConvolution(
self.inplanes, self.inplanes, kernel_size=3, stride=3, dimension=D)
self.bn5 = ME.MinkowskiBatchNorm(self.inplanes)
self.glob_avg = ME.MinkowskiGlobalMaxPooling()
self.final = ME.MinkowskiLinear(self.inplanes, out_channels, bias=True)
def __init__(self, cfg, name='deeplab_chain'):
super(DeepLab_Chain, self).__init__()
self.model_cfg = cfg
self.net = DeepLabUNet(cfg, name='deeplab_unet')
self.F = self.model_cfg.get('num_filters', 16)
self.num_classes = self.model_cfg.get('num_classes', 5)
self.segmentation = ME.MinkowskiLinear(self.F * 2, self.num_classes)
def __init__(self, cfg, name='uresnet_chain'):
super(UResNet_Chain, self).__init__()
self.model_cfg = cfg
self.net = UResNet(cfg, name='uresnet')
self.F = self.model_cfg.get('num_filters', 16)
self.num_classes = self.model_cfg.get('num_classes', 5)
self.segmentation = ME.MinkowskiLinear(self.F, self.num_classes)
def __init__(self,
in_features,
out_features,
stride=1,
dilation=1,
downsample=None,
bn_momentum=0.1,
leakiness=0.0,
dimension=-1):
super(ResNetBlock, self).__init__()
assert dimension > 0
if in_features != out_features:
self.residual = ME.MinkowskiLinear(in_features, out_features)
else:
self.residual = Identity()
self.conv1 = ME.MinkowskiConvolution(in_features,
out_features,
kernel_size=3,
stride=stride,
dilation=dilation,
dimension=dimension)
self.norm1 = ME.MinkowskiBatchNorm(out_features, momentum=bn_momentum)
self.conv2 = ME.MinkowskiConvolution(out_features,
out_features,
kernel_size=3,
stride=1,
dilation=dilation,
dimension=dimension)
self.norm2 = ME.MinkowskiBatchNorm(out_features, momentum=bn_momentum)
self.leaky_relu = MinkowskiLeakyReLU(negative_slope=leakiness)
def __init__(self, cfg, name='acnn'):
super(ACNN_Chain, self).__init__()
self.model_cfg = cfg['modules'][name]
self.net = ACNN(cfg, name='acnn')
self.num_classes = self.model_cfg.get('num_classes', 5)
self.segmentation = ME.MinkowskiLinear(self.net.outputFeatures,
self.num_classes)
print('Total Number of Trainable Parameters = {}'.format(
sum(p.numel() for p in self.parameters() if p.requires_grad)))
def network_initialization(self, in_channels, out_channels, D):
field_ch = 32
field_ch2 = 64
self.field_network = nn.Sequential(
ME.MinkowskiSinusoidal(in_channels, field_ch),
ME.MinkowskiBatchNorm(field_ch),
ME.MinkowskiReLU(inplace=True),
ME.MinkowskiLinear(field_ch, field_ch),
ME.MinkowskiBatchNorm(field_ch),
ME.MinkowskiReLU(inplace=True),
ME.MinkowskiToSparseTensor(),
)
self.field_network2 = nn.Sequential(
ME.MinkowskiSinusoidal(field_ch + in_channels, field_ch2),
ME.MinkowskiBatchNorm(field_ch2),
ME.MinkowskiReLU(inplace=True),
ME.MinkowskiLinear(field_ch2, field_ch2),
ME.MinkowskiBatchNorm(field_ch2),
ME.MinkowskiReLU(inplace=True),
ME.MinkowskiToSparseTensor(),
)
ResNetBase.network_initialization(self, field_ch2, out_channels, D)
def __init__(self,
in_features,
out_features,
kernel_sizes=None,
dilations=None,
mode='avg',
D=3):
super(SPP, self).__init__()
if mode == 'avg':
self.pool_fn = ME.MinkowskiAvgPooling
elif mode == 'max':
self.pool_fn = ME.MinkowskiMaxPooling
elif mode == 'sum':
self.pool_fn = ME.MinkowskiSumPooling
else:
raise ValueError("Invalid pooling mode, must be one of \
'sum', 'max' or 'average'")
self.unpool_fn = ME.MinkowskiPoolingTranspose
# Include global pooling as first modules.
self.pool = [ME.MinkowskiGlobalPooling(dimension=D)]
self.unpool = [ME.MinkowskiBroadcast(dimension=D)]
multiplier = 1
# Define subregion poolings
self.spp = []
if kernel_sizes is not None:
if isinstance(dilations, int):
dilations = [dilations for _ in range(len(kernel_sizes))]
elif isinstance(dilations, list):
assert len(kernel_sizes) == len(dilations)
else:
raise ValueError("Invalid input to dilations, must be either \
int or list of ints")
multiplier = len(kernel_sizes) + 1 # Additional 1 for globalPool
for k, d in zip(kernel_sizes, dilations):
pooling_layer = self.pool_fn(kernel_size=k,
dilation=d,
stride=k,
dimension=D)
unpooling_layer = self.unpool_fn(kernel_size=k,
dilation=d,
stride=k,
dimension=D)
self.pool.append(pooling_layer)
self.unpool.append(unpooling_layer)
self.pool = nn.Sequential(*self.pool)
self.unpool = nn.Sequential(*self.unpool)
self.linear = ME.MinkowskiLinear(in_features * multiplier,
out_features)
def network_initialization(self, in_channels, out_channels, D):
self.inplanes = self.INIT_DIM
self.conv1 = nn.Sequential(
ME.MinkowskiConvolution(in_channels,
self.inplanes,
kernel_size=3,
stride=2,
dimension=D),
ME.MinkowskiBatchNorm(self.inplanes),
ME.MinkowskiReLU(inplace=True),
ME.MinkowskiMaxPooling(kernel_size=2, stride=2, dimension=D),
)
self.layer1 = self._make_layer(self.BLOCK,
self.PLANES[0],
self.LAYERS[0],
stride=2)
self.layer2 = self._make_layer(self.BLOCK,
self.PLANES[1],
self.LAYERS[1],
stride=2)
self.layer3 = self._make_layer(self.BLOCK,
self.PLANES[2],
self.LAYERS[2],
stride=2)
self.layer4 = self._make_layer(self.BLOCK,
self.PLANES[3],
self.LAYERS[3],
stride=2)
self.conv5 = nn.Sequential(
ME.MinkowskiDropout(),
ME.MinkowskiConvolution(self.inplanes,
self.inplanes,
kernel_size=3,
stride=3,
dimension=D),
ME.MinkowskiBatchNorm(self.inplanes),
ME.MinkowskiGELU(),
)
self.glob_pool = ME.MinkowskiGlobalMaxPooling()
self.final = ME.MinkowskiLinear(self.inplanes, out_channels, bias=True)
def __init__(self, in_feat, out_feat, D):
super(ExampleNetwork, self).__init__(D)
self.net = nn.Sequential(
ME.MinkowskiConvolution(in_channels=in_feat,
out_channels=64,
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D), ME.MinkowskiBatchNorm(64),
ME.MinkowskiReLU(),
ME.MinkowskiConvolution(in_channels=64,
out_channels=128,
kernel_size=3,
stride=2,
dimension=D), ME.MinkowskiBatchNorm(128),
ME.MinkowskiReLU(), ME.MinkowskiGlobalPooling(dimension=D),
ME.MinkowskiLinear(128, out_feat))
def __init__(self, in_feat, out_feat, D):
super(ExampleNetwork, self).__init__(D)
self.conv1 = ME.MinkowskiConvolution(in_channels=in_feat,
out_channels=64,
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.bn1 = ME.MinkowskiBatchNorm(64)
self.conv2 = ME.MinkowskiConvolution(in_channels=64,
out_channels=128,
kernel_size=3,
stride=2,
dimension=D)
self.bn2 = ME.MinkowskiBatchNorm(128)
self.pooling = ME.MinkowskiGlobalPooling(dimension=D)
self.linear = ME.MinkowskiLinear(128, out_feat)
def network_initialization(
self, in_channel, out_channel, channels, embedding_channel, kernel_size, D=3,
):
self.mlp1 = self.get_mlp_block(in_channel, channels[0])
self.conv1 = self.get_conv_block(
channels[0], channels[1], kernel_size=kernel_size, stride=1,
)
self.conv2 = self.get_conv_block(
channels[1], channels[2], kernel_size=kernel_size, stride=2,
)
self.conv3 = self.get_conv_block(
channels[2], channels[3], kernel_size=kernel_size, stride=2,
)
self.conv4 = self.get_conv_block(
channels[3], channels[4], kernel_size=kernel_size, stride=2,
)
self.conv5 = nn.Sequential(
self.get_conv_block(
channels[1] + channels[2] + channels[3] + channels[4],
embedding_channel // 4,
kernel_size=3,
stride=2,
),
self.get_conv_block(
embedding_channel // 4, embedding_channel // 2, kernel_size=3, stride=2,
),
self.get_conv_block(
embedding_channel // 2, embedding_channel, kernel_size=3, stride=2,
),
)
self.pool = ME.MinkowskiMaxPooling(kernel_size=3, stride=2, dimension=D)
self.global_max_pool = ME.MinkowskiGlobalMaxPooling()
self.global_avg_pool = ME.MinkowskiGlobalAvgPooling()
self.final = nn.Sequential(
self.get_mlp_block(embedding_channel * 2, 512),
ME.MinkowskiDropout(),
self.get_mlp_block(512, 512),
ME.MinkowskiLinear(512, out_channel, bias=True),
)
def __init__(self,
in_channel,
out_channel,
embedding_channel=1024,
dimension=3):
ME.MinkowskiNetwork.__init__(self, dimension)
self.conv1 = nn.Sequential(
ME.MinkowskiLinear(3, 64, bias=False),
ME.MinkowskiBatchNorm(64),
ME.MinkowskiReLU(),
)
self.conv2 = nn.Sequential(
ME.MinkowskiLinear(64, 64, bias=False),
ME.MinkowskiBatchNorm(64),
ME.MinkowskiReLU(),
)
self.conv3 = nn.Sequential(
ME.MinkowskiLinear(64, 64, bias=False),
ME.MinkowskiBatchNorm(64),
ME.MinkowskiReLU(),
)
self.conv4 = nn.Sequential(
ME.MinkowskiLinear(64, 128, bias=False),
ME.MinkowskiBatchNorm(128),
ME.MinkowskiReLU(),
)
self.conv5 = nn.Sequential(
ME.MinkowskiLinear(128, embedding_channel, bias=False),
ME.MinkowskiBatchNorm(embedding_channel),
ME.MinkowskiReLU(),
)
self.max_pool = ME.MinkowskiGlobalMaxPooling()
self.linear1 = nn.Sequential(
ME.MinkowskiLinear(embedding_channel, 512, bias=False),
ME.MinkowskiBatchNorm(512),
ME.MinkowskiReLU(),
)
self.dp1 = ME.MinkowskiDropout()
self.linear2 = ME.MinkowskiLinear(512, out_channel, bias=True)
def __init__(self, in_features, out_features, dimension=3, leakiness=0.0):
super(AtrousIIBlock, self).__init__()
assert dimension > 0
self.D = dimension
if in_features != out_features:
self.residual = ME.MinkowskiLinear(in_features, out_features)
else:
self.residual = Identity()
self.conv1 = ME.MinkowskiConvolution(in_features,
out_features,
kernel_size=3,
stride=1,
dilation=1,
dimension=self.D)
self.norm1 = ME.MinkowskiInstanceNorm(out_features, dimension=self.D)
self.conv2 = ME.MinkowskiConvolution(out_features,
out_features,
kernel_size=3,
stride=1,
dilation=3,
dimension=self.D)
self.norm2 = ME.MinkowskiInstanceNorm(out_features, dimension=self.D)
self.leaky_relu = MinkowskiLeakyReLU(negative_slope=leakiness)
def __init__(self,
in_channel,
out_channel,
num_class,
embedding_channel=1024,
dimension=3):
ME.MinkowskiNetwork.__init__(self, dimension)
# The normal channel for Modelnet is 3, for scannet is 6, for scanobjnn is 0
normal_channel = 3
# in_channel = normal_channel+3 # normal ch + xyz
self.normal_channel = normal_channel
self.input_mlp = nn.Sequential(
ME.MinkowskiConvolution(in_channel, 32, kernel_size=1,
dimension=3), ME.MinkowskiBatchNorm(32),
ME.MinkowskiReLU(),
ME.MinkowskiConvolution(32, 32, kernel_size=1, dimension=3),
ME.MinkowskiBatchNorm(32))
self.in_dims = [32, 64, 128, 256]
self.out_dims = [64, 128, 256, 512]
self.neighbor_ks = [16, 32, 64, 16, 16]
# self.neighbor_ks = [8, 8, 8, 8, 8]
self.PTBlock0 = PTBlock(in_dim=self.in_dims[0],
n_sample=self.neighbor_ks[0])
self.TDLayer1 = TDLayer(input_dim=self.in_dims[0],
out_dim=self.out_dims[0])
self.PTBlock1 = PTBlock(in_dim=self.out_dims[0],
n_sample=self.neighbor_ks[1])
self.TDLayer2 = TDLayer(input_dim=self.in_dims[1],
out_dim=self.out_dims[1])
self.PTBlock2 = PTBlock(in_dim=self.out_dims[1],
n_sample=self.neighbor_ks[1])
self.TDLayer3 = TDLayer(input_dim=self.in_dims[2],
out_dim=self.out_dims[2])
self.PTBlock3 = PTBlock(in_dim=self.out_dims[2],
n_sample=self.neighbor_ks[2])
self.TDLayer4 = TDLayer(input_dim=self.in_dims[3],
out_dim=self.out_dims[3])
self.PTBlock4 = PTBlock(in_dim=self.out_dims[3],
n_sample=self.neighbor_ks[4])
self.middle_linear = ME.MinkowskiConvolution(self.out_dims[3],
self.out_dims[3],
kernel_size=1,
dimension=3)
self.PTBlock_middle = PTBlock(in_dim=self.out_dims[3],
n_sample=self.neighbor_ks[4])
self.TULayer5 = TULayer(input_dim=self.out_dims[3],
out_dim=self.in_dims[3])
self.PTBlock5 = PTBlock(in_dim=self.in_dims[3],
n_sample=self.neighbor_ks[4])
self.TULayer6 = TULayer(input_dim=self.out_dims[2],
out_dim=self.in_dims[2])
self.PTBlock6 = PTBlock(in_dim=self.in_dims[2],
n_sample=self.neighbor_ks[3])
self.TULayer7 = TULayer(input_dim=self.out_dims[1],
out_dim=self.in_dims[1])
self.PTBlock7 = PTBlock(in_dim=self.in_dims[1],
n_sample=self.neighbor_ks[2])
self.TULayer8 = TULayer(input_dim=self.out_dims[0],
out_dim=self.in_dims[0])
self.PTBlock8 = PTBlock(in_dim=self.in_dims[0],
n_sample=self.neighbor_ks[1])
self.fc = nn.Sequential(
# ME.MinkowskiLinear(32, 32),
ME.MinkowskiLinear(self.out_dims[3], 32),
ME.MinkowskiDropout(0.4),
ME.MinkowskiLinear(32, num_class))
self.global_avg_pool = ME.MinkowskiGlobalAvgPooling()
def __init__(self,
config,
in_channel,
out_channel,
final_dim=96,
dimension=3):
ME.MinkowskiNetwork.__init__(self, dimension)
# The normal channel for Modelnet is 3, for scannet is 6, for scanobjnn is 0
normal_channel = 3
self.CONV_TYPE = ConvType.SPATIAL_HYPERCUBE
# self.dims = np.array([32, 64, 128, 256])
self.dims = np.array([32, 64, 128, 256])
# self.neighbor_ks = np.array([4, 4, 4, 4])
# self.neighbor_ks = np.array([12, 12, 12, 12])
self.neighbor_ks = np.array([16, 16, 16, 16])
# self.neighbor_ks = np.array([32, 32, 32, 32])
# self.neighbor_ks = np.array([8, 8, 16, 16])
# self.neighbor_ks = np.array([32, 32, 32, 32])
# self.neighbor_ks = np.array([8, 12, 16, 24])
# self.neighbor_ks = np.array([8, 8, 8, 8])
self.final_dim = final_dim
stem_dim = self.dims[0]
if config.xyz_input:
in_channel = normal_channel + in_channel
else:
in_channel = in_channel
# pixel size 1
self.stem1 = nn.Sequential(
ME.MinkowskiConvolution(in_channel,
stem_dim,
kernel_size=config.ks,
dimension=3),
ME.MinkowskiBatchNorm(stem_dim),
ME.MinkowskiReLU(),
)
# self.stem2 = TDLayer(input_dim=self.dims[0], out_dim=self.dims[0])
self.stem2 = nn.Sequential(
ME.MinkowskiConvolution(stem_dim,
stem_dim,
kernel_size=3,
dimension=3,
stride=2),
ME.MinkowskiBatchNorm(stem_dim),
ME.MinkowskiReLU(),
)
# window_beta = 5
window_beta = None
base_r = 5
self.ALPHA_BLENDING_MID_TR = False
self.ALPHA_BLENDING_FIRST_TR = True
# self.PTBlock1 = PTBlock(in_dim=self.dims[0], hidden_dim = self.dims[0], n_sample=self.neighbor_ks[0], skip_attn=False, r=base_r, kernel_size=config.ks, window_beta=window_beta)
self.PTBlock2 = PTBlock(in_dim=self.dims[1],
hidden_dim=self.dims[1],
n_sample=self.neighbor_ks[1],
skip_attn=False,
r=2 * base_r,
kernel_size=config.ks,
window_beta=window_beta)
self.PTBlock3 = PTBlock(in_dim=self.dims[2],
hidden_dim=self.dims[2],
n_sample=self.neighbor_ks[2],
skip_attn=False,
r=2 * base_r,
kernel_size=config.ks,
window_beta=window_beta)
self.PTBlock4 = PTBlock(in_dim=self.dims[3],
hidden_dim=self.dims[3],
n_sample=self.neighbor_ks[3],
skip_attn=False,
r=4 * base_r,
kernel_size=config.ks,
window_beta=window_beta)
if self.ALPHA_BLENDING_MID_TR:
self.PTBlock4_branch = PTBlock(in_dim=self.dims[3],
hidden_dim=self.dims[3],
n_sample=self.neighbor_ks[3],
skip_attn=True,
r=4 * base_r,
kernel_size=config.ks)
self.PTBlock5_branch = PTBlock(in_dim=128,
hidden_dim=128,
n_sample=self.neighbor_ks[3],
skip_attn=True,
r=2 * base_r,
kernel_size=config.ks) # out: 256
if self.ALPHA_BLENDING_FIRST_TR:
self.PTBlock1_branch = PTBlock(in_dim=self.dims[0],
hidden_dim=self.dims[0],
n_sample=self.neighbor_ks[0],
skip_attn=False,
r=base_r,
kernel_size=config.ks,
window_beta=window_beta)
self.PTBlock5 = PTBlock(in_dim=128,
hidden_dim=128,
n_sample=self.neighbor_ks[3],
skip_attn=False,
r=2 * base_r,
kernel_size=config.ks,
window_beta=window_beta) # out: 256
self.PTBlock6 = PTBlock(in_dim=128,
hidden_dim=128,
n_sample=self.neighbor_ks[2],
skip_attn=False,
r=2 * base_r,
kernel_size=config.ks,
window_beta=window_beta) # out: 128
self.PTBlock7 = PTBlock(in_dim=96,
hidden_dim=96,
n_sample=self.neighbor_ks[1],
skip_attn=False,
r=base_r,
kernel_size=config.ks,
window_beta=window_beta) # out: 64
# self.PTBlock5 = StackedPTBlock(in_dim=128, hidden_dim=128,n_sample=self.neighbor_ks[3], skip_attn=False, r=2*base_r, kernel_size=config.ks) # out: 256
# self.PTBlock6 = StackedPTBlock(in_dim=128, hidden_dim=128, n_sample=self.neighbor_ks[2], skip_attn=False, r=2*base_r, kernel_size=config.ks) # out: 128
# self.PTBlock7 = StackedPTBlock(in_dim=96, hidden_dim=96, n_sample=self.neighbor_ks[1], skip_attn=False, r=base_r, kernel_size=config.ks) # out: 64
# BLOCK_TYPE = SingleConv
BLOCK_TYPE = BasicBlock
self.PTBlock1 = self._make_layer(block=BLOCK_TYPE,
inplanes=self.dims[0],
planes=self.dims[0],
num_blocks=1)
# self.PTBlock2 = self._make_layer(block=BLOCK_TYPE, inplanes=self.dims[1], planes=self.dims[1], num_blocks=1)
# self.PTBlock3 = self._make_layer(block=BLOCK_TYPE, inplanes=self.dims[2], planes=self.dims[2], num_blocks=1)
# self.PTBlock4 = self._make_layer(block=BLOCK_TYPE, inplanes=self.dims[3], planes=self.dims[3], num_blocks=1)
# self.PTBlock5 = self._make_layer(block=BLOCK_TYPE, inplanes=128, planes=128, num_blocks=1)
# self.PTBlock6 = self._make_layer(block=BLOCK_TYPE, inplanes=128, planes=128, num_blocks=1)
# self.PTBlock7 = self._make_layer(block=BLOCK_TYPE, inplanes=96, planes=96, num_blocks=1)
# pixel size 2
self.TDLayer1 = TDLayer(input_dim=self.dims[0],
out_dim=self.dims[1]) # strided conv
# pixel size 4
self.TDLayer2 = TDLayer(input_dim=self.dims[1], out_dim=self.dims[2])
# pixel size 8
self.TDLayer3 = TDLayer(input_dim=self.dims[2], out_dim=self.dims[3])
# pixel size 16: PTBlock4
# pixel size 8
# self.TULayer5 = TULayer(input_a_dim=self.dims[3], input_b_dim = self.dims[2], out_dim=self.dims[2]) # out: 256//2 + 128 = 256
self.TULayer5 = ResNetLikeTU(input_a_dim=256,
input_b_dim=128,
out_dim=128) # out: 256//2 + 128 = 256
# pixel size 4
# self.TULayer6 = TULayer(input_a_dim=self.dims[2], input_b_dim = self.dims[1], out_dim=self.dims[1]) # out: 256//2 + 64 = 192
self.TULayer6 = ResNetLikeTU(input_a_dim=128,
input_b_dim=64,
out_dim=128) # out: 256//2 + 64 = 192
# pixel size 2
# self.TULayer7 = TULayer(input_a_dim=self.dims[1], input_b_dim = self.dims[0], out_dim=self.dims[0]) # 128 // 2 + 32 = 96
self.TULayer7 = ResNetLikeTU(input_a_dim=128,
input_b_dim=32,
out_dim=96) # 128 // 2 + 32 = 96
# pixel size 1
# self.PTBlock8 = PTBlock(in_dim=self.dims[0], hidden_dim=self.dims[0], n_sample=self.neighbor_ks[1]) # 32
# self.TULayer8 = TULayer(input_a_dim=self.dims[0], input_b_dim = self.dims[0], out_dim=self.dims[0]) # 64 // 2 + 32
self.TULayer8 = ResNetLikeTU(input_a_dim=96,
input_b_dim=32,
out_dim=96) # 64 // 2 + 32
self.fc = ME.MinkowskiLinear(96, out_channel)
def __init__(self):
nn.Module.__init__(self)
# Input sparse tensor must have tensor stride 128.
ch = self.CHANNELS
# Block 1
self.block1 = nn.Sequential(
ME.MinkowskiConvolution(1, ch[0], kernel_size=3, stride=2, dimension=3),
ME.MinkowskiBatchNorm(ch[0]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(ch[0], ch[0], kernel_size=3, dimension=3),
ME.MinkowskiBatchNorm(ch[0]),
ME.MinkowskiELU(),
)
self.block2 = nn.Sequential(
ME.MinkowskiConvolution(ch[0], ch[1], kernel_size=3, stride=2, dimension=3),
ME.MinkowskiBatchNorm(ch[1]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(ch[1], ch[1], kernel_size=3, dimension=3),
ME.MinkowskiBatchNorm(ch[1]),
ME.MinkowskiELU(),
)
self.block3 = nn.Sequential(
ME.MinkowskiConvolution(ch[1], ch[2], kernel_size=3, stride=2, dimension=3),
ME.MinkowskiBatchNorm(ch[2]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(ch[2], ch[2], kernel_size=3, dimension=3),
ME.MinkowskiBatchNorm(ch[2]),
ME.MinkowskiELU(),
)
self.block4 = nn.Sequential(
ME.MinkowskiConvolution(ch[2], ch[3], kernel_size=3, stride=2, dimension=3),
ME.MinkowskiBatchNorm(ch[3]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(ch[3], ch[3], kernel_size=3, dimension=3),
ME.MinkowskiBatchNorm(ch[3]),
ME.MinkowskiELU(),
)
self.block5 = nn.Sequential(
ME.MinkowskiConvolution(ch[3], ch[4], kernel_size=3, stride=2, dimension=3),
ME.MinkowskiBatchNorm(ch[4]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(ch[4], ch[4], kernel_size=3, dimension=3),
ME.MinkowskiBatchNorm(ch[4]),
ME.MinkowskiELU(),
)
self.block6 = nn.Sequential(
ME.MinkowskiConvolution(ch[4], ch[5], kernel_size=3, stride=2, dimension=3),
ME.MinkowskiBatchNorm(ch[5]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(ch[5], ch[5], kernel_size=3, dimension=3),
ME.MinkowskiBatchNorm(ch[5]),
ME.MinkowskiELU(),
)
self.block7 = nn.Sequential(
ME.MinkowskiConvolution(ch[5], ch[6], kernel_size=3, stride=2, dimension=3),
ME.MinkowskiBatchNorm(ch[6]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(ch[6], ch[6], kernel_size=3, dimension=3),
ME.MinkowskiBatchNorm(ch[6]),
ME.MinkowskiELU(),
)
self.global_pool = ME.MinkowskiGlobalPooling()
self.linear_mean = ME.MinkowskiLinear(ch[6], ch[6], bias=True)
self.linear_log_var = ME.MinkowskiLinear(ch[6], ch[6], bias=True)
self.weight_initialization()
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.MinkowskiReLU(),
)
def __init__(self, config, in_channel, out_channel, final_dim=96, dimension=3):
ME.MinkowskiNetwork.__init__(self, dimension)
# The normal channel for Modelnet is 3, for scannet is 6, for scanobjnn is 0
normal_channel = 3
self.CONV_TYPE = ConvType.SPATIAL_HYPERCUBE
# self.dims = np.array([32, 64, 128, 256])
self.dims = np.array([64, 64, 256, 256])
# self.neighbor_ks = np.array([12, 12, 12, 12])
self.neighbor_ks = np.array([16, 16, 16, 16])
# self.neighbor_ks = np.array([8, 12, 16, 24])
# self.neighbor_ks = np.array([8, 8, 8, 8])
self.final_dim = final_dim
stem_dim = 32
if config.xyz_input:
in_channel = normal_channel + in_channel
else:
in_channel = in_channel
# pixel size 1
self.stem1 = nn.Sequential(
ME.MinkowskiConvolution(in_channel, stem_dim, kernel_size=config.ks, dimension=3),
ME.MinkowskiBatchNorm(stem_dim),
ME.MinkowskiReLU(),
)
# self.stem2 = TDLayer(input_dim=self.dims[0], out_dim=self.dims[0])
# DEBUG ONLY!!
self.stem2 = nn.Sequential(
ME.MinkowskiConvolution(stem_dim, stem_dim*2, kernel_size=4, dimension=3, stride=4),
ME.MinkowskiBatchNorm(stem_dim*2),
ME.MinkowskiReLU(),
)
base_r = 5
self.PTBlock1 = PTBlock(in_dim=self.dims[0], hidden_dim = self.dims[0], n_sample=self.neighbor_ks[0], skip_knn=False, r=base_r, kernel_size=config.ks)
self.PTBlock2 = PTBlock(in_dim=self.dims[1], hidden_dim = self.dims[1], n_sample=self.neighbor_ks[1], skip_knn=False, r=2*base_r, kernel_size=config.ks)
self.PTBlock3 = PTBlock(in_dim=self.dims[2],hidden_dim = self.dims[2], n_sample=self.neighbor_ks[2], skip_knn=False, r=2*base_r, kernel_size=config.ks)
self.PTBlock4 = PTBlock(in_dim=self.dims[3], hidden_dim = self.dims[3], n_sample=self.neighbor_ks[3], skip_knn=False, r=4*base_r, kernel_size=config.ks)
# self.PTBlock5 = PTBlock(in_dim=self.dims[2], hidden_dim = self.dims[2], n_sample=self.neighbor_ks[3], skip_knn=False, r=2*base_r, kernel_size=config.ks) # out: 256
# self.PTBlock6 = PTBlock(in_dim=self.dims[1], hidden_dim=self.dims[1], n_sample=self.neighbor_ks[2], skip_knn=False, r=2*base_r, kernel_size=config.ks) # out: 128
# self.PTBlock7 = PTBlock(in_dim=self.dims[0], hidden_dim=self.dims[0], n_sample=self.neighbor_ks[1], skip_knn=False, r=base_r, kernel_size=config.ks) # out: 64
self.PTBlock5 = PTBlock(in_dim=256, hidden_dim=256,n_sample=self.neighbor_ks[3], skip_knn=False, r=2*base_r, kernel_size=config.ks) # out: 256
self.PTBlock6 = PTBlock(in_dim=128, hidden_dim=128, n_sample=self.neighbor_ks[2], skip_knn=False, r=2*base_r, kernel_size=config.ks) # out: 128
self.PTBlock7 = PTBlock(in_dim=128, hidden_dim=128, n_sample=self.neighbor_ks[1], skip_knn=False, r=base_r, kernel_size=config.ks) # out: 64
# self.PTBlock1 = StackedPTBlock(in_dim=self.dims[0], hidden_dim = self.dims[0], n_sample=self.neighbor_ks[0], skip_knn=False, r=base_r, kernel_size=config.ks)
# self.PTBlock2 = StackedPTBlock(in_dim=self.dims[1], hidden_dim = self.dims[1], n_sample=self.neighbor_ks[1], skip_knn=False, r=2*base_r, kernel_size=config.ks)
# self.PTBlock3 = StackedPTBlock(in_dim=self.dims[2],hidden_dim = self.dims[2], n_sample=self.neighbor_ks[2], skip_knn=False, r=2*base_r, kernel_size=config.ks)
# self.PTBlock4 = StackedPTBlock(in_dim=self.dims[3], hidden_dim = self.dims[3], n_sample=self.neighbor_ks[3], skip_knn=False, r=4*base_r, kernel_size=config.ks)
# self.PTBlock5 = StackedPTBlock(in_dim=self.dims[2], hidden_dim = self.dims[2], n_sample=self.neighbor_ks[3], skip_knn=False, r=2*base_r, kernel_size=config.ks) # out: 256
# self.PTBlock6 = StackedPTBlock(in_dim=self.dims[1], hidden_dim=self.dims[1], n_sample=self.neighbor_ks[2], skip_knn=False, r=2*base_r, kernel_size=config.ks) # out: 128
# self.PTBlock7 = StackedPTBlock(in_dim=self.dims[0], hidden_dim=self.dims[0], n_sample=self.neighbor_ks[1], skip_knn=False, r=base_r, kernel_size=config.ks) # out: 64
# self.PTBlock1 = self._make_layer(block=BasicBlock, inplanes=self.dims[0], planes=self.dims[0], num_blocks=2)
# self.PTBlock2 = self._make_layer(block=BasicBlock, inplanes=self.dims[1], planes=self.dims[1], num_blocks=2)
# self.PTBlock3 = self._make_layer(block=BasicBlock, inplanes=self.dims[2], planes=self.dims[2], num_blocks=2)
# self.PTBlock4 = self._make_layer(block=BasicBlock, inplanes=self.dims[3], planes=self.dims[3], num_blocks=2)
# self.PTBlock5 = self._make_layer(block=BasicBlock, inplanes=self.dims[2], planes=self.dims[2], num_blocks=2)
# self.PTBlock6 = self._make_layer(block=BasicBlock, inplanes=self.dims[1], planes=self.dims[1], num_blocks=2)
# self.PTBlock7 = self._make_layer(block=BasicBlock, inplanes=self.dims[0], planes=self.dims[0], num_blocks=2)
# pixel size 2
# self.TDLayer1 = TDLayer(input_dim=self.dims[0], out_dim=self.dims[1]) # strided conv
# pixel size 4
self.TDLayer2 = TDLayer(input_dim=self.dims[1], out_dim=self.dims[2])
# pixel size 8
# self.TDLayer3 = TDLayer(input_dim=self.dims[2], out_dim=self.dims[3])
# pixel size 16: PTBlock4
# pixel size 8
# self.TULayer5 = TULayer(input_a_dim=self.dims[3], input_b_dim = self.dims[2], out_dim=self.dims[2]) # out: 256//2 + 128 = 256
# self.TULayer5 = ResNetLikeTU(input_a_dim=256, input_b_dim = 128, out_dim=128) # out: 256//2 + 128 = 256
# pixel size 4
# self.TULayer6 = TULayer(input_a_dim=self.dims[2], input_b_dim = self.dims[1], out_dim=self.dims[1]) # out: 256//2 + 64 = 192
self.TULayer6 = ResNetLikeTU(input_a_dim=256, input_b_dim = 64, out_dim=128) # out: 256//2 + 64 = 192
# pixel size 2
# self.TULayer7 = TULayer(input_a_dim=self.dims[1], input_b_dim = self.dims[0], out_dim=self.dims[0]) # 128 // 2 + 32 = 96
# self.TULayer7 = ResNetLikeTU(input_a_dim=128, input_b_dim = 32, out_dim=96) # 128 // 2 + 32 = 96
# pixel size 1
# self.PTBlock8 = PTBlock(in_dim=self.dims[0], hidden_dim=self.dims[0], n_sample=self.neighbor_ks[1]) # 32
# self.TULayer8 = TULayer(input_a_dim=self.dims[0], input_b_dim = self.dims[0], out_dim=self.dims[0]) # 64 // 2 + 32
self.TULayer8 = ResNetLikeTU(input_a_dim=128, input_b_dim = 32, out_dim=96) # 64 // 2 + 32
self.fc = ME.MinkowskiLinear(96, out_channel)
def __init__(self,
config,
in_channel,
out_channel,
final_dim=96,
dimension=3):
ME.MinkowskiNetwork.__init__(
self, dimension
) # The normal channel for Modelnet is 3, for scannet is 6, for scanobjnn is 0
normal_channel = 3 # the RGB
self.CONV_TYPE = ConvType.SPATIAL_HYPERCUBE
self.dims = np.array([32, 64, 128, 256, 512])
self.neighbor_ks = np.array([32, 32, 32, 32, 32]) // 2
self.final_dim = final_dim
stem_dim = self.dims[0]
in_channel = normal_channel + 3 # normal ch + xyz
self.normal_channel = normal_channel
# pixel size 1
self.stem1 = nn.Sequential(
ME.MinkowskiConvolution(in_channel,
stem_dim,
kernel_size=1,
dimension=3),
ME.MinkowskiBatchNorm(stem_dim),
ME.MinkowskiReLU(),
)
# when using split-scnee, no stride here
split_scene = True
if split_scene:
self.stem2 = nn.Sequential(
ME.MinkowskiConvolution(stem_dim,
stem_dim,
kernel_size=1,
dimension=3,
stride=1),
ME.MinkowskiBatchNorm(stem_dim),
ME.MinkowskiReLU(),
)
# does the spatial downsampling
# pixel size 2
else:
self.stem2 = nn.Sequential(
ME.MinkowskiConvolution(stem_dim,
stem_dim,
kernel_size=2,
dimension=3,
stride=2),
ME.MinkowskiBatchNorm(stem_dim),
ME.MinkowskiReLU(),
)
base_r = 4
self.PTBlock0 = PTBlock(in_dim=self.dims[0],
hidden_dim=self.dims[0],
n_sample=self.neighbor_ks[0],
skip_knn=False,
r=base_r)
self.PTBlock1 = PTBlock(in_dim=self.dims[1],
hidden_dim=self.dims[1],
n_sample=self.neighbor_ks[1],
skip_knn=False,
r=base_r)
self.PTBlock2 = PTBlock(in_dim=self.dims[2],
hidden_dim=self.dims[2],
n_sample=self.neighbor_ks[2],
skip_knn=False,
r=2 * base_r)
self.PTBlock3 = PTBlock(in_dim=self.dims[3],
hidden_dim=self.dims[3],
n_sample=self.neighbor_ks[3],
skip_knn=False,
r=int(2 * base_r))
self.PTBlock4 = PTBlock(in_dim=self.dims[4],
hidden_dim=self.dims[4],
n_sample=self.neighbor_ks[3],
skip_knn=False,
r=int(4 * base_r))
self.PTBlock_middle = PTBlock(in_dim=self.dims[4],
hidden_dim=self.dims[4],
n_sample=self.neighbor_ks[3],
skip_knn=False,
r=int(4 * base_r))
self.PTBlock5 = PTBlock(in_dim=self.dims[3],
hidden_dim=self.dims[3],
n_sample=self.neighbor_ks[3],
skip_knn=False,
r=2 * base_r) # out: 256
self.PTBlock6 = PTBlock(in_dim=self.dims[2],
hidden_dim=self.dims[2],
n_sample=self.neighbor_ks[2],
skip_knn=False,
r=2 * base_r) # out: 128
self.PTBlock7 = PTBlock(in_dim=self.dims[1],
hidden_dim=self.dims[1],
n_sample=self.neighbor_ks[1],
skip_knn=False,
r=base_r) # out: 64
self.PTBlock8 = PTBlock(in_dim=self.dims[0],
hidden_dim=self.dims[0],
n_sample=self.neighbor_ks[1],
skip_knn=False,
r=base_r) # out: 64
# self.PTBlock0 = PTBlock(in_dim=self.dims[0], hidden_dim = self.dims[0], n_sample=self.neighbor_ks[0], skip_knn=True, r=base_r)
# self.PTBlock1 = PTBlock(in_dim=self.dims[1], hidden_dim = self.dims[1], n_sample=self.neighbor_ks[1], skip_knn=True, r=2*base_r)
# self.PTBlock2 = PTBlock(in_dim=self.dims[2],hidden_dim = self.dims[2], n_sample=self.neighbor_ks[2], skip_knn=True, r=2*base_r)
# self.PTBlock3 = PTBlock(in_dim=self.dims[3], hidden_dim = self.dims[3], n_sample=self.neighbor_ks[3], skip_knn=True, r=int(4*base_r))
# self.PTBlock4 = PTBlock(in_dim=self.dims[4], hidden_dim = self.dims[4], n_sample=self.neighbor_ks[3], skip_knn=True, r=int(16*base_r))
# self.PTBlock_middle = PTBlock(in_dim=self.dims[4], hidden_dim = self.dims[4], n_sample=self.neighbor_ks[3], skip_knn=True, r=int(16*base_r))
# self.PTBlock5 = PTBlock(in_dim=self.dims[3], hidden_dim = self.dims[3], n_sample=self.neighbor_ks[3], skip_knn=True, r=4*base_r) # out: 256
# self.PTBlock6 = PTBlock(in_dim=self.dims[2], hidden_dim=self.dims[2], n_sample=self.neighbor_ks[2], skip_knn=True, r=2*base_r) # out: 128
# self.PTBlock7 = PTBlock(in_dim=self.dims[1], hidden_dim=self.dims[1], n_sample=self.neighbor_ks[1], skip_knn=True, r=2*base_r) # out: 64
# self.PTBlock8 = PTBlock(in_dim=self.dims[0], hidden_dim=self.dims[0], n_sample=self.neighbor_ks[1], skip_knn=True, r=base_r) # out: 64
# self.PTBlock1 = self._make_layer(block=BasicBlock, inplanes=self.dims[0], planes=self.dims[0], num_blocks=2)
# self.PTBlock2 = self._make_layer(block=BasicBlock, inplanes=self.dims[1], planes=self.dims[1], num_blocks=2)
# self.PTBlock3 = self._make_layer(block=BasicBlock, inplanes=self.dims[2], planes=self.dims[2], num_blocks=2)
# self.PTBlock4 = self._make_layer(block=BasicBlock, inplanes=self.dims[3], planes=self.dims[3], num_blocks=2)
# self.PTBlock_middle = self._make_layer(block=BasicBlock, inplanes=self.dims[4], planes=self.dims[4], num_blocks=2)
# self.PTBlock5 = self._make_layer(block=BasicBlock, inplanes=self.dims[3], planes=self.dims[3], num_blocks=2)
# self.PTBlock6 = self._make_layer(block=BasicBlock, inplanes=self.dims[2], planes=self.dims[2], num_blocks=2)
# self.PTBlock7 = self._make_layer(block=BasicBlock, inplanes=self.dims[1], planes=self.dims[1], num_blocks=2)
# self.PTBlock8 = self._make_layer(block=BasicBlock, inplanes=self.dims[0], planes=self.dims[0], num_blocks=2)
TD_kernel_size = [4, 8, 12,
16] # only applied when using non-PointTRlike
# pixel size 2
self.TDLayer1 = TDLayer(input_dim=self.dims[0],
out_dim=self.dims[1],
kernel_size=TD_kernel_size[0]) # strided conv
# pixel size 4
self.TDLayer2 = TDLayer(input_dim=self.dims[1],
out_dim=self.dims[2],
kernel_size=TD_kernel_size[1])
# pixel size 8
self.TDLayer3 = TDLayer(input_dim=self.dims[2],
out_dim=self.dims[3],
kernel_size=TD_kernel_size[2])
# pixel size 16: PTBlock4
self.TDLayer4 = TDLayer(input_dim=self.dims[3],
out_dim=self.dims[4],
kernel_size=TD_kernel_size[3])
self.middle_linear = ME.MinkowskiConvolution(self.dims[4],
self.dims[4],
kernel_size=1,
dimension=3)
# pixel size 8
self.TULayer5 = TULayer(
input_a_dim=self.dims[4],
input_b_dim=self.dims[3],
out_dim=self.dims[3]) # out: 256//2 + 128 = 256
# pixel size 4
self.TULayer6 = TULayer(input_a_dim=self.dims[3],
input_b_dim=self.dims[2],
out_dim=self.dims[2]) # out: 256//2 + 64 = 192
# pixel size 2
self.TULayer7 = TULayer(input_a_dim=self.dims[2],
input_b_dim=self.dims[1],
out_dim=self.dims[1]) # 128 // 2 + 32 = 96
self.TULayer8 = TULayer(input_a_dim=self.dims[1],
input_b_dim=self.dims[0],
out_dim=self.dims[0]) # 128 // 2 + 32 = 96
self.final_dim = 32
if split_scene:
self.final_conv = nn.Sequential(
ME.MinkowskiConvolution(self.dims[0],
self.final_dim,
kernel_size=1,
stride=1,
dimension=3),
ME.MinkowskiDropout(0.4),
)
else:
self.final_conv = nn.Sequential(
ME.MinkowskiConvolutionTranspose(self.dims[0],
self.final_dim,
kernel_size=2,
stride=2,
dimension=3),
ME.MinkowskiDropout(0.4),
)
self.fc = ME.MinkowskiLinear(self.final_dim, out_channel)
def __init__(self, cfg, name='fpn'):
super(FPN, self).__init__(cfg)
model_cfg = cfg['modules'][name]
# UResNet Configurations
self.reps = model_cfg.get('reps', 2)
self.depth = model_cfg.get('depth', 5)
self.num_filters = model_cfg.get('num_filters', 16)
self.nPlanes = [i * self.num_filters for i in range(1, self.depth + 1)]
# self.kernel_size = cfg.get('kernel_size', 3)
# self.downsample = cfg.get(downsample, 2)
self.input_kernel = model_cfg.get('input_kernel', 3)
# Initialize Input Layer
self.input_layer = ME.MinkowskiConvolution(
in_channels=self.num_input,
out_channels=self.num_filters,
kernel_size=self.input_kernel,
stride=1,
dimension=self.D)
# Initialize Encoder
self.encoding_conv = []
self.encoding_block = []
for i, F in enumerate(self.nPlanes):
m = []
for _ in range(self.reps):
m.append(
ResNetBlock(F,
F,
dimension=self.D,
leakiness=self.leakiness))
m = nn.Sequential(*m)
self.encoding_block.append(m)
m = []
if i < self.depth - 1:
m.append(ME.MinkowskiBatchNorm(F))
m.append(MinkowskiLeakyReLU(negative_slope=self.leakiness))
m.append(
ME.MinkowskiConvolution(in_channels=self.nPlanes[i],
out_channels=self.nPlanes[i + 1],
kernel_size=2,
stride=2,
dimension=self.D))
m = nn.Sequential(*m)
self.encoding_conv.append(m)
self.encoding_conv = nn.Sequential(*self.encoding_conv)
self.encoding_block = nn.Sequential(*self.encoding_block)
# Lateral Connections
self.lateral = []
for i, F in enumerate(self.nPlanes[-2::-1]):
self.lateral.append(ME.MinkowskiLinear(F, F))
self.lateral = nn.Sequential(*self.lateral)
# Initialize Decoder
self.decoding_block = []
self.decoding_conv = []
for i in range(self.depth - 2, -1, -1):
m = []
m.append(ME.MinkowskiBatchNorm(self.nPlanes[i + 1]))
m.append(MinkowskiLeakyReLU(negative_slope=self.leakiness))
m.append(
ME.MinkowskiConvolutionTranspose(in_channels=self.nPlanes[i +
1],
out_channels=self.nPlanes[i],
kernel_size=2,
stride=2,
dimension=self.D))
m = nn.Sequential(*m)
self.decoding_conv.append(m)
m = []
for j in range(self.reps):
m.append(
ResNetBlock(self.nPlanes[i],
self.nPlanes[i],
dimension=self.D,
leakiness=self.leakiness))
m = nn.Sequential(*m)
self.decoding_block.append(m)
self.decoding_block = nn.Sequential(*self.decoding_block)
self.decoding_conv = nn.Sequential(*self.decoding_conv)
print('Total Number of Trainable Parameters = {}'.format(
sum(p.numel() for p in self.parameters() if p.requires_grad)))
def __init__(self,
in_features,
out_features,
dimension=3,
leakiness=0.0,
cardinality=4,
depth=1,
dilations=None,
kernel_sizes=3,
strides=1):
super(ResNeXtBlock, self).__init__()
assert dimension > 0
assert cardinality > 0
assert (in_features % cardinality == 0
and out_features % cardinality == 0)
self.D = dimension
nIn = in_features // cardinality
nOut = out_features // cardinality
self.dilations = []
if dilations is None:
# Default
self.dilations = [3**i for i in range(cardinality)]
elif isinstance(dilations, int):
self.dilations = [dilations for _ in range(cardinality)]
elif isinstance(dilations, list):
assert len(dilations) == cardinality
self.dilations = dilations
else:
raise ValueError(
'Invalid type for input strides, must be int or list!')
self.kernels = []
if isinstance(kernel_sizes, int):
self.kernels = [kernel_sizes for _ in range(cardinality)]
elif isinstance(kernel_sizes, list):
assert len(kernel_sizes) == cardinality
self.kernels = kernel_sizes
else:
raise ValueError(
'Invalid type for input strides, must be int or list!')
self.strides = []
if isinstance(strides, int):
self.strides = [strides for _ in range(cardinality)]
elif isinstance(strides, list):
assert len(strides) == cardinality
self.strides = strides
else:
raise ValueError(
'Invalid type for input strides, must be int or list!')
# For each path, generate sequentials
self.paths = []
for i in range(cardinality):
m = []
m.append(ME.MinkowskiLinear(in_features, nIn))
for j in range(depth):
in_C = (nIn if j == 0 else nOut)
m.append(
ME.MinkowskiConvolution(in_channels=in_C,
out_channels=nOut,
kernel_size=self.kernels[i],
stride=self.strides[i],
dilation=self.dilations[i],
dimension=self.D))
m.append(ME.MinkowskiBatchNorm(nOut))
m.append(MinkowskiLeakyReLU(negative_slope=leakiness))
m = nn.Sequential(*m)
self.paths.append(m)
self.paths = nn.Sequential(*self.paths)
self.linear = ME.MinkowskiLinear(out_features, out_features)
# Skip Connection
if in_features != out_features:
self.residual = ME.MinkowskiLinear(in_features, out_features)
else:
self.residual = Identity()
