def __init__(self, in_nchannel, out_nchannel, D):
super(UNet, self).__init__(D)
self.block1 = torch.nn.Sequential(
ME.MinkowskiConvolution(
in_channels=in_nchannel,
out_channels=8,
kernel_size=3,
stride=1,
dimension=D),
ME.MinkowskiBatchNorm(8))
self.block2 = torch.nn.Sequential(
ME.MinkowskiConvolution(
in_channels=8,
out_channels=16,
kernel_size=3,
stride=2,
dimension=D),
ME.MinkowskiBatchNorm(16),
)
self.block3 = torch.nn.Sequential(
ME.MinkowskiConvolution(
in_channels=16,
out_channels=32,
kernel_size=3,
stride=2,
dimension=D),
ME.MinkowskiBatchNorm(32))
self.block3_tr = torch.nn.Sequential(
ME.MinkowskiConvolutionTranspose(
in_channels=32,
out_channels=16,
kernel_size=3,
stride=2,
dimension=D),
ME.MinkowskiBatchNorm(16))
self.block2_tr = torch.nn.Sequential(
ME.MinkowskiConvolutionTranspose(
in_channels=32,
out_channels=16,
kernel_size=3,
stride=2,
dimension=D),
ME.MinkowskiBatchNorm(16))
self.conv1_tr = ME.MinkowskiConvolution(
in_channels=24,
out_channels=out_nchannel,
kernel_size=1,
stride=1,
dimension=D)
def network_initialization(self, in_channels, config, D):
up_kernel_size = 3
self.conv_up1 = nn.Sequential(
ME.MinkowskiConvolutionTranspose(
in_channels[0], in_channels[0], kernel_size=up_kernel_size, stride=2,
generate_new_coords=True, dimension=3),
ME.MinkowskiBatchNorm(in_channels[0]),
ME.MinkowskiELU())
self.conv_up2 = nn.Sequential(
ME.MinkowskiConvolutionTranspose(
in_channels[1], in_channels[0], kernel_size=up_kernel_size, stride=2,
generate_new_coords=True, dimension=3),
ME.MinkowskiBatchNorm(in_channels[0]),
ME.MinkowskiELU())
self.conv_up3 = nn.Sequential(
ME.MinkowskiConvolutionTranspose(
in_channels[2], in_channels[1], kernel_size=up_kernel_size, stride=2,
generate_new_coords=True, dimension=3),
ME.MinkowskiBatchNorm(in_channels[1]),
ME.MinkowskiELU())
self.conv_up4 = nn.Sequential(
ME.MinkowskiConvolutionTranspose(
in_channels[3], in_channels[2], kernel_size=up_kernel_size, stride=2,
generate_new_coords=True, dimension=3),
ME.MinkowskiBatchNorm(in_channels[2]),
ME.MinkowskiELU())
self.conv_feat1 = nn.Sequential(
ME.MinkowskiConvolution(
in_channels[0], config.upsample_feat_size, kernel_size=1, dimension=3),
ME.MinkowskiBatchNorm(config.upsample_feat_size),
ME.MinkowskiELU())
self.conv_feat2 = nn.Sequential(
ME.MinkowskiConvolution(
in_channels[1], config.upsample_feat_size, kernel_size=1, dimension=3),
ME.MinkowskiBatchNorm(config.upsample_feat_size),
ME.MinkowskiELU())
self.conv_feat3 = nn.Sequential(
ME.MinkowskiConvolution(
in_channels[2], config.upsample_feat_size, kernel_size=1, dimension=3),
ME.MinkowskiBatchNorm(config.upsample_feat_size),
ME.MinkowskiELU())
self.conv_feat4 = nn.Sequential(
ME.MinkowskiConvolution(
in_channels[3], config.upsample_feat_size, kernel_size=1, dimension=3),
ME.MinkowskiBatchNorm(config.upsample_feat_size),
ME.MinkowskiELU())
def __init__(self, input_a_dim, input_b_dim, out_dim, kernel_size=2):
super().__init__()
'''
Deconv x_a
concat with x_b
then apply output-projection
'''
self.input_a_dim = input_a_dim
self.input_b_dim = input_b_dim
self.out_dim = out_dim
self.conv_a = nn.Sequential(
ME.MinkowskiConvolutionTranspose(in_channels=input_a_dim,
out_channels=input_a_dim,
kernel_size=4,
stride=4,
dimension=3),
ME.MinkowskiBatchNorm(input_a_dim),
ME.MinkowskiReLU(),
)
self.conv_proj = nn.Sequential(
ME.MinkowskiConvolution(in_channels=input_a_dim + input_b_dim,
out_channels=out_dim,
kernel_size=3,
stride=1,
dimension=3),
ME.MinkowskiBatchNorm(out_dim),
ME.MinkowskiReLU(),
)
def conv_tr(in_planes,
out_planes,
kernel_size,
upsample_stride=1,
dilation=1,
bias=False,
conv_type=ConvType.HYPERCUBE,
D=-1):
assert D > 0, 'Dimension must be a positive integer'
region_type, axis_types, kernel_size = convert_conv_type(conv_type, kernel_size, D)
kernel_generator = ME.KernelGenerator(
kernel_size,
upsample_stride,
dilation,
region_type=region_type,
axis_types=axis_types,
dimension=D)
return ME.MinkowskiConvolutionTranspose(
in_channels=in_planes,
out_channels=out_planes,
kernel_size=kernel_size,
stride=upsample_stride,
dilation=dilation,
has_bias=bias,
kernel_generator=kernel_generator,
dimension=D)
def __init__(self,
up_conv_nn=[],
kernel_size=3,
stride=1,
dilation=1,
has_bias=False,
activation=ME.MinkowskiReLU(inplace=True),
bn_momentum=0.01,
dimension=-1,
**kwargs):
"""
Block convolution which consists of a convolution a batch norm a
block operation and an activation.
the block operation is usually a resnetBlock
"""
# instantiate convolution
# instantiate batchnorm
# instantiate block
# activation
super(SimpleBlockUp, self).__init__()
self.conv_tr = ME.MinkowskiConvolutionTranspose(
up_conv_nn[0],
up_conv_nn[1],
kernel_size=kernel_size,
stride=stride,
dilation=dilation,
dimension=dimension)
self.bn = ME.MinkowskiBatchNorm(up_conv_nn[1], momentum=bn_momentum)
self.block = BasicBlock(up_conv_nn[1],
up_conv_nn[1],
bn_momentum=bn_momentum,
dimension=dimension)
self.activation = activation
def conv_tr(in_channels,
out_channels,
kernel_size,
stride=1,
dilation=1,
bias=False,
region_type=ME.RegionType.HYPER_CUBE,
dimension=-1):
assert dimension > 0, 'Dimension must be a positive integer'
kernel_generator = ME.KernelGenerator(kernel_size,
stride,
dilation,
is_transpose=True,
region_type=region_type,
dimension=dimension)
kernel_generator = ME.KernelGenerator(kernel_size,
stride,
dilation,
is_transpose=True,
region_type=region_type,
dimension=dimension)
return ME.MinkowskiConvolutionTranspose(in_channels=in_channels,
out_channels=out_channels,
kernel_size=kernel_size,
stride=stride,
dilation=dilation,
bias=bias,
kernel_generator=kernel_generator,
dimension=dimension)
def test_sum(self):
coords, colors, pcd = load_file("1.ply")
device = "cuda"
D = 3
batch_size = 16
voxel_size = 0.02
channels = [3, 64, 128]
dcoords = torch.from_numpy(np.floor(coords / voxel_size)).int()
bcoords = batched_coordinates([dcoords for i in range(batch_size)])
in_feats = torch.rand(len(bcoords), 3).to(0)
layer = MinkowskiStackSum(
ME.MinkowskiConvolution(
channels[0],
channels[1],
kernel_size=3,
stride=1,
dimension=3,
),
nn.Sequential(
ME.MinkowskiConvolution(
channels[0],
channels[1],
kernel_size=3,
stride=2,
dimension=3,
),
ME.MinkowskiStackSum(
nn.Identity(),
nn.Sequential(
ME.MinkowskiConvolution(
channels[1],
channels[2],
kernel_size=3,
stride=2,
dimension=3,
),
ME.MinkowskiConvolutionTranspose(
channels[2],
channels[1],
kernel_size=3,
stride=1,
dimension=3,
),
ME.MinkowskiPoolingTranspose(
kernel_size=2, stride=2, dimension=D
),
),
),
ME.MinkowskiPoolingTranspose(kernel_size=2, stride=2, dimension=D),
),
).cuda()
for i in range(1000):
torch.cuda.empty_cache()
sinput = ME.SparseTensor(in_feats, coordinates=bcoords, device=device)
layer(sinput)
def __init__(self, inc, outc, ks=3, stride=1, D=3):
super(BasicDeconvolutionBlock, self).__init__()
self.net = nn.Sequential(
ME.MinkowskiConvolutionTranspose(
inc, outc, kernel_size=ks, stride=stride, dimension=D),
ME.MinkowskiBatchNorm(outc),
ME.MinkowskiReLU(True))
nn.init.constant_(self.net[1].bn.weight, 1.0)
nn.init.constant_(self.net[1].bn.bias, 0.0)
def __init__(self, in_nchannel, out_nchannel, D):
super(UNet, self).__init__(D)
self.conv1 = ME.MinkowskiConvolution(
in_channels=in_nchannel,
out_channels=8,
kernel_size=3,
stride=1,
dimension=D)
self.bn1 = ME.MinkowskiBatchNorm(8)
self.conv2 = ME.MinkowskiConvolution(
in_channels=8,
out_channels=16,
kernel_size=3,
stride=2,
dimension=D)
self.bn2 = ME.MinkowskiBatchNorm(16)
self.conv3 = ME.MinkowskiConvolution(
in_channels=16,
out_channels=32,
kernel_size=3,
stride=2,
dimension=D)
self.bn3 = ME.MinkowskiBatchNorm(32)
self.conv4 = ME.MinkowskiConvolutionTranspose(
in_channels=32,
out_channels=16,
kernel_size=3,
stride=2,
dimension=D)
self.bn4 = ME.MinkowskiBatchNorm(16)
self.conv5 = ME.MinkowskiConvolutionTranspose(
in_channels=32,
out_channels=16,
kernel_size=3,
stride=2,
dimension=D)
self.bn5 = ME.MinkowskiBatchNorm(16)
self.conv6 = ME.MinkowskiConvolution(
in_channels=24,
out_channels=out_nchannel,
kernel_size=1,
stride=1,
dimension=D)
def __init__(self, inc, outc, ks=3, stride=1, D=3):
super(BasicDeconvolutionBlock4d, self).__init__()
self.net = nn.Sequential(
ME.MinkowskiConvolutionTranspose(
inc, outc, stride=stride, dimension=D,
kernel_generator=ME.KernelGenerator(
kernel_size=(ks, ks, ks, 1), dimension=D,
region_type=ME.RegionType.HYPERCUBE)),
ME.MinkowskiBatchNorm(outc),
ME.MinkowskiReLU(True))
nn.init.constant_(self.net[1].bn.weight, 1.0)
nn.init.constant_(self.net[1].bn.bias, 0.0)
def __init__(self, in_channels, out_channels, D=3):
nn.Module.__init__(self)
self.net = nn.Sequential(
ME.MinkowskiConvolution(in_channels, 32, 3, dimension=D),
ME.MinkowskiBatchNorm(32),
ME.MinkowskiReLU(),
ME.MinkowskiConvolution(32, 64, 3, stride=2, dimension=D),
ME.MinkowskiBatchNorm(64),
ME.MinkowskiReLU(),
ME.MinkowskiConvolutionTranspose(64, 32, 3, stride=2, dimension=D),
ME.MinkowskiBatchNorm(32),
ME.MinkowskiReLU(),
ME.MinkowskiConvolution(32, out_channels, kernel_size=1, dimension=D),
)
def __init__(self, in_nchannel, out_nchannel, D):
ME.MinkowskiNetwork.__init__(self, D)
channels = [in_nchannel, 16, 32]
self.net = nn.Sequential(
ME.MinkowskiStackSum(
ME.MinkowskiConvolution(
channels[0],
channels[1],
kernel_size=3,
stride=1,
dimension=D,
),
nn.Sequential(
ME.MinkowskiConvolution(
channels[0],
channels[1],
kernel_size=3,
stride=2,
dimension=D,
),
ME.MinkowskiStackSum(
nn.Identity(),
nn.Sequential(
ME.MinkowskiConvolution(
channels[1],
channels[2],
kernel_size=3,
stride=2,
dimension=D,
),
ME.MinkowskiConvolutionTranspose(
channels[2],
channels[1],
kernel_size=3,
stride=1,
dimension=D,
),
ME.MinkowskiPoolingTranspose(
kernel_size=2, stride=2, dimension=D
),
),
),
ME.MinkowskiPoolingTranspose(kernel_size=2, stride=2, dimension=D),
),
),
ME.MinkowskiToFeature(),
nn.Linear(channels[1], out_nchannel, bias=True),
)
def __init__(self, down_channels, skip_channels, out_channels, D=3):
super().__init__()
self.deconv = nn.Sequential(
ME.MinkowskiConvolutionTranspose(down_channels,
down_channels,
kernel_size=3,
stride=2,
dimension=D),
ME.MinkowskiReLU(inplace=True),
)
self.conv = nn.Sequential(
ME.MinkowskiConvolution(skip_channels + down_channels,
out_channels,
kernel_size=3,
dimension=D),
ME.MinkowskiReLU(inplace=True),
)
def conv_on_coords():
in_channels, out_channels, D = 2, 3, 2
coords, feats, labels = data_loader(in_channels, batch_size=1)
# Create input with tensor stride == 4
strided_coords4, tensor_stride4 = get_random_coords(tensor_stride=4)
strided_coords2, tensor_stride2 = get_random_coords(tensor_stride=2)
input = ME.SparseTensor(
feats=torch.rand(len(strided_coords4), in_channels), #
coords=strided_coords4,
tensor_stride=tensor_stride4)
cm = input.coords_man
# Convolution transpose and generate new coordinates
conv_tr = ME.MinkowskiConvolutionTranspose(in_channels,
out_channels,
kernel_size=3,
stride=2,
has_bias=False,
dimension=D)
pool_tr = ME.MinkowskiPoolingTranspose(in_channels,
out_channels,
kernel_size=2,
stride=2,
has_bias=False,
dimension=D)
# If the there is no coordinates defined for the tensor stride, it will create one
# tensor stride 4 -> conv_tr with stride 2 -> tensor stride 2
output1 = conv_tr(input)
# output1 = pool_tr(input)
# convolution on the specified coords
output2 = conv_tr(input, coords)
# output2 = pool_tr(input, coords)
# convolution on the specified coords with tensor stride == 2
coords_key = cm.create_coords_key(strided_coords2, tensor_stride=2)
output3 = conv_tr(input, coords_key)
# output3 = pool_tr(input, coords_key)
# convolution on the coordinates of a sparse tensor
output4 = conv_tr(input, output1)
def conv():
in_channels, out_channels, D = 2, 3, 2
bp()
coords, feats, labels = data_loader(in_channels, batch_size=1)
# Convolution
input = ME.SparseTensor(feats=feats, coords=coords)
conv = ME.MinkowskiConvolution(in_channels,
out_channels,
kernel_size=3,
stride=2,
has_bias=False,
dimension=D)
output = conv(input)
print('Input:')
print_sparse_tensor(input)
print('Output:')
print_sparse_tensor(output)
# Convolution transpose and generate new coordinates
strided_coords, tensor_stride = get_random_coords()
bp()
input = ME.SparseTensor(
feats=torch.rand(len(strided_coords), in_channels), #
coords=strided_coords,
tensor_stride=tensor_stride)
conv_tr = ME.MinkowskiConvolutionTranspose(in_channels,
out_channels,
kernel_size=3,
stride=2,
has_bias=False,
dimension=D)
output = conv_tr(input)
print('\nInput:')
print_sparse_tensor(input)
print('Convolution Transpose Output:')
print_sparse_tensor(output)
def __init__(self,
up_conv_nn,
kernel_size,
stride,
dilation,
dimension=3,
bn_momentum=0.01,
norm_type=NormType.BATCH_NORM,
block_norm_type=NormType.BATCH_NORM,
**kwargs):
ME.MinkowskiNetwork.__init__(self, dimension)
self.conv = ME.MinkowskiConvolutionTranspose(
in_channels=up_conv_nn[0],
out_channels=up_conv_nn[1],
kernel_size=kernel_size,
stride=stride,
dilation=dilation,
has_bias=False,
dimension=dimension,
)
if len(up_conv_nn) == 3:
self.final = ME.MinkowskiConvolution(
in_channels=up_conv_nn[1],
out_channels=up_conv_nn[2],
kernel_size=kernel_size,
stride=stride,
dilation=dilation,
has_bias=True,
dimension=dimension,
)
else:
self.norm = get_norm(norm_type,
up_conv_nn[1],
bn_momentum=bn_momentum,
D=dimension)
self.block = get_block(block_norm_type,
up_conv_nn[1],
up_conv_nn[1],
bn_momentum=bn_momentum,
D=dimension)
self.final = None
def __init__(self, input_dim, out_dim, k=3):
super().__init__()
'''
Transition Up Layer
npoint: number of input points
nsample: k in kNN, default 3
in_dim: feature dimension of the input feature x (output of the PCTLayer)
out_dim: feature dimension of the TDLayer
'''
self.k = k
self.input_dim = input_dim
self.out_dim = out_dim
self.linear_1 = nn.Conv1d(input_dim, out_dim, 1)
self.linear_2 = nn.Conv1d(out_dim, out_dim, 1)
self.conv = ME.MinkowskiConvolutionTranspose(in_channels=input_dim,
out_channels=out_dim,
kernel_size=3,
stride=2,
dimension=3)
def __init__(self,
up_conv_nn=[],
kernel_sizes=[],
strides=[],
dilations=[],
kernel_size=3,
stride=1,
dilation=1,
norm_layer=ME.MinkowskiBatchNorm,
activation=ME.MinkowskiReLU,
bn_momentum=0.1,
dimension=-1,
up_stride=2,
skip=True,
**kwargs):
self.skip = skip
super(ResnetBlockUp, self).__init__(
up_conv_nn[0],
up_conv_nn[1],
up_conv_nn[2],
kernel_sizes=kernel_sizes,
strides=strides,
dilations=dilations,
kernel_size=kernel_size,
stride=stride,
dilation=dilation,
norm_layer=norm_layer,
activation=activation,
bn_momentum=bn_momentum,
dimension=dimension,
)
self.upsample = ME.MinkowskiConvolutionTranspose(up_conv_nn[0],
up_conv_nn[2],
kernel_size=2,
stride=up_stride,
dimension=dimension)
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])
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])
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])
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])
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] + self.PLANES[2] * self.BLOCK.expansion
self.block5 = self._make_layer(self.BLOCK, self.PLANES[4],
self.LAYERS[4])
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
self.block6 = self._make_layer(self.BLOCK, self.PLANES[5],
self.LAYERS[5])
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
self.block7 = self._make_layer(self.BLOCK, self.PLANES[6],
self.LAYERS[6])
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
self.block8 = self._make_layer(self.BLOCK, self.PLANES[7],
self.LAYERS[7])
if ME.__version__.split('.')[1] < '5':
self.final = ME.MinkowskiConvolution(self.PLANES[7] *
self.BLOCK.expansion,
out_channels,
kernel_size=1,
has_bias=True,
dimension=D)
else:
self.final = ME.MinkowskiConvolution(self.PLANES[7] *
self.BLOCK.expansion,
out_channels,
kernel_size=1,
bias=True,
dimension=D)
self.relu = ME.MinkowskiReLU(inplace=True)
def __init__(self,
in_channels=3,
out_channels=32,
bn_momentum=0.1,
normalize_feature=None,
conv1_kernel_size=None,
D=3):
ME.MinkowskiNetwork.__init__(self, D)
NORM_TYPE = self.NORM_TYPE
BLOCK_NORM_TYPE = self.BLOCK_NORM_TYPE
CHANNELS = self.CHANNELS
TR_CHANNELS = self.TR_CHANNELS
self.normalize_feature = normalize_feature
self.conv1 = ME.MinkowskiConvolution(
in_channels=in_channels,
out_channels=CHANNELS[1],
kernel_size=conv1_kernel_size,
stride=1,
dilation=1,
has_bias=False,
dimension=D)
self.norm1 = get_norm(NORM_TYPE, CHANNELS[1], bn_momentum=bn_momentum, D=D)
self.block1 = get_block(
BLOCK_NORM_TYPE, CHANNELS[1], CHANNELS[1], bn_momentum=bn_momentum, D=D)
self.conv2 = ME.MinkowskiConvolution(
in_channels=CHANNELS[1],
out_channels=CHANNELS[2],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm2 = get_norm(NORM_TYPE, CHANNELS[2], bn_momentum=bn_momentum, D=D)
self.block2 = get_block(
BLOCK_NORM_TYPE, CHANNELS[2], CHANNELS[2], bn_momentum=bn_momentum, D=D)
self.conv3 = ME.MinkowskiConvolution(
in_channels=CHANNELS[2],
out_channels=CHANNELS[3],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm3 = get_norm(NORM_TYPE, CHANNELS[3], bn_momentum=bn_momentum, D=D)
self.block3 = get_block(
BLOCK_NORM_TYPE, CHANNELS[3], CHANNELS[3], bn_momentum=bn_momentum, D=D)
self.conv4 = ME.MinkowskiConvolution(
in_channels=CHANNELS[3],
out_channels=CHANNELS[4],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm4 = get_norm(NORM_TYPE, CHANNELS[4], bn_momentum=bn_momentum, D=D)
self.block4 = get_block(
BLOCK_NORM_TYPE, CHANNELS[4], CHANNELS[4], bn_momentum=bn_momentum, D=D)
self.conv4_tr = ME.MinkowskiConvolutionTranspose(
in_channels=CHANNELS[4],
out_channels=TR_CHANNELS[4],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm4_tr = get_norm(NORM_TYPE, TR_CHANNELS[4], bn_momentum=bn_momentum, D=D)
self.block4_tr = get_block(
BLOCK_NORM_TYPE, TR_CHANNELS[4], TR_CHANNELS[4], bn_momentum=bn_momentum, D=D)
self.conv3_tr = ME.MinkowskiConvolutionTranspose(
in_channels=CHANNELS[3] + TR_CHANNELS[4],
out_channels=TR_CHANNELS[3],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm3_tr = get_norm(NORM_TYPE, TR_CHANNELS[3], bn_momentum=bn_momentum, D=D)
self.block3_tr = get_block(
BLOCK_NORM_TYPE, TR_CHANNELS[3], TR_CHANNELS[3], bn_momentum=bn_momentum, D=D)
self.conv2_tr = ME.MinkowskiConvolutionTranspose(
in_channels=CHANNELS[2] + TR_CHANNELS[3],
out_channels=TR_CHANNELS[2],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm2_tr = get_norm(NORM_TYPE, TR_CHANNELS[2], bn_momentum=bn_momentum, D=D)
self.block2_tr = get_block(
BLOCK_NORM_TYPE, TR_CHANNELS[2], TR_CHANNELS[2], bn_momentum=bn_momentum, D=D)
self.conv1_tr = ME.MinkowskiConvolution(
in_channels=CHANNELS[1] + TR_CHANNELS[2],
out_channels=TR_CHANNELS[1],
kernel_size=1,
stride=1,
dilation=1,
has_bias=False,
dimension=D)
# self.block1_tr = BasicBlockBN(TR_CHANNELS[1], TR_CHANNELS[1], bn_momentum=bn_momentum, D=D)
self.final = ME.MinkowskiConvolution(
in_channels=TR_CHANNELS[1],
out_channels=out_channels,
kernel_size=1,
stride=1,
dilation=1,
has_bias=True,
dimension=D)
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.MinkowskiConvolutionTranspose(ch[0],
ch[0],
kernel_size=2,
stride=2,
generate_new_coords=True,
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(),
ME.MinkowskiConvolutionTranspose(ch[0],
ch[1],
kernel_size=2,
stride=2,
generate_new_coords=True,
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.block1_cls = ME.MinkowskiConvolution(ch[1],
1,
kernel_size=1,
has_bias=True,
dimension=3)
# Block 2
self.block2 = nn.Sequential(
ME.MinkowskiConvolutionTranspose(ch[1],
ch[2],
kernel_size=2,
stride=2,
generate_new_coords=True,
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.block2_cls = ME.MinkowskiConvolution(ch[2],
1,
kernel_size=1,
has_bias=True,
dimension=3)
# Block 3
self.block3 = nn.Sequential(
ME.MinkowskiConvolutionTranspose(ch[2],
ch[3],
kernel_size=2,
stride=2,
generate_new_coords=True,
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.block3_cls = ME.MinkowskiConvolution(ch[3],
1,
kernel_size=1,
has_bias=True,
dimension=3)
# Block 4
self.block4 = nn.Sequential(
ME.MinkowskiConvolutionTranspose(ch[3],
ch[4],
kernel_size=2,
stride=2,
generate_new_coords=True,
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.block4_cls = ME.MinkowskiConvolution(ch[4],
1,
kernel_size=1,
has_bias=True,
dimension=3)
# Block 5
self.block5 = nn.Sequential(
ME.MinkowskiConvolutionTranspose(ch[4],
ch[5],
kernel_size=2,
stride=2,
generate_new_coords=True,
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.block5_cls = ME.MinkowskiConvolution(ch[5],
1,
kernel_size=1,
has_bias=True,
dimension=3)
# Block 6
self.block6 = nn.Sequential(
ME.MinkowskiConvolutionTranspose(ch[5],
ch[6],
kernel_size=2,
stride=2,
generate_new_coords=True,
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.block6_cls = ME.MinkowskiConvolution(ch[6],
1,
kernel_size=1,
has_bias=True,
dimension=3)
# pruning
self.pruning = ME.MinkowskiPruning()
def __init__(self, resolution, in_nchannel=512):
nn.Module.__init__(self)
self.resolution = resolution
# Input sparse tensor must have tensor stride 128.
enc_ch = self.ENC_CHANNELS
dec_ch = self.DEC_CHANNELS
# Encoder
self.enc_block_s1 = nn.Sequential(
ME.MinkowskiConvolution(1,
enc_ch[0],
kernel_size=3,
stride=1,
dimension=3),
ME.MinkowskiBatchNorm(enc_ch[0]),
ME.MinkowskiELU(),
)
self.enc_block_s1s2 = nn.Sequential(
ME.MinkowskiConvolution(enc_ch[0],
enc_ch[1],
kernel_size=2,
stride=2,
dimension=3),
ME.MinkowskiBatchNorm(enc_ch[1]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(enc_ch[1],
enc_ch[1],
kernel_size=3,
dimension=3),
ME.MinkowskiBatchNorm(enc_ch[1]),
ME.MinkowskiELU(),
)
self.enc_block_s2s4 = nn.Sequential(
ME.MinkowskiConvolution(enc_ch[1],
enc_ch[2],
kernel_size=2,
stride=2,
dimension=3),
ME.MinkowskiBatchNorm(enc_ch[2]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(enc_ch[2],
enc_ch[2],
kernel_size=3,
dimension=3),
ME.MinkowskiBatchNorm(enc_ch[2]),
ME.MinkowskiELU(),
)
self.enc_block_s4s8 = nn.Sequential(
ME.MinkowskiConvolution(enc_ch[2],
enc_ch[3],
kernel_size=2,
stride=2,
dimension=3),
ME.MinkowskiBatchNorm(enc_ch[3]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(enc_ch[3],
enc_ch[3],
kernel_size=3,
dimension=3),
ME.MinkowskiBatchNorm(enc_ch[3]),
ME.MinkowskiELU(),
)
self.enc_block_s8s16 = nn.Sequential(
ME.MinkowskiConvolution(enc_ch[3],
enc_ch[4],
kernel_size=2,
stride=2,
dimension=3),
ME.MinkowskiBatchNorm(enc_ch[4]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(enc_ch[4],
enc_ch[4],
kernel_size=3,
dimension=3),
ME.MinkowskiBatchNorm(enc_ch[4]),
ME.MinkowskiELU(),
)
self.enc_block_s16s32 = nn.Sequential(
ME.MinkowskiConvolution(enc_ch[4],
enc_ch[5],
kernel_size=2,
stride=2,
dimension=3),
ME.MinkowskiBatchNorm(enc_ch[5]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(enc_ch[5],
enc_ch[5],
kernel_size=3,
dimension=3),
ME.MinkowskiBatchNorm(enc_ch[5]),
ME.MinkowskiELU(),
)
self.enc_block_s32s64 = nn.Sequential(
ME.MinkowskiConvolution(enc_ch[5],
enc_ch[6],
kernel_size=2,
stride=2,
dimension=3),
ME.MinkowskiBatchNorm(enc_ch[6]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(enc_ch[6],
enc_ch[6],
kernel_size=3,
dimension=3),
ME.MinkowskiBatchNorm(enc_ch[6]),
ME.MinkowskiELU(),
)
# Decoder
self.dec_block_s64s32 = nn.Sequential(
ME.MinkowskiConvolutionTranspose(enc_ch[6],
dec_ch[5],
kernel_size=4,
stride=2,
generate_new_coords=True,
dimension=3),
ME.MinkowskiBatchNorm(dec_ch[5]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(dec_ch[5],
dec_ch[5],
kernel_size=3,
dimension=3),
ME.MinkowskiBatchNorm(dec_ch[5]),
ME.MinkowskiELU(),
)
self.dec_s32_cls = ME.MinkowskiConvolution(dec_ch[5],
1,
kernel_size=1,
has_bias=True,
dimension=3)
self.dec_block_s32s16 = nn.Sequential(
ME.MinkowskiConvolutionTranspose(enc_ch[5],
dec_ch[4],
kernel_size=2,
stride=2,
generate_new_coords=True,
dimension=3),
ME.MinkowskiBatchNorm(dec_ch[4]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(dec_ch[4],
dec_ch[4],
kernel_size=3,
dimension=3),
ME.MinkowskiBatchNorm(dec_ch[4]),
ME.MinkowskiELU(),
)
self.dec_s16_cls = ME.MinkowskiConvolution(dec_ch[4],
1,
kernel_size=1,
has_bias=True,
dimension=3)
self.dec_block_s16s8 = nn.Sequential(
ME.MinkowskiConvolutionTranspose(dec_ch[4],
dec_ch[3],
kernel_size=2,
stride=2,
generate_new_coords=True,
dimension=3),
ME.MinkowskiBatchNorm(dec_ch[3]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(dec_ch[3],
dec_ch[3],
kernel_size=3,
dimension=3),
ME.MinkowskiBatchNorm(dec_ch[3]),
ME.MinkowskiELU(),
)
self.dec_s8_cls = ME.MinkowskiConvolution(dec_ch[3],
1,
kernel_size=1,
has_bias=True,
dimension=3)
self.dec_block_s8s4 = nn.Sequential(
ME.MinkowskiConvolutionTranspose(dec_ch[3],
dec_ch[2],
kernel_size=2,
stride=2,
generate_new_coords=True,
dimension=3),
ME.MinkowskiBatchNorm(dec_ch[2]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(dec_ch[2],
dec_ch[2],
kernel_size=3,
dimension=3),
ME.MinkowskiBatchNorm(dec_ch[2]),
ME.MinkowskiELU(),
)
self.dec_s4_cls = ME.MinkowskiConvolution(dec_ch[2],
1,
kernel_size=1,
has_bias=True,
dimension=3)
self.dec_block_s4s2 = nn.Sequential(
ME.MinkowskiConvolutionTranspose(dec_ch[2],
dec_ch[1],
kernel_size=2,
stride=2,
generate_new_coords=True,
dimension=3),
ME.MinkowskiBatchNorm(dec_ch[1]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(dec_ch[1],
dec_ch[1],
kernel_size=3,
dimension=3),
ME.MinkowskiBatchNorm(dec_ch[1]),
ME.MinkowskiELU(),
)
self.dec_s2_cls = ME.MinkowskiConvolution(dec_ch[1],
1,
kernel_size=1,
has_bias=True,
dimension=3)
self.dec_block_s2s1 = nn.Sequential(
ME.MinkowskiConvolutionTranspose(dec_ch[1],
dec_ch[0],
kernel_size=2,
stride=2,
generate_new_coords=True,
dimension=3),
ME.MinkowskiBatchNorm(dec_ch[0]),
ME.MinkowskiELU(),
ME.MinkowskiConvolution(dec_ch[0],
dec_ch[0],
kernel_size=3,
dimension=3),
ME.MinkowskiBatchNorm(dec_ch[0]),
ME.MinkowskiELU(),
)
self.dec_s1_cls = ME.MinkowskiConvolution(dec_ch[0],
1,
kernel_size=1,
has_bias=True,
dimension=3)
# pruning
self.pruning = ME.MinkowskiPruning()
def __init__(self,
in_channels=3,
out_channels=32,
bn_momentum=0.1,
normalize_feature=None,
conv1_kernel_size=None,
D=3):
super(SimpleNet, self).__init__(D)
NORM_TYPE = self.NORM_TYPE
CHANNELS = self.CHANNELS
TR_CHANNELS = self.TR_CHANNELS
self.normalize_feature = normalize_feature
self.conv1 = ME.MinkowskiConvolution(
in_channels=in_channels,
out_channels=CHANNELS[1],
kernel_size=conv1_kernel_size,
stride=1,
dilation=1,
has_bias=False,
dimension=D)
self.norm1 = get_norm(NORM_TYPE, CHANNELS[1], bn_momentum=bn_momentum, D=D)
self.conv2 = ME.MinkowskiConvolution(
in_channels=CHANNELS[1],
out_channels=CHANNELS[2],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm2 = get_norm(NORM_TYPE, CHANNELS[2], bn_momentum=bn_momentum, D=D)
self.conv3 = ME.MinkowskiConvolution(
in_channels=CHANNELS[2],
out_channels=CHANNELS[3],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm3 = get_norm(NORM_TYPE, CHANNELS[3], bn_momentum=bn_momentum, D=D)
self.conv3_tr = ME.MinkowskiConvolutionTranspose(
in_channels=CHANNELS[3],
out_channels=TR_CHANNELS[3],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm3_tr = get_norm(NORM_TYPE, TR_CHANNELS[3], bn_momentum=bn_momentum, D=D)
self.conv2_tr = ME.MinkowskiConvolutionTranspose(
in_channels=CHANNELS[2] + TR_CHANNELS[3],
out_channels=TR_CHANNELS[2],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm2_tr = get_norm(NORM_TYPE, TR_CHANNELS[2], bn_momentum=bn_momentum, D=D)
self.conv1_tr = ME.MinkowskiConvolution(
in_channels=CHANNELS[1] + TR_CHANNELS[2],
out_channels=TR_CHANNELS[1],
kernel_size=3,
stride=1,
dilation=1,
has_bias=False,
dimension=D)
self.norm1_tr = get_norm(NORM_TYPE, TR_CHANNELS[1], bn_momentum=bn_momentum, D=D)
self.final = ME.MinkowskiConvolution(
in_channels=TR_CHANNELS[1],
out_channels=out_channels,
kernel_size=1,
stride=1,
dilation=1,
has_bias=True,
dimension=D)
def __init__(self, input_a_dim, input_b_dim, out_dim, k=3):
super().__init__()
'''
Transition Up Layer
npoint: number of input points
nsample: k in kNN, default 3
input_a_dim: feature dimension of the input a(needs upsampling)
input_b_dim: feature dimension of the input b (directly concat)
out_dim: feature dimension of the TDLayer(fixed as the input_a_dim // 2) + input_b_dim
'''
self.k = k
self.input_a_dim = input_a_dim
self.input_b_dim = input_b_dim
self.intermediate_dim = (input_a_dim // 2) + input_b_dim
self.out_dim = out_dim
self.POINT_TR_LIKE = False
self.SUM_FEATURE = True # only used when POINTTR_LIKE is False, somehow have some peoblems
# -------- Point TR like -----------
if self.POINT_TR_LIKE:
self.linear_a = nn.Linear(input_a_dim, out_dim)
self.linear_b = nn.Linear(input_b_dim, out_dim)
else:
if self.SUM_FEATURE:
self.conv_a = nn.Sequential(
ME.MinkowskiConvolutionTranspose(in_channels=input_a_dim,
out_channels=out_dim,
kernel_size=2,
stride=2,
dimension=3),
ME.MinkowskiBatchNorm(out_dim),
ME.MinkowskiReLU(),
)
self.conv_b = nn.Sequential(
ME.MinkowskiConvolution(in_channels=input_b_dim,
out_channels=out_dim,
kernel_size=1,
stride=1,
dimension=3),
ME.MinkowskiBatchNorm(out_dim),
ME.MinkowskiReLU(),
)
else:
self.conv = ME.MinkowskiConvolutionTranspose(
in_channels=input_a_dim,
out_channels=input_a_dim // 2,
kernel_size=2,
stride=2,
dimension=3)
self.bn = ME.MinkowskiBatchNorm(self.input_a_dim // 2)
self.relu = ME.MinkowskiReLU()
# -----------------------------------------
self.out_conv = ME.MinkowskiConvolution(
in_channels=input_a_dim // 2 + input_b_dim,
out_channels=out_dim,
kernel_size=3,
stride=1,
dimension=3)
self.out_bn = ME.MinkowskiBatchNorm(self.out_dim)
self.out_relu = ME.MinkowskiReLU()
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)
def network_initialization(self, in_channels, out_channels, D):
assert len(self.layers) == len(self.planes)
assert len(self.planes) == self.num_bottom_up
self.convs = nn.ModuleList(
) # Bottom-up convolutional blocks with stride=2
self.bn = nn.ModuleList() # Bottom-up BatchNorms
self.blocks = nn.ModuleList() # Bottom-up blocks
self.tconvs = nn.ModuleList() # Top-down tranposed convolutions
self.conv1x1 = nn.ModuleList(
) # 1x1 convolutions in lateral connections
# The first convolution is special case, with kernel size = 5
self.inplanes = self.planes[0]
self.conv0 = ME.MinkowskiConvolution(
in_channels,
self.inplanes,
kernel_size=self.conv0_kernel_size,
dimension=D)
self.bn0 = ME.MinkowskiBatchNorm(self.inplanes)
for plane, layer in zip(self.planes, self.layers):
self.convs.append(
ME.MinkowskiConvolution(self.inplanes,
self.inplanes,
kernel_size=2,
stride=2,
dimension=D))
self.bn.append(ME.MinkowskiBatchNorm(self.inplanes))
self.blocks.append(self._make_layer(self.block, plane, layer))
# Lateral connections
for i in range(self.num_top_down):
self.conv1x1.append(
ME.MinkowskiConvolution(self.planes[-1 - i],
self.lateral_dim,
kernel_size=1,
stride=1,
dimension=D))
self.tconvs.append(
ME.MinkowskiConvolutionTranspose(self.lateral_dim,
self.lateral_dim,
kernel_size=2,
stride=2,
dimension=D))
# There's one more lateral connection than top-down TConv blocks
if self.num_top_down < self.num_bottom_up:
# Lateral connection from Conv block 1 or above
self.conv1x1.append(
ME.MinkowskiConvolution(self.planes[-1 - self.num_top_down],
self.lateral_dim,
kernel_size=1,
stride=1,
dimension=D))
else:
# Lateral connection from Con0 block
self.conv1x1.append(
ME.MinkowskiConvolution(self.planes[0],
self.lateral_dim,
kernel_size=1,
stride=1,
dimension=D))
self.relu = ME.MinkowskiReLU(inplace=True)
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,config,D=3): ME.MinkowskiNetwork.__init__(self, D) NORM_TYPE = self.NORM_TYPE BLOCK_NORM_TYPE = self.BLOCK_NORM_TYPE CHANNELS = self.CHANNELS TR_CHANNELS = self.TR_CHANNELS bn_momentum = config.bn_momentum self.normalize_feature = config.normalize_feature self.voxel_size = config.voxel_size self.conv1 = ME.MinkowskiConvolution( in_channels=config.in_feats_dim, out_channels=CHANNELS[1], kernel_size=config.conv1_kernel_size, stride=1, dilation=1, bias=False, dimension=D) self.norm1 = get_norm(NORM_TYPE, CHANNELS[1], bn_momentum=bn_momentum, D=D) self.block1 = get_block( BLOCK_NORM_TYPE, CHANNELS[1], CHANNELS[1], bn_momentum=bn_momentum, D=D) self.conv2 = ME.MinkowskiConvolution( in_channels=CHANNELS[1], out_channels=CHANNELS[2], kernel_size=3, stride=2, dilation=1, bias=False, dimension=D) self.norm2 = get_norm(NORM_TYPE, CHANNELS[2], bn_momentum=bn_momentum, D=D) self.block2 = get_block( BLOCK_NORM_TYPE, CHANNELS[2], CHANNELS[2], bn_momentum=bn_momentum, D=D) self.conv3 = ME.MinkowskiConvolution( in_channels=CHANNELS[2], out_channels=CHANNELS[3], kernel_size=3, stride=2, dilation=1, bias=False, dimension=D) self.norm3 = get_norm(NORM_TYPE, CHANNELS[3], bn_momentum=bn_momentum, D=D) self.block3 = get_block( BLOCK_NORM_TYPE, CHANNELS[3], CHANNELS[3], bn_momentum=bn_momentum, D=D) self.conv4 = ME.MinkowskiConvolution( in_channels=CHANNELS[3], out_channels=CHANNELS[4], kernel_size=3, stride=2, dilation=1, bias=False, dimension=D) self.norm4 = get_norm(NORM_TYPE, CHANNELS[4], bn_momentum=bn_momentum, D=D) self.block4 = get_block( BLOCK_NORM_TYPE, CHANNELS[4], CHANNELS[4], bn_momentum=bn_momentum, D=D) # adapt input tensor here self.conv4_tr = ME.MinkowskiConvolutionTranspose( in_channels = config.gnn_feats_dim + 2, out_channels=TR_CHANNELS[4], kernel_size=3, stride=2, dilation=1, bias=False, dimension=D) self.norm4_tr = get_norm(NORM_TYPE, TR_CHANNELS[4], bn_momentum=bn_momentum, D=D) self.block4_tr = get_block( BLOCK_NORM_TYPE, TR_CHANNELS[4], TR_CHANNELS[4], bn_momentum=bn_momentum, D=D) self.conv3_tr = ME.MinkowskiConvolutionTranspose( in_channels=CHANNELS[3] + TR_CHANNELS[4], out_channels=TR_CHANNELS[3], kernel_size=3, stride=2, dilation=1, bias=False, dimension=D) self.norm3_tr = get_norm(NORM_TYPE, TR_CHANNELS[3], bn_momentum=bn_momentum, D=D) self.block3_tr = get_block( BLOCK_NORM_TYPE, TR_CHANNELS[3], TR_CHANNELS[3], bn_momentum=bn_momentum, D=D) self.conv2_tr = ME.MinkowskiConvolutionTranspose( in_channels=CHANNELS[2] + TR_CHANNELS[3], out_channels=TR_CHANNELS[2], kernel_size=3, stride=2, dilation=1, bias=False, dimension=D) self.norm2_tr = get_norm(NORM_TYPE, TR_CHANNELS[2], bn_momentum=bn_momentum, D=D) self.block2_tr = get_block( BLOCK_NORM_TYPE, TR_CHANNELS[2], TR_CHANNELS[2], bn_momentum=bn_momentum, D=D) self.conv1_tr = ME.MinkowskiConvolution( in_channels=CHANNELS[1] + TR_CHANNELS[2], out_channels=TR_CHANNELS[1], kernel_size=1, stride=1, dilation=1, bias=False, dimension=D) self.final = ME.MinkowskiConvolution( in_channels=TR_CHANNELS[1], out_channels=config.out_feats_dim + 2, kernel_size=1, stride=1, dilation=1, bias=True, dimension=D) ############# # Overlap attention module self.epsilon = torch.nn.Parameter(torch.tensor(-5.0)) self.bottle = nn.Conv1d(CHANNELS[4], config.gnn_feats_dim,kernel_size=1,bias=True) self.gnn = GCN(config.num_head,config.gnn_feats_dim, config.dgcnn_k, config.nets) self.proj_gnn = nn.Conv1d(config.gnn_feats_dim,config.gnn_feats_dim,kernel_size=1, bias=True) self.proj_score = nn.Conv1d(config.gnn_feats_dim,1,kernel_size=1,bias=True)
def __init__(self, in_channels=3, out_channels=32, bn_momentum=0.1, D=3, config=None):
ME.MinkowskiNetwork.__init__(self, D)
NORM_TYPE = self.NORM_TYPE
bn_momentum = config.bn_momentum
CHANNELS = self.CHANNELS
TR_CHANNELS = self.TR_CHANNELS
self.normalize_feature = config.normalize_feature
self.conv1 = ME.MinkowskiConvolution(
in_channels=in_channels,
out_channels=CHANNELS[1],
kernel_size=config.conv1_kernel_size,
stride=1,
dilation=1,
has_bias=False,
dimension=D)
self.norm1 = get_norm(NORM_TYPE, CHANNELS[1], bn_momentum=bn_momentum, D=D)
self.conv2 = ME.MinkowskiConvolution(
in_channels=CHANNELS[1],
out_channels=CHANNELS[2],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm2 = get_norm(NORM_TYPE, CHANNELS[2], bn_momentum=bn_momentum, D=D)
self.conv3 = ME.MinkowskiConvolution(
in_channels=CHANNELS[2],
out_channels=CHANNELS[3],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm3 = get_norm(NORM_TYPE, CHANNELS[3], bn_momentum=bn_momentum, D=D)
self.conv4 = ME.MinkowskiConvolution(
in_channels=CHANNELS[3],
out_channels=CHANNELS[4],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm4 = get_norm(NORM_TYPE, CHANNELS[4], bn_momentum=bn_momentum, D=D)
self.conv5 = ME.MinkowskiConvolution(
in_channels=CHANNELS[4],
out_channels=CHANNELS[5],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm5 = get_norm(NORM_TYPE, CHANNELS[5], bn_momentum=bn_momentum, D=D)
self.conv5_tr = ME.MinkowskiConvolutionTranspose(
in_channels=CHANNELS[5],
out_channels=TR_CHANNELS[5],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm5_tr = get_norm(NORM_TYPE, TR_CHANNELS[5], bn_momentum=bn_momentum, D=D)
self.conv4_tr = ME.MinkowskiConvolutionTranspose(
in_channels=CHANNELS[4] + TR_CHANNELS[5],
out_channels=TR_CHANNELS[4],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm4_tr = get_norm(NORM_TYPE, TR_CHANNELS[4], bn_momentum=bn_momentum, D=D)
self.conv3_tr = ME.MinkowskiConvolutionTranspose(
in_channels=CHANNELS[3] + TR_CHANNELS[4],
out_channels=TR_CHANNELS[3],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm3_tr = get_norm(NORM_TYPE, TR_CHANNELS[3], bn_momentum=bn_momentum, D=D)
self.conv2_tr = ME.MinkowskiConvolutionTranspose(
in_channels=CHANNELS[2] + TR_CHANNELS[3],
out_channels=TR_CHANNELS[2],
kernel_size=3,
stride=2,
dilation=1,
has_bias=False,
dimension=D)
self.norm2_tr = get_norm(NORM_TYPE, TR_CHANNELS[2], bn_momentum=bn_momentum, D=D)
self.conv1_tr = ME.MinkowskiConvolution(
in_channels=CHANNELS[1] + TR_CHANNELS[2],
out_channels=TR_CHANNELS[1],
kernel_size=1,
stride=1,
dilation=1,
has_bias=True,
dimension=D)
def __init__(self, cfg, name='uresnext'):
super(UResNeXt, self).__init__(cfg)
self.model_cfg = cfg['modules'][name]
# Configurations
self.reps = self.model_cfg.get('reps', 2)
self.depth = self.model_cfg.get('depth', 5)
self.num_filters = self.model_cfg.get('num_filters', 16)
self.cardinality = self.model_cfg.get('cardinality', 4)
assert (self.num_filters % self.cardinality == 0)
self.nPlanes = [i * self.num_filters for i in range(1, self.depth + 1)]
self.input_kernel = self.model_cfg.get('input_kernel', 3)
# Initialize Input Layer
self.input_layer = AtrousIIBlock(self.num_input,
self.num_filters,
leakiness=self.leakiness)
# Initialize Encoder
self.encoding_conv = []
self.encoding_block = []
for i, F in enumerate(self.nPlanes):
m = []
for _ in range(self.reps):
m.append(
ResNeXtBlock(F,
F,
dimension=self.D,
leakiness=self.leakiness,
cardinality=4,
dilations=[1, 1, 3, 9]))
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)
# 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(
ResNeXtBlock(self.nPlanes[i] * (2 if j == 0 else 1),
self.nPlanes[i],
dimension=self.D,
cardinality=4,
dilations=[1, 1, 3, 9],
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)))
