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.MinkowskiGlobalPooling(dimension=D)
self.final = ME.MinkowskiLinear(self.inplanes, out_channels, bias=True)
def __init__(self):
super(PointNetFeature, self).__init__()
k = self.KERNEL_SIZES
s = self.STRIDES
c = self.CONV_CHANNELS
self.stn = STN3d(D=3)
self.block1 = nn.Sequential(
ME.MinkowskiConvolution(6,
c[0],
kernel_size=k[0],
stride=s[0],
has_bias=False,
dimension=3),
ME.MinkowskiInstanceNorm(c[0]), ME.MinkowskiReLU())
self.block2 = nn.Sequential(
ME.MinkowskiConvolution(c[0],
c[1],
kernel_size=k[1],
stride=s[1],
has_bias=False,
dimension=3),
ME.MinkowskiInstanceNorm(c[1]), ME.MinkowskiReLU())
self.block3 = nn.Sequential(
ME.MinkowskiConvolution(c[1],
c[2],
kernel_size=k[2],
stride=s[2],
has_bias=False,
dimension=3),
ME.MinkowskiInstanceNorm(c[2]), ME.MinkowskiReLU())
self.avgpool = ME.MinkowskiGlobalPooling()
self.concat = ME.MinkowskiBroadcastConcatenation()
def __init__(self):
super(PointNetFeature, self).__init__()
k = self.KERNEL_SIZES
s = self.STRIDES
c = self.CONV_CHANNELS
self.stn = STN3d(D=3)
self.conv1 = ME.MinkowskiConvolution(6,
c[0],
kernel_size=k[0],
stride=s[0],
has_bias=False,
dimension=3)
self.conv2 = ME.MinkowskiConvolution(c[0],
c[1],
kernel_size=k[1],
stride=s[1],
has_bias=False,
dimension=3)
self.conv3 = ME.MinkowskiConvolution(c[1],
c[2],
kernel_size=k[2],
stride=s[2],
has_bias=False,
dimension=3)
self.bn1 = ME.MinkowskiInstanceNorm(c[0], dimension=3)
self.bn2 = ME.MinkowskiInstanceNorm(c[1], dimension=3)
self.bn3 = ME.MinkowskiInstanceNorm(c[2], dimension=3)
self.relu = ME.MinkowskiReLU(inplace=True)
self.avgpool = ME.MinkowskiGlobalPooling()
self.concat = ME.MinkowskiBroadcastConcatenation()
def network_initialization(
self, in_channel, out_channel, channels, embedding_channel, kernel_size, D=2,
):
self.conv1_1 = self.get_conv_block(in_channel, channels[0], kernel_size=kernel_size, stride=1)
self.conv1_2 = self.get_conv_block(channels[0], channels[1], kernel_size=kernel_size, stride=1)
self.pool1 = ME.MinkowskiMaxPooling(kernel_size=3, stride=2, dimension=D)
self.conv2_1 = self.get_conv_block(channels[1], channels[2], kernel_size=kernel_size, stride=1)
self.conv2_2 = self.get_conv_block(channels[2], channels[3], kernel_size=kernel_size, stride=1)
self.pool2 = ME.MinkowskiMaxPooling(kernel_size=3, stride=2, dimension=D)
self.conv3_1 = self.get_conv_block(channels[3], channels[4], kernel_size=kernel_size, stride=1)
self.conv3_2 = self.get_conv_block(channels[4], channels[5], kernel_size=kernel_size, stride=1)
self.conv3_3 = self.get_conv_block(channels[5], channels[6], kernel_size=kernel_size, stride=1)
self.pool3 = ME.MinkowskiMaxPooling(kernel_size=3, stride=2, dimension=D)
self.conv4_1 = self.get_conv_block(channels[6], channels[7], kernel_size=kernel_size, stride=1)
self.conv4_2 = self.get_conv_block(channels[7], channels[8], kernel_size=kernel_size, stride=1)
self.conv4_3 = self.get_conv_block(channels[8], channels[9], kernel_size=kernel_size, stride=1)
self.pool4 = ME.MinkowskiMaxPooling(kernel_size=3, stride=2, dimension=D)
self.conv5_1 = self.get_conv_block(channels[9], channels[10], kernel_size=kernel_size, stride=1)
self.conv5_2 = self.get_conv_block(channels[10], channels[11], kernel_size=kernel_size, stride=1)
self.conv5_3 = self.get_conv_block(channels[11], channels[12], kernel_size=kernel_size, stride=1)
self.pool5 = ME.MinkowskiMaxPooling(kernel_size=3, stride=2, dimension=D)
self.global_pool = ME.MinkowskiGlobalPooling()
self.final = nn.Sequential(
self.get_mlp_block(512, 512),
ME.MinkowskiDropout(),
self.get_mlp_block(512, 512),
self.get_mlp_block(512, 2),
#ME.MinkowskiFunctional.softmax(),
)
def network_initialization(self, in_channels, out_channels, config, D):
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)
dilations = config.dilations
bn_momentum = config.bn_momentum
self.inplanes = self.INIT_DIM
self.conv1 = conv(
in_channels,
self.inplanes,
kernel_size=space_n_time_m(config.conv1_kernel_size, 1),
stride=1,
D=D)
self.bn1 = get_norm(NormType.BATCH_NORM, self.inplanes, D=self.D, bn_momentum=bn_momentum)
self.relu = ME.MinkowskiReLU(inplace=True)
self.pool = sum_pool(kernel_size=space_n_time_m(2, 1), stride=space_n_time_m(2, 1), D=D)
self.layer1 = self._make_layer(
self.BLOCK,
self.PLANES[0],
self.LAYERS[0],
stride=space_n_time_m(2, 1),
dilation=space_n_time_m(dilations[0], 1))
self.layer2 = self._make_layer(
self.BLOCK,
self.PLANES[1],
self.LAYERS[1],
stride=space_n_time_m(2, 1),
dilation=space_n_time_m(dilations[1], 1))
self.layer3 = self._make_layer(
self.BLOCK,
self.PLANES[2],
self.LAYERS[2],
stride=space_n_time_m(2, 1),
dilation=space_n_time_m(dilations[2], 1))
self.layer4 = self._make_layer(
self.BLOCK,
self.PLANES[3],
self.LAYERS[3],
stride=space_n_time_m(2, 1),
dilation=space_n_time_m(dilations[3], 1))
if self.NETWORK_TYPE == NetworkType.CLASSIFICATION:
self.glob_avg = ME.MinkowskiGlobalPooling(dimension=D)
if self.HAS_LAST_BLOCK:
self.final1 = nn.Linear(self.inplanes, self.inplanes, bias=False)
self.bnfinal1 = nn.BatchNorm1d(self.inplanes)
self.final2 = nn.Linear(self.inplanes, self.inplanes, bias=False)
self.bnfinal2 = nn.BatchNorm1d(self.inplanes)
self.final = nn.Linear(self.inplanes, out_channels, bias=True)
else:
self.final = conv(
self.PLANES[3] * self.BLOCK.expansion, out_channels, kernel_size=1, bias=True, D=D)
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, D=3):
super(STN3d, self).__init__()
k = self.KERNEL_SIZES
s = self.STRIDES
c = self.CONV_CHANNELS
self.block1 = nn.Sequential(
ME.MinkowskiConvolution(3,
c[0],
kernel_size=k[0],
stride=s[0],
has_bias=False,
dimension=3),
ME.MinkowskiInstanceNorm(c[0]), ME.MinkowskiReLU())
self.block2 = nn.Sequential(
ME.MinkowskiConvolution(c[0],
c[1],
kernel_size=k[1],
stride=s[1],
has_bias=False,
dimension=3),
ME.MinkowskiInstanceNorm(c[1]), ME.MinkowskiReLU())
self.block3 = nn.Sequential(
ME.MinkowskiConvolution(c[1],
c[2],
kernel_size=k[2],
stride=s[2],
has_bias=False,
dimension=3),
ME.MinkowskiInstanceNorm(c[2]), ME.MinkowskiReLU())
# Use the kernelsize 1 convolution for linear layers. If kernel size ==
# 1, minkowski engine internally uses a linear function.
self.block4 = nn.Sequential(
ME.MinkowskiConvolution(c[2],
c[3],
kernel_size=1,
has_bias=False,
dimension=3),
ME.MinkowskiInstanceNorm(c[3]), ME.MinkowskiReLU())
self.block5 = nn.Sequential(
ME.MinkowskiConvolution(c[3],
c[4],
kernel_size=1,
has_bias=False,
dimension=3),
ME.MinkowskiInstanceNorm(c[4]), ME.MinkowskiReLU())
self.fc6 = ME.MinkowskiConvolution(c[4],
9,
kernel_size=1,
has_bias=True,
dimension=3)
self.avgpool = ME.MinkowskiGlobalPooling()
self.broadcast = ME.MinkowskiBroadcast()
def __init__(self, in_channels):
super(PointNetfeat, self).__init__()
self.pool = ME.MinkowskiGlobalPooling()
self.broadcast = ME.MinkowskiBroadcast()
self.stn = STN3d(D=3)
self.conv1 = nn.Conv1d(in_channels + 3, 128, 1, bias=False)
self.conv2 = nn.Conv1d(128, 256, 1, bias=False)
self.bn1 = nn.BatchNorm1d(128)
self.bn2 = nn.BatchNorm1d(256)
self.relu = nn.ReLU()
def __init__(self, D=3):
super(STN3d, self).__init__()
k = self.KERNEL_SIZES
s = self.STRIDES
c = self.CONV_CHANNELS
self.conv1 = ME.MinkowskiConvolution(3,
c[0],
kernel_size=k[0],
stride=s[0],
has_bias=False,
dimension=3)
self.conv2 = ME.MinkowskiConvolution(c[0],
c[1],
kernel_size=k[1],
stride=s[1],
has_bias=False,
dimension=3)
self.conv3 = ME.MinkowskiConvolution(c[1],
c[2],
kernel_size=k[2],
stride=s[2],
has_bias=False,
dimension=3)
# Use the kernelsize 1 convolution for linear layers. If kernel size ==
# 1, minkowski engine internally uses a linear function.
self.fc4 = ME.MinkowskiConvolution(c[2],
c[3],
kernel_size=1,
has_bias=False,
dimension=3)
self.fc5 = ME.MinkowskiConvolution(c[3],
c[4],
kernel_size=1,
has_bias=False,
dimension=3)
self.fc6 = ME.MinkowskiConvolution(c[4],
9,
kernel_size=1,
has_bias=True,
dimension=3)
self.relu = ME.MinkowskiReLU(inplace=True)
self.avgpool = ME.MinkowskiGlobalPooling()
self.broadcast = ME.MinkowskiBroadcast()
self.bn1 = ME.MinkowskiInstanceNorm(c[0], dimension=3)
self.bn2 = ME.MinkowskiInstanceNorm(c[1], dimension=3)
self.bn3 = ME.MinkowskiInstanceNorm(c[2], dimension=3)
self.bn4 = ME.MinkowskiInstanceNorm(c[3], dimension=3)
self.bn5 = ME.MinkowskiInstanceNorm(c[4], dimension=3)
def __init__(self, channels, gamma=2, b=1):
super().__init__()
t = int(abs((np.log2(channels) + b) / gamma))
k_size = t if t % 2 else t + 1
self.avg_pool = ME.MinkowskiGlobalPooling()
self.conv = nn.Conv1d(1,
1,
kernel_size=k_size,
padding=(k_size - 1) // 2,
bias=False)
self.sigmoid = nn.Sigmoid()
self.broadcast_mul = ME.MinkowskiBroadcastMultiplication()
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 __init__(self, D=3):
super(STN3d, self).__init__()
self.conv1 = nn.Conv1d(3, 64, 1, bias=False)
self.conv2 = nn.Conv1d(64, 128, 1, bias=False)
self.conv3 = nn.Conv1d(128, 256, 1, bias=False)
self.fc1 = nn.Linear(256, 128, bias=False)
self.fc2 = nn.Linear(128, 64, bias=False)
self.fc3 = nn.Linear(64, 9)
self.relu = nn.ReLU()
self.pool = ME.MinkowskiGlobalPooling()
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(256)
self.bn4 = nn.BatchNorm1d(128)
self.bn5 = nn.BatchNorm1d(64)
self.broadcast = ME.MinkowskiBroadcast()
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 __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, config={}, **kwargs):
MEEncoder.__init__(self, config=config, **kwargs)
# need square and power of 2 image size input
power = math.log(self.config.input_size[0], 2)
assert (power % 1 == 0.0) and (
power > 3
), "Dumoulin Encoder needs a power of 2 as image input size (>=16)"
# need square image input
assert torch.all(
torch.tensor([
self.config.input_size[i] == self.config.input_size[0]
for i in range(1, len(self.config.input_size))
])), "Dumoulin Encoder needs a square image input size"
assert self.config.n_conv_layers == power, "The number of convolutional layers in DumoulinEncoder must be log2(input_size) "
# network architecture
if self.config.hidden_channels is None:
self.config.hidden_channels = 8
hidden_channels = self.config.hidden_channels
kernels_size = [4, 4] * self.config.n_conv_layers
strides = [1, 2] * self.config.n_conv_layers
pads = [0, 1] * self.config.n_conv_layers
dils = [1, 1] * self.config.n_conv_layers
# feature map size
feature_map_sizes = conv_output_sizes(self.config.input_size,
2 * self.config.n_conv_layers,
kernels_size, strides, pads,
dils)
# local feature
## convolutional layers
self.local_feature_shape = (
int(hidden_channels * math.pow(2, self.config.feature_layer + 1)),
feature_map_sizes[2 * self.config.feature_layer + 1][0],
feature_map_sizes[2 * self.config.feature_layer + 1][1])
self.lf = nn.Sequential()
for conv_layer_id in range(self.config.feature_layer + 1):
if conv_layer_id == 0:
self.lf.add_module(
"conv_{}".format(conv_layer_id),
nn.Sequential(
ME.MinkowskiConvolution(
self.config.n_channels,
hidden_channels,
kernel_size=kernels_size[2 * conv_layer_id],
stride=strides[2 * conv_layer_id],
dilation=dils[2 * conv_layer_id],
dimension=self.spatial_dims,
bias=True),
ME.MinkowskiBatchNorm(hidden_channels),
ME.MinkowskiELU(inplace=True),
ME.MinkowskiConvolution(
hidden_channels,
2 * hidden_channels,
kernel_size=kernels_size[2 * conv_layer_id + 1],
stride=strides[2 * conv_layer_id + 1],
dilation=dils[2 * conv_layer_id + 1],
dimension=self.spatial_dims,
bias=True),
ME.MinkowskiBatchNorm(2 * hidden_channels),
ME.MinkowskiELU(inplace=True),
))
else:
self.lf.add_module(
"conv_{}".format(conv_layer_id),
nn.Sequential(
ME.MinkowskiConvolution(
hidden_channels,
hidden_channels,
kernel_size=kernels_size[2 * conv_layer_id],
stride=strides[2 * conv_layer_id],
dilation=dils[2 * conv_layer_id],
dimension=self.spatial_dims,
bias=True),
ME.MinkowskiBatchNorm(hidden_channels),
ME.MinkowskiELU(inplace=True),
ME.MinkowskiConvolution(
hidden_channels,
2 * hidden_channels,
kernel_size=kernels_size[2 * conv_layer_id + 1],
stride=strides[2 * conv_layer_id + 1],
dilation=dils[2 * conv_layer_id + 1],
dimension=self.spatial_dims,
bias=True),
ME.MinkowskiBatchNorm(2 * hidden_channels),
ME.MinkowskiELU(inplace=True),
))
hidden_channels *= 2
self.lf.out_connection_type = ("conv", hidden_channels)
# global feature
self.gf = nn.Sequential()
## convolutional layers
for conv_layer_id in range(self.config.feature_layer + 1,
self.config.n_conv_layers):
self.gf.add_module(
"conv_{}".format(conv_layer_id),
nn.Sequential(
ME.MinkowskiConvolution(
hidden_channels,
hidden_channels,
kernel_size=kernels_size[2 * conv_layer_id],
stride=strides[2 * conv_layer_id],
dilation=dils[2 * conv_layer_id],
dimension=self.spatial_dims,
bias=True),
ME.MinkowskiBatchNorm(hidden_channels),
ME.MinkowskiELU(inplace=True),
ME.MinkowskiConvolution(
hidden_channels,
2 * hidden_channels,
kernel_size=kernels_size[2 * conv_layer_id + 1],
stride=strides[2 * conv_layer_id + 1],
dilation=dils[2 * conv_layer_id + 1],
dimension=self.spatial_dims,
bias=True),
ME.MinkowskiBatchNorm(2 * hidden_channels),
ME.MinkowskiELU(inplace=True),
))
hidden_channels *= 2
self.gf.out_connection_type = ("conv", hidden_channels)
# encoding feature
if self.config.encoder_conditional_type == "gaussian":
self.add_module(
"ef",
nn.Sequential(
ME.MinkowskiConvolution(hidden_channels,
hidden_channels,
kernel_size=1,
stride=1,
dilation=1,
dimension=self.spatial_dims,
bias=True),
ME.MinkowskiBatchNorm(hidden_channels),
ME.MinkowskiELU(inplace=True),
ME.MinkowskiConvolution(hidden_channels,
2 * self.config.n_latents,
kernel_size=1,
stride=1,
dilation=1,
dimension=self.spatial_dims,
bias=True),
))
elif self.config.encoder_conditional_type == "deterministic":
self.add_module(
"ef",
nn.Sequential(
ME.MinkowskiConvolution(hidden_channels,
hidden_channels,
kernel_size=1,
stride=1,
dilation=1,
dimension=self.spatial_dims,
bias=True),
ME.MinkowskiBatchNorm(hidden_channels),
ME.MinkowskiELU(inplace=True),
ME.MinkowskiConvolution(hidden_channels,
self.config.n_latents,
kernel_size=1,
stride=1,
dilation=1,
dimension=self.spatial_dims,
bias=True),
))
# global pool
self.global_pool = ME.MinkowskiGlobalPooling()
# attention feature
if self.config.use_attention:
self.add_module(
"af",
ME.MinkowskiConvolution(hidden_channels,
4 * self.config.n_latents,
kernel_size=1,
stride=1,
dilation=1,
dimension=self.spatial_dims,
bias=True))
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()
# 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.MinkowskiGlobalPooling(mode=ME.GlobalPoolingMode.AUTO,
dimension=3)
# 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, coords=coordinates).to(device)
timer.tic()
soutput = pool(sinput)
timer.toc()
print(
# 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.MinkowskiGlobalPooling(mode=ME.GlobalPoolingMode.AUTO)
# 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, coords=coordinates).to(device)
timer.tic()
soutput = pool(sinput)
timer.toc()
print(
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)