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 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 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, fc_dim=64, in_channels=3, out_channels=40, D=3):
super(SparseResNet18, self).__init__()
original_resnet = ResNet18(in_channels=in_channels,
out_channels=out_channels,
D=D)
checkpoint_path = os.path.join(
os.path.split(os.path.join(os.getcwd(), __file__))[0],
'resnet18_pretrained.pth')
#Load modelnet40 pretrained weights
checkpoint = torch.load(checkpoint_path)
original_resnet.load_state_dict(checkpoint['state_dict'])
self.features = nn.Sequential(*list(original_resnet.children())[:-4])
self.fc = ME.MinkowskiConvolution(in_channels=512,
out_channels=fc_dim,
kernel_size=3,
stride=1,
dimension=D)
self.glob_avg = ME.MinkowskiGlobalMaxPooling(dimension=D)
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):
super().__init__()
self.f = ME.MinkowskiGlobalMaxPooling()
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.MinkowskiGlobalMaxPooling()
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.conv0p1s1 = 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.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.relu = MinkowskiReLU(inplace=True)
# add a classification head here
self.clf_glob_avg = ME.MinkowskiGlobalPooling(dimension=D)
self.clf_glob_max=ME.MinkowskiGlobalMaxPooling(dimension=D)
self.clf_conv0 = conv(
256,
512,
kernel_size=3,
stride=2,
dilation=1,
conv_type=self.NON_BLOCK_CONV_TYPE,
D=D)
self.clf_bn0 = get_norm(self.NORM_TYPE, 512, D, bn_momentum=bn_momentum)
self.clf_conv1 = conv(
512,
512,
kernel_size=3,
stride=2,
dilation=1,
conv_type=self.NON_BLOCK_CONV_TYPE,
D=D)
self.clf_bn1 = get_norm(self.NORM_TYPE, 512, D, bn_momentum=bn_momentum)
self.clf_conv2 = conv(
512,
config['clf_num_labels'],
kernel_size=1,
stride=1,
dilation=1,
conv_type=self.NON_BLOCK_CONV_TYPE,
D=D)
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)
