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, 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 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 __init__(self, in_channels, out_channels, config, D=3, **kwargs):
super(FCNet, self).__init__(in_channels, out_channels, config, D)
reps = 2
# Output of the first conv concated to conv6
self.conv1p1s1 = ME.MinkowskiConvolution(in_channels=in_channels,
out_channels=self.PLANES[0],
pixel_dist=1,
kernel_size=3,
has_bias=False,
dimension=D)
self.bn1 = nn.BatchNorm1d(self.PLANES[0])
self.block1 = FCNBlocks(self.PLANES[0], pixel_dist=1, reps=reps, D=D)
self.conv2p1s2 = ME.MinkowskiConvolution(in_channels=self.PLANES[0],
out_channels=self.PLANES[1],
pixel_dist=1,
kernel_size=2,
stride=2,
has_bias=False,
dimension=D)
self.bn2 = nn.BatchNorm1d(self.PLANES[1])
self.block2 = FCNBlocks(self.PLANES[1], pixel_dist=2, reps=reps, D=D)
self.unpool2 = ME.MinkowskiPoolingTranspose(pixel_dist=2,
kernel_size=2,
stride=2,
dimension=D)
self.conv3p2s2 = ME.MinkowskiConvolution(in_channels=self.PLANES[1],
out_channels=self.PLANES[2],
pixel_dist=2,
kernel_size=2,
stride=2,
has_bias=False,
dimension=D)
self.bn3 = nn.BatchNorm1d(self.PLANES[2])
self.block3 = FCNBlocks(self.PLANES[2], pixel_dist=4, reps=reps, D=D)
self.unpool3 = ME.MinkowskiPoolingTranspose(pixel_dist=4,
kernel_size=4,
stride=4,
dimension=D)
self.conv4p4s2 = ME.MinkowskiConvolution(in_channels=self.PLANES[2],
out_channels=self.PLANES[3],
pixel_dist=4,
kernel_size=2,
stride=2,
has_bias=False,
dimension=D)
self.bn4 = nn.BatchNorm1d(self.PLANES[3])
self.block4 = FCNBlocks(self.PLANES[3], pixel_dist=8, reps=reps, D=D)
self.unpool4 = ME.MinkowskiPoolingTranspose(pixel_dist=8,
kernel_size=8,
stride=8,
dimension=D)
self.relu = nn.ReLU(inplace=True)
self.final = ME.MinkowskiConvolution(in_channels=sum(self.PLANES[:4]),
out_channels=out_channels,
pixel_dist=1,
kernel_size=1,
stride=1,
has_bias=True,
dimension=D)
def __init__(self,
in_channels=3,
out_channels=32,
bn_momentum=0.1,
conv1_kernel_size=3,
normalize_feature=False,
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
DEPTHS = self.DEPTHS
REGION_TYPE = self.REGION_TYPE
self.normalize_feature = normalize_feature
self.conv1 = conv(
in_channels=in_channels,
out_channels=CHANNELS[1],
kernel_size=conv1_kernel_size,
stride=1,
dilation=1,
has_bias=False,
region_type=REGION_TYPE,
dimension=D)
self.norm1 = get_norm(NORM_TYPE, CHANNELS[1], bn_momentum=bn_momentum, dimension=D)
self.block1 = nn.Sequential(*[
get_block(
BLOCK_NORM_TYPE,
CHANNELS[1],
CHANNELS[1],
bn_momentum=bn_momentum,
region_type=REGION_TYPE,
dimension=D) for d in range(DEPTHS[1])
])
self.pool2 = ME.MinkowskiSumPooling(kernel_size=2, stride=2, dimension=D)
self.conv2 = conv(
in_channels=CHANNELS[1],
out_channels=CHANNELS[2],
kernel_size=1,
stride=1,
dilation=1,
has_bias=False,
dimension=D)
self.norm2 = get_norm(NORM_TYPE, CHANNELS[2], bn_momentum=bn_momentum, dimension=D)
self.block2 = nn.Sequential(*[
get_block(
BLOCK_NORM_TYPE,
CHANNELS[2],
CHANNELS[2],
bn_momentum=bn_momentum,
region_type=REGION_TYPE,
dimension=D) for d in range(DEPTHS[2])
])
self.pool3 = ME.MinkowskiSumPooling(kernel_size=2, stride=2, dimension=D)
self.conv3 = conv(
in_channels=CHANNELS[2],
out_channels=CHANNELS[3],
kernel_size=1,
stride=1,
dilation=1,
has_bias=False,
dimension=D)
self.norm3 = get_norm(NORM_TYPE, CHANNELS[3], bn_momentum=bn_momentum, dimension=D)
self.block3 = nn.Sequential(*[
get_block(
BLOCK_NORM_TYPE,
CHANNELS[3],
CHANNELS[3],
bn_momentum=bn_momentum,
region_type=REGION_TYPE,
dimension=D) for d in range(DEPTHS[3])
])
self.pool3_tr = ME.MinkowskiPoolingTranspose(kernel_size=2, stride=2, dimension=D)
self.conv3_tr = conv_tr(
in_channels=CHANNELS[3],
out_channels=TR_CHANNELS[3],
kernel_size=1,
stride=1,
dilation=1,
has_bias=False,
dimension=D)
self.norm3_tr = get_norm(
NORM_TYPE, TR_CHANNELS[3], bn_momentum=bn_momentum, dimension=D)
self.block3_tr = nn.Sequential(*[
get_block(
BLOCK_NORM_TYPE,
TR_CHANNELS[3],
TR_CHANNELS[3],
bn_momentum=bn_momentum,
region_type=REGION_TYPE,
dimension=D) for d in range(DEPTHS[-3])
])
self.pool2_tr = ME.MinkowskiPoolingTranspose(kernel_size=2, stride=2, dimension=D)
self.conv2_tr = conv_tr(
in_channels=CHANNELS[2] + TR_CHANNELS[3],
out_channels=TR_CHANNELS[2],
kernel_size=1,
stride=1,
dilation=1,
has_bias=False,
region_type=REGION_TYPE,
dimension=D)
self.norm2_tr = get_norm(
NORM_TYPE, TR_CHANNELS[2], bn_momentum=bn_momentum, dimension=D)
self.block2_tr = nn.Sequential(*[
get_block(
BLOCK_NORM_TYPE,
TR_CHANNELS[2],
TR_CHANNELS[2],
bn_momentum=bn_momentum,
region_type=REGION_TYPE,
dimension=D) for d in range(DEPTHS[-2])
])
self.conv1_tr = conv_tr(
in_channels=CHANNELS[1] + TR_CHANNELS[2],
out_channels=TR_CHANNELS[1],
kernel_size=1,
stride=1,
dilation=1,
has_bias=False,
region_type=REGION_TYPE,
dimension=D)
# self.block1_tr = BasicBlockBN(TR_CHANNELS[1], TR_CHANNELS[1], bn_momentum=bn_momentum, dimension=D)
self.final = conv(
in_channels=TR_CHANNELS[1],
out_channels=out_channels,
kernel_size=1,
stride=1,
dilation=1,
has_bias=True,
dimension=D)
def network_initialization(self, in_channels, out_channels, config, D):
# Setup net_metadata
dilations = self.DILATIONS
bn_momentum = config.bn_momentum
def space_n_time_m(n, m):
return n if D == 3 else [n, n, n, m]
if D == 4:
self.OUT_PIXEL_DIST = space_n_time_m(self.OUT_PIXEL_DIST, 1)
# Output of the first conv concated to conv6
self.inplanes = self.INIT_DIM
self.conv1p1s1 = conv(in_channels,
self.inplanes,
kernel_size=space_n_time_m(
config.conv1_kernel_size, 1),
stride=1,
dilation=1,
conv_type=self.NON_BLOCK_CONV_TYPE,
D=D)
self.bn1 = get_norm(self.NORM_TYPE,
self.PLANES[0],
D,
bn_momentum=bn_momentum)
self.block1 = self._make_layer(self.BLOCK,
self.PLANES[0],
self.LAYERS[0],
dilation=dilations[0],
norm_type=self.NORM_TYPE,
bn_momentum=bn_momentum)
self.conv2p1s2 = conv(self.inplanes,
self.inplanes,
kernel_size=space_n_time_m(2, 1),
stride=space_n_time_m(2, 1),
dilation=1,
conv_type=self.NON_BLOCK_CONV_TYPE,
D=D)
self.bn2 = get_norm(self.NORM_TYPE,
self.inplanes,
D,
bn_momentum=bn_momentum)
self.block2 = self._make_layer(self.BLOCK,
self.PLANES[1],
self.LAYERS[1],
dilation=dilations[1],
norm_type=self.NORM_TYPE,
bn_momentum=bn_momentum)
self.conv3p2s2 = conv(self.inplanes,
self.inplanes,
kernel_size=space_n_time_m(2, 1),
stride=space_n_time_m(2, 1),
dilation=1,
conv_type=self.NON_BLOCK_CONV_TYPE,
D=D)
self.bn3 = get_norm(self.NORM_TYPE,
self.inplanes,
D,
bn_momentum=bn_momentum)
self.block3 = self._make_layer(self.BLOCK,
self.PLANES[2],
self.LAYERS[2],
dilation=dilations[2],
norm_type=self.NORM_TYPE,
bn_momentum=bn_momentum)
self.conv4p4s2 = conv(self.inplanes,
self.inplanes,
kernel_size=space_n_time_m(2, 1),
stride=space_n_time_m(2, 1),
dilation=1,
conv_type=self.NON_BLOCK_CONV_TYPE,
D=D)
self.bn4 = get_norm(self.NORM_TYPE,
self.inplanes,
D,
bn_momentum=bn_momentum)
self.block4 = self._make_layer(self.BLOCK,
self.PLANES[3],
self.LAYERS[3],
dilation=dilations[3],
norm_type=self.NORM_TYPE,
bn_momentum=bn_momentum)
self.pool_tr4 = ME.MinkowskiPoolingTranspose(kernel_size=8,
stride=8,
dimension=D)
out_pool4 = self.inplanes
self.convtr4p8s2 = conv_tr(self.inplanes,
self.PLANES[4],
kernel_size=space_n_time_m(2, 1),
upsample_stride=space_n_time_m(2, 1),
dilation=1,
bias=False,
conv_type=self.NON_BLOCK_CONV_TYPE,
D=D)
self.bntr4 = get_norm(self.NORM_TYPE,
self.PLANES[4],
D,
bn_momentum=bn_momentum)
self.inplanes = self.PLANES[4] + self.PLANES[2] * self.BLOCK.expansion
self.block5 = self._make_layer(self.BLOCK,
self.PLANES[4],
self.LAYERS[4],
dilation=dilations[4],
norm_type=self.NORM_TYPE,
bn_momentum=bn_momentum)
self.pool_tr5 = ME.MinkowskiPoolingTranspose(kernel_size=4,
stride=4,
dimension=D)
out_pool5 = self.inplanes
self.convtr5p4s2 = conv_tr(self.inplanes,
self.PLANES[5],
kernel_size=space_n_time_m(2, 1),
upsample_stride=space_n_time_m(2, 1),
dilation=1,
bias=False,
conv_type=self.NON_BLOCK_CONV_TYPE,
D=D)
self.bntr5 = get_norm(self.NORM_TYPE,
self.PLANES[5],
D,
bn_momentum=bn_momentum)
self.inplanes = self.PLANES[5] + self.PLANES[1] * self.BLOCK.expansion
self.block6 = self._make_layer(self.BLOCK,
self.PLANES[5],
self.LAYERS[5],
dilation=dilations[5],
norm_type=self.NORM_TYPE,
bn_momentum=bn_momentum)
self.pool_tr6 = ME.MinkowskiPoolingTranspose(kernel_size=2,
stride=2,
dimension=D)
out_pool6 = self.inplanes
self.convtr6p2s2 = conv_tr(self.inplanes,
self.PLANES[6],
kernel_size=space_n_time_m(2, 1),
upsample_stride=space_n_time_m(2, 1),
dilation=1,
bias=False,
conv_type=self.NON_BLOCK_CONV_TYPE,
D=D)
self.bntr6 = get_norm(self.NORM_TYPE,
self.PLANES[6],
D,
bn_momentum=bn_momentum)
self.relu = MinkowskiReLU(inplace=True)
self.final = nn.Sequential(
conv(out_pool5 + out_pool6 + self.PLANES[6] +
self.PLANES[0] * self.BLOCK.expansion,
512,
kernel_size=1,
bias=False,
D=D), ME.MinkowskiBatchNorm(512), ME.MinkowskiReLU(),
conv(512, out_channels, kernel_size=1, bias=True, D=D))
