def __init__(self, input_dim, p=3, eps=1e-6):
super(GeM, self).__init__()
self.input_dim = input_dim
# Same output number of channels as input number of channels
self.output_dim = self.input_dim
self.p = nn.Parameter(torch.ones(1) * p)
self.eps = eps
self.f = ME.MinkowskiGlobalAvgPooling()
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, p=3, eps=1e-6):
super(GeM, self).__init__()
self.p = nn.Parameter(torch.ones(1) * p)
self.eps = eps
self.f = ME.MinkowskiGlobalAvgPooling()
def __init__(self):
super().__init__()
self.f = ME.MinkowskiGlobalAvgPooling()
def __init__(self, input_dim):
super().__init__()
self.input_dim = input_dim
# Same output number of channels as input number of channels
self.output_dim = self.input_dim
self.f = ME.MinkowskiGlobalAvgPooling()
# Normalize features and create a sparse tensor
return features, coordinates
if __name__ == '__main__':
config = parser.parse_args()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# Define a model and load the weights
feats = [3, 8, 16, 32, 64, 128]
features, coordinates = generate_input_sparse_tensor(
config.file_name,
voxel_size=config.voxel_size,
batch_size=config.batch_size)
pool = ME.MinkowskiGlobalAvgPooling()
# Measure time
print('Forward')
for feat in feats:
timer = Timer()
features = torch.rand(len(coordinates), feat).to(device)
# Feed-forward pass and get the prediction
for i in range(20):
sinput = ME.SparseTensor(features,
coordinates=coordinates,
device=device)
timer.tic()
soutput = pool(sinput)
def network_initialization(self, in_channels, out_channels, D):
# Setup net_metadata
dilations = self.DILATIONS
bn_momentum = 0.02
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
conv1_kernel_size = 3
self.inplanes = self.INIT_DIM
self.conv0p1s1 = conv(in_channels,
self.inplanes,
kernel_size=space_n_time_m(conv1_kernel_size, 1),
stride=1,
dilation=1,
conv_type=self.NON_BLOCK_CONV_TYPE,
D=D)
self.bn0 = get_norm(self.NORM_TYPE,
self.inplanes,
D,
bn_momentum=bn_momentum)
self.conv1p1s2 = 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.bn1 = get_norm(self.NORM_TYPE,
self.inplanes,
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.conv2p2s2 = 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.conv3p4s2 = 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.conv4p8s2 = 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.convtr4p16s2 = 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.convtr5p8s2 = 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.convtr6p4s2 = 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.inplanes = self.PLANES[6] + self.PLANES[0] * self.BLOCK.expansion
self.block7 = self._make_layer(self.BLOCK,
self.PLANES[6],
self.LAYERS[6],
dilation=dilations[6],
norm_type=self.NORM_TYPE,
bn_momentum=bn_momentum)
self.convtr7p2s2 = conv_tr(self.inplanes,
self.PLANES[7],
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.bntr7 = get_norm(self.NORM_TYPE,
self.PLANES[7],
D,
bn_momentum=bn_momentum)
self.inplanes = self.PLANES[7] + self.INIT_DIM
self.block8 = self._make_layer(self.BLOCK,
self.PLANES[7],
self.LAYERS[7],
dilation=dilations[7],
norm_type=self.NORM_TYPE,
bn_momentum=bn_momentum)
self.all_feat_names = [
"en0",
"en1",
"en2",
"en3",
"en4",
"plane4",
"plane5",
"plane6",
"plane7",
]
self.relu = MinkowskiReLU(inplace=True)
self.maxpool = ME.MinkowskiGlobalMaxPooling()
self.avgpool = ME.MinkowskiGlobalAvgPooling()
if self.use_mlp:
self.head = SMLP(self.mlp_dim)
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()
