def triple_conv(in_channels, out_channels, indice_key=None,activation=None):
# activation_fun = nn.ReLU(inplace=True)
# if activation=="lrelu":
# activation_fun=nn.LeakyReLU(0.1,inplace=True)
# elif activation=="mish":
# activation_fun=Mish()
return spconv.SparseSequential(
SubMConv3d(in_channels, out_channels, 3, indice_key=indice_key),
BatchNorm1d(out_channels),
activation(),
SubMConv3d(out_channels, out_channels, 3, indice_key=indice_key),
BatchNorm1d(out_channels),
activation(),
SubMConv3d(out_channels, out_channels, 3, indice_key=indice_key),
BatchNorm1d(out_channels,),
activation(),
)
def double_conv(in_channels, out_channels, indice_key=None):
return spconv.SparseSequential(
spconv.SubMConv3d(in_channels,
out_channels,
3,
bias=False,
indice_key=indice_key),
nn.BatchNorm1d(out_channels, eps=1e-3, momentum=0.01),
nn.ReLU(),
spconv.SubMConv3d(out_channels,
out_channels,
3,
bias=False,
indice_key=indice_key),
nn.BatchNorm1d(out_channels, eps=1e-3, momentum=0.01),
nn.ReLU(),
)
def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
kernel_size, stride):
super().__init__()
self.net = spconv.SparseSequential(
spconv.SparseConv3d(in_channels,
out_channels,
kernel_size,
stride,
indice_key="cp0",
bias=False),
spconv.SparseInverseConv3d(out_channels,
in_channels,
kernel_size,
indice_key="cp0",
bias=False),
)
self.todense = spconv.ToDense()
self.shape = shape
def createSparseConvBlock(in_channels, out_channels, kernel_size, indice_key=None, stride=1, padding=1,
conv_type='subm'):
if conv_type == 'subm':
conv = spconv.SubMConv3d(in_channels, out_channels, kernel_size, bias=False, indice_key=indice_key)
elif conv_type == 'spconv':
conv = spconv.SparseConv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding,
bias=False, indice_key=indice_key)
else:
raise NotImplementedError
m = spconv.SparseSequential(
conv,
nn.BatchNorm1d(out_channels, eps=1e-3, momentum=0.01),
nn.ReLU(),
)
return m
def __init__(self, num_input_features):
super(VxNet, self).__init__()
self.conv0 = double_conv(num_input_features, 16, 'subm0')
self.down0 = stride_conv(16, 32, 'down0')
self.conv1 = double_conv(32, 32, 'subm1')
self.down1 = stride_conv(32, 64, 'down1')
self.conv2 = triple_conv(64, 64, 'subm2')
self.down2 = stride_conv(64, 64, 'down2')
self.conv3 = triple_conv(64, 64, 'subm3') # middle line
self.extra_conv = spconv.SparseSequential(
spconv.SparseConv3d(64, 64, (1, 1, 1), (1, 1, 1),
bias=False), # shape no change
nn.BatchNorm1d(64, eps=1e-3, momentum=0.01),
nn.ReLU())
def _make_layer(self, inplanes, planes, num_blocks, idx, stride=1):
print("NUM BLOCKS:", num_blocks, "STRIDE:", stride)
if self._use_norm:
if self._use_groupnorm:
BatchNorm2d = change_default_args(
num_groups=self._num_groups, eps=1e-3)(GroupNorm)
else:
BatchNorm2d = change_default_args(
eps=1e-3, momentum=0.01)(nn.BatchNorm1d)
Conv2d = change_default_args(bias=False)(spconv.SparseConv2d)
SubMConv2d = change_default_args(bias=False)(spconv.SubMConv2d)
ConvTranspose2d = change_default_args(bias=False)(
spconv.SparseConvTranspose2d)
else:
BatchNorm2d = Empty
Conv2d = change_default_args(bias=True)(spconv.SparseConv2d)
SubMConv2d = change_default_args(bias=True)(spconv.SubMConv2d)
ConvTranspose2d = change_default_args(bias=True)(
spconv.SparseConvTranspose2d)
block = spconv.SparseSequential()
wide_conv_schedule = [None, 7, 9]
filter_radius = wide_conv_schedule[idx]
if filter_radius is not None:
print("Filter radius:", filter_radius)
block.add(SubMConv2d(inplanes, inplanes, filter_radius))
# Substitute for later convs.
num_blocks -= 1
block.add(Conv2d(inplanes, planes, 2, stride=stride))
block.add(BatchNorm2d(planes))
block.add(nn.ReLU())
for j in range(num_blocks):
block.add(SubMConv2d(planes, planes, 3, padding=1))
block.add(BatchNorm2d(planes)),
block.add(nn.ReLU())
# block.add(PrintLayer(2 + j))
return block, planes
def __init__(self, model_cfg, input_channels, grid_size, **kwargs):
super().__init__()
self.model_cfg = model_cfg
norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
self.sparse_shape = grid_size[::-1] + [1, 0, 0]
self.conv_input = spconv.SparseSequential(
spconv.SubMConv3d(input_channels, 16, 3, padding=1, bias=False, indice_key='subm1'),
norm_fn(16),
nn.ReLU(),
)
block = post_act_block
self.conv1 = spconv.SparseSequential(
SparseBasicBlock(16, 16, norm_fn=norm_fn, indice_key='res1'),
SparseBasicBlock(16, 16, norm_fn=norm_fn, indice_key='res1'),
)
self.conv2 = spconv.SparseSequential(
# [1600, 1408, 41] <- [800, 704, 21]
block(16, 32, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv2', conv_type='spconv'),
SparseBasicBlock(32, 32, norm_fn=norm_fn, indice_key='res2'),
SparseBasicBlock(32, 32, norm_fn=norm_fn, indice_key='res2'),
)
self.conv3 = spconv.SparseSequential(
# [800, 704, 21] <- [400, 352, 11]
block(32, 64, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv3', conv_type='spconv'),
SparseBasicBlock(64, 64, norm_fn=norm_fn, indice_key='res3'),
SparseBasicBlock(64, 64, norm_fn=norm_fn, indice_key='res3'),
)
self.conv4 = spconv.SparseSequential(
# [400, 352, 11] <- [200, 176, 5]
block(64, 128, 3, norm_fn=norm_fn, stride=2, padding=(0, 1, 1), indice_key='spconv4', conv_type='spconv'),
SparseBasicBlock(128, 128, norm_fn=norm_fn, indice_key='res4'),
SparseBasicBlock(128, 128, norm_fn=norm_fn, indice_key='res4'),
)
last_pad = 0
last_pad = self.model_cfg.get('last_pad', last_pad)
self.conv_out = spconv.SparseSequential(
# [200, 150, 5] -> [200, 150, 2]
spconv.SparseConv3d(128, 128, (3, 1, 1), stride=(2, 1, 1), padding=last_pad,
bias=False, indice_key='spconv_down2'),
norm_fn(128),
nn.ReLU(),
)
self.num_point_features = 128
def post_act_block(in_channels, out_channels, kernel_size, indice_key=None, stride=1, padding=0,
conv_type='subm', norm_fn=None):
if conv_type == 'subm':
conv = spconv.SubMConv3d(in_channels, out_channels, kernel_size, bias=False, indice_key=indice_key)
elif conv_type == 'spconv':
conv = spconv.SparseConv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding,
bias=False, indice_key=indice_key)
elif conv_type == 'inverseconv':
conv = spconv.SparseInverseConv3d(in_channels, out_channels, kernel_size, indice_key=indice_key, bias=False)
else:
raise NotImplementedError
m = spconv.SparseSequential(
conv,
norm_fn(out_channels),
nn.ReLU(),
)
return m
def __init__(self, shape): super().__init__() # For some reason the batch normalization is screwing things up idk why. # I heard it fixes itself when you do model.eval rather than just getting a single output self.net = spconv.SparseSequential( spconv.SparseConv3d(1, 32, 4), nn.LeakyReLU(), spconv.SparseConv3d(32, 32, 4), nn.LeakyReLU(), spconvpool.SparseMaxPool(3, 5), spconv.SparseConv3d(32, 32, 4), nn.LeakyReLU(), spconv.SparseConv3d(32, 32, 4), nn.LeakyReLU(), spconvpool.SparseMaxPool(3, 5), spconv.SparseConv3d(32, 32, 4), nn.LeakyReLU(), spconv.SparseConv3d(32, 32, 4), nn.LeakyReLU(), spconv.SparseConv3d(32, 1, 4), nn.LeakyReLU(), spconvpool.SparseMaxPool(3, 5), spconv.ToDense(), nn.Flatten(), nn.Linear(14688, 1000), nn.LeakyReLU(), nn.Linear(1000, 1000), nn.LeakyReLU(), nn.Linear(1000, 1), nn.LeakyReLU(), # nn.Sigmoid() ) self.shape = shape
def __init__(self,
num_layers,
ndim,
shape,
in_channels,
out_channels,
kernel_size,
stride,
padding,
dilation,
algo=spconv.ConvAlgo.BatchGemmGather):
super().__init__()
layers = [
spconv.SparseConv3d(in_channels,
out_channels,
kernel_size,
stride,
padding=padding,
dilation=dilation,
bias=False,
use_hash=False,
algo=algo)
]
for i in range(1, num_layers):
layers.append(
spconv.SparseConv3d(out_channels,
out_channels,
kernel_size,
stride,
padding=padding,
dilation=dilation,
bias=False,
use_hash=False,
algo=algo))
self.net = spconv.SparseSequential(*layers, )
# self.grid = torch.full([3, *shape], -1, dtype=torch.int32).cuda()
self.grid = None
self.shape = shape
def __init__(self, input_channels):
super().__init__()
norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
self.conv_input = spconv.SparseSequential(
spconv.SubMConv3d(input_channels, 16, 3, padding=1, bias=False, indice_key='subm1'),
norm_fn(16),
nn.ReLU(),
)
block = self.post_act_block
self.conv1 = spconv.SparseSequential(
block(16, 16, 3, norm_fn=norm_fn, padding=1, indice_key='subm1'),
)
self.conv2 = spconv.SparseSequential(
# [1600, 1408, 41] <- [800, 704, 21]
block(16, 32, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv2', conv_type='spconv'),
block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
)
self.conv3 = spconv.SparseSequential(
# [800, 704, 21] <- [400, 352, 11]
block(32, 64, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv3', conv_type='spconv'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
)
self.conv4 = spconv.SparseSequential(
# [400, 352, 11] <- [200, 176, 5]
block(64, 64, 3, norm_fn=norm_fn, stride=2, padding=(0, 1, 1), indice_key='spconv4', conv_type='spconv'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
)
last_pad = 0 if cfg.DATA_CONFIG.VOXEL_GENERATOR.VOXEL_SIZE[-1] in [0.1, 0.2] else (1, 0, 0)
self.conv_out = spconv.SparseSequential(
# [200, 150, 5] -> [200, 150, 2]
spconv.SparseConv3d(64, 128, (3, 1, 1), stride=(2, 1, 1), padding=last_pad,
bias=False, indice_key='spconv_down2'),
norm_fn(128),
nn.ReLU(),
)
def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
kernel_size, stride, padding, dilation):
super().__init__()
layers = [
spconv.SparseConvTranspose3d(in_channels,
out_channels,
kernel_size,
stride,
padding=padding,
dilation=dilation,
bias=False)
]
for i in range(1, num_layers):
layers.append(
spconv.SparseConvTranspose3d(out_channels,
out_channels,
kernel_size,
stride,
padding=padding,
dilation=dilation,
bias=False))
self.net = spconv.SparseSequential(*layers, )
self.shape = shape
def generate_block(self,
in_dim,
out_dim,
ksize,
stride,
padding,
do_subm=True):
block = spconv.SparseSequential(
spconv.SubMConv3d(in_channels=in_dim,
out_channels=out_dim,
kernel_size=1,
stride=1,
indice_key="subm0"),
nn.BatchNorm1d(num_features=out_dim),
nn.LeakyReLU(),
spconv.SubMConv3d(in_channels=in_dim,
out_channels=out_dim,
kernel_size=3,
stride=1,
padding=1,
indice_key='subm0')
if do_subm else spconv.SparseConv3d(in_channels=in_dim,
out_channels=out_dim,
kernel_size=3,
stride=1,
padding=1),
nn.BatchNorm1d(num_features=out_dim),
nn.LeakyReLU(),
spconv.SubMConv3d(in_channels=in_dim,
out_channels=out_dim,
kernel_size=1,
stride=1,
indice_key="subm0"),
nn.BatchNorm1d(num_features=out_dim),
spconv.ToDense(),
)
return block
def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
kernel_size, stride, padding, dilation):
super().__init__()
layers = [
spconv.SubMConv3d(in_channels,
out_channels,
kernel_size,
stride,
padding=padding,
dilation=dilation,
bias=False)
]
for i in range(1, num_layers):
layers.append(
spconv.SubMConv3d(out_channels,
out_channels,
kernel_size,
stride,
padding=padding,
dilation=dilation,
bias=False))
self.net = spconv.SparseSequential(*layers, )
self.grid = torch.full([3, *shape], -1, dtype=torch.int32).cuda()
self.shape = shape
def __init__(self,
arg_dict):
super(UNet_Like, self).__init__()
self.in_key = arg_dict['in_key']
self.out_key = arg_dict['out_key']
self.init_channel = int(arg_dict['init_channel'])
self.channels = eval(arg_dict['channels'])
assert isinstance(self.channels, list) and len(self.channels) > 1, "invalid channels"
self.se_mode = arg_dict['se_mode']
self.feature_concat = True if arg_dict['feature_concat'] == 'True' else False
self.use_norm = True if arg_dict['use_norm'] == 'True' else False
if self.use_norm:
BatchNorm1d = change_default_args(
eps=1e-3, momentum=0.01)(nn.BatchNorm1d)
SpConv3d = change_default_args(bias=False)(spconv.SparseConv3d)
SubMConv3d = change_default_args(bias=False)(spconv.SubMConv3d)
SpInverseConv3d = change_default_args(bias=False)(spconv.SparseInverseConv3d)
else:
BatchNorm1d = Empty
SpConv3d = change_default_args(bias=True)(spconv.SparseConv3d)
SubMConv3d = change_default_args(bias=True)(spconv.SubMConv3d)
SpInverseConv3d = change_default_args(bias=True)(spconv.SparseInverseConv3d)
self.stem = spconv.SparseSequential(
BasicBlock(self.init_channel, self.channels[0], 2, 'subm_0', self.se_mode, use_norm=self.use_norm),
BasicBlock(self.channels[0], self.channels[0], 2, 'subm_0', self.se_mode, use_norm=self.use_norm)
)
for level in range(1, len(self.channels)):
block = spconv.SparseSequential(
SpConv3d(self.channels[level - 1], self.channels[level], 3, 2, padding=1, indice_key=f'down_{level}'),
BatchNorm1d(self.channels[level]),
nn.ReLU(inplace=True),
BasicBlock(self.channels[level], self.channels[level], 2, f'subm_{level}', self.se_mode,
use_norm=self.use_norm),
BasicBlock(self.channels[level], self.channels[level], 2, f'subm_{level}', self.se_mode,
use_norm=self.use_norm)
)
self.add_module(f'encoder_{level}', block)
for level in range(1, len(self.channels)):
block = spconv.SparseSequential(
SpInverseConv3d(self.channels[level], self.channels[level - 1], 3, indice_key=f'down_{level}'),
BatchNorm1d(self.channels[level - 1]),
nn.ReLU(inplace=True),
BasicBlock(self.channels[level - 1], self.channels[level - 1], 2, f'subm_{level - 1}', self.se_mode,
use_norm=self.use_norm),
BasicBlock(self.channels[level - 1], self.channels[level - 1], 2, f'subm_{level - 1}', self.se_mode,
use_norm=self.use_norm)
)
self.add_module(f'decoder_{level}', block)
if self.feature_concat:
for level in range(len(self.channels) - 1):
block = spconv.SparseSequential(
SubMConv3d(self.channels[level] * 2, self.channels[level], 1),
BatchNorm1d(self.channels[level]),
nn.ReLU(inplace=True)
)
self.add_module(f'conv1x1_{level}', block)
def U(nPlanes): #Recursive function
m = spconv.SparseSequential()
if len(nPlanes) == 1:
for _ in range(reps):
m.add(
spconv.SparseBasicBlock(nPlanes[0],
nPlanes[0],
3,
indice_key="subm{}".format(
len(nPlanes))))
else:
m = spconv.SparseSequential()
for _ in range(reps):
m.add(
spconv.SparseBasicBlock(nPlanes[0],
nPlanes[0],
3,
indice_key="subm{}".format(
len(nPlanes))))
m.add(spconv.ConcatTable().add(spconv.Identity()).add(
spconv.SparseSequential().add(
spconv.SparseConv3d(
nPlanes[0],
nPlanes[1],
downsample[0],
stride=downsample[1],
bias=False,
indice_key="conv{}".format(len(nPlanes)))).add(
nn.BatchNorm1d(
nPlanes[1], eps=1e-3, momentum=0.01)).add(
nn.ReLU()).add(U(nPlanes[1:])).add(
spconv.SparseInverseConv3d(
nPlanes[1],
nPlanes[0],
downsample[0],
bias=False,
indice_key="conv{}".format(
len(nPlanes)))).add(
nn.BatchNorm1d(
nPlanes[0],
eps=1e-3,
momentum=0.01)).add(
nn.ReLU())))
m.add(spconv.JoinTable())
for i in range(reps):
m.add(
spconv.SubMConv3d(
nPlanes[0] * 2,
nPlanes[0],
3,
bias=False,
indice_key="end_pp{}".format(len(nPlanes)))).add(
nn.BatchNorm1d(nPlanes[0], eps=1e-3,
momentum=0.01)).add(nn.ReLU())
m.add(
spconv.SparseBasicBlock(nPlanes[0],
nPlanes[0],
3,
indice_key="end_pp{}".format(
len(nPlanes))))
return m
def __init__(self,
num_input_features=128,
norm_cfg=None,
name="SpMiddleResNetFHD",
**kwargs):
super(SpMiddleResNetFHD, self).__init__()
self.name = name
self.dcn = None
self.zero_init_residual = False
if norm_cfg is None:
norm_cfg = dict(type="BN1d", eps=1e-3, momentum=0.01)
# input: # [1600, 1200, 41]
self.conv_input = spconv.SparseSequential(
SubMConv3d(num_input_features,
16,
3,
bias=False,
indice_key="res0"),
build_norm_layer(norm_cfg, 16)[1], nn.ReLU(inplace=True))
self.conv1 = spconv.SparseSequential(
SparseBasicBlock(16, 16, norm_cfg=norm_cfg, indice_key="res0"),
SparseBasicBlock(16, 16, norm_cfg=norm_cfg, indice_key="res0"),
)
self.conv2 = spconv.SparseSequential(
SparseConv3d(16, 32, 3, 2, padding=1,
bias=False), # [1600, 1200, 41] -> [800, 600, 21]
build_norm_layer(norm_cfg, 32)[1],
nn.ReLU(inplace=True),
SparseBasicBlock(32, 32, norm_cfg=norm_cfg, indice_key="res1"),
SparseBasicBlock(32, 32, norm_cfg=norm_cfg, indice_key="res1"),
)
self.conv3 = spconv.SparseSequential(
SparseConv3d(32, 64, 3, 2, padding=1,
bias=False), # [800, 600, 21] -> [400, 300, 11]
build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(inplace=True),
SparseBasicBlock(64, 64, norm_cfg=norm_cfg, indice_key="res2"),
SparseBasicBlock(64, 64, norm_cfg=norm_cfg, indice_key="res2"),
)
self.conv4 = spconv.SparseSequential(
SparseConv3d(64, 128, 3, 2, padding=[0, 1, 1],
bias=False), # [400, 300, 11] -> [200, 150, 5]
build_norm_layer(norm_cfg, 128)[1],
nn.ReLU(inplace=True),
SparseBasicBlock(128, 128, norm_cfg=norm_cfg, indice_key="res3"),
SparseBasicBlock(128, 128, norm_cfg=norm_cfg, indice_key="res3"),
)
self.extra_conv = spconv.SparseSequential(
SparseConv3d(128, 128, (3, 1, 1), (2, 1, 1),
bias=False), # [200, 150, 5] -> [200, 150, 2]
build_norm_layer(norm_cfg, 128)[1],
nn.ReLU(),
)
def __init__(self,
num_input_features=128,
norm_cfg=None,
name="SpMiddleFHD",
**kwargs):
super(SpMiddleFHD, self).__init__()
self.name = name
self.dcn = None
self.zero_init_residual = False
if norm_cfg is None:
norm_cfg = dict(type="BN1d", eps=1e-3, momentum=0.01)
self.middle_conv = spconv.SparseSequential(
SubMConv3d(num_input_features,
16,
3,
bias=False,
indice_key="subm0"),
build_norm_layer(norm_cfg, 16)[1],
nn.ReLU(),
SubMConv3d(16, 16, 3, bias=False, indice_key="subm0"),
build_norm_layer(norm_cfg, 16)[1],
nn.ReLU(),
SparseConv3d(16, 32, 3, 2, padding=1,
bias=False), # [41, 1600, 1408] -> [21, 800, 704]
build_norm_layer(norm_cfg, 32)[1],
nn.ReLU(),
SubMConv3d(32, 32, 3, indice_key="subm1", bias=False),
build_norm_layer(norm_cfg, 32)[1],
nn.ReLU(),
SubMConv3d(32, 32, 3, indice_key="subm1", bias=False),
build_norm_layer(norm_cfg, 32)[1],
nn.ReLU(),
SparseConv3d(32, 64, 3, 2, padding=1,
bias=False), # [21, 800, 704] -> [11, 400, 352]
build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm2", bias=False),
build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm2", bias=False),
build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm2", bias=False),
build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SparseConv3d(64, 64, 3, 2, padding=[0, 1, 1],
bias=False), # [11, 400, 352] -> [5, 200, 176]
build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm3", bias=False),
build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm3", bias=False),
build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm3", bias=False),
build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SparseConv3d(64, 64, (3, 1, 1), (2, 1, 1),
bias=False), # [5, 200, 176] -> [2, 200, 176]
build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
)
def __init__(self,
num_input_features=128,
norm_cfg=None,
name="SpMiddleFHD",
**kwargs):
super(SpMiddleFHDNobn, self).__init__()
self.name = name
self.dcn = None
self.zero_init_residual = False
if norm_cfg is None:
norm_cfg = dict(type="BN1d", eps=1e-3, momentum=0.01)
self.middle_conv = spconv.SparseSequential(
SubMConv3d(num_input_features,
16,
3,
bias=True,
indice_key="subm0"),
# build_norm_layer(norm_cfg, 16)[1],
nn.ReLU(),
SubMConv3d(16, 16, 3, bias=True, indice_key="subm0"),
# build_norm_layer(norm_cfg, 16)[1],
nn.ReLU(),
SparseConv3d(16, 32, 3, 2, padding=1,
bias=True), # [1600, 1200, 41] -> [800, 600, 21]
# build_norm_layer(norm_cfg, 32)[1],
nn.ReLU(),
SubMConv3d(32, 32, 3, indice_key="subm1", bias=True),
# build_norm_layer(norm_cfg, 32)[1],
nn.ReLU(),
SubMConv3d(32, 32, 3, indice_key="subm1", bias=True),
# build_norm_layer(norm_cfg, 32)[1],
nn.ReLU(),
SparseConv3d(32, 64, 3, 2, padding=1,
bias=True), # [800, 600, 21] -> [400, 300, 11]
# build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm2", bias=True),
# build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm2", bias=True),
# build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm2", bias=True),
# build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SparseConv3d(64, 64, 3, 2, padding=[0, 1, 1],
bias=True), # [400, 300, 11] -> [200, 150, 5]
# build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm3", bias=True),
# build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm3", bias=True),
# build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm3", bias=True),
# build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SparseConv3d(64, 64, (3, 1, 1), (2, 1, 1),
bias=True), # [200, 150, 5] -> [200, 150, 2]
# build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
)
def __init__(self, cfg):
super().__init__()
input_c = cfg.input_channel
m = cfg.m
self.m = m
classes = cfg.classes
block_reps = cfg.block_reps
block_residual = cfg.block_residual
self.cluster_radius = cfg.cluster_radius
self.cluster_meanActive = cfg.cluster_meanActive
self.cluster_shift_meanActive = cfg.cluster_shift_meanActive
self.cluster_npoint_thre = cfg.cluster_npoint_thre
self.score_scale = cfg.score_scale
self.score_fullscale = cfg.score_fullscale
self.mode = cfg.score_mode
self.prepare_epochs = cfg.prepare_epochs
self.pretrain_path = cfg.pretrain_path
self.pretrain_module = cfg.pretrain_module
self.fix_module = cfg.fix_module
norm_fn = functools.partial(nn.BatchNorm1d, eps=1e-4, momentum=0.1)
if block_residual:
block = ResidualBlock
else:
block = VGGBlock
if cfg.use_coords:
input_c += 3
#### backbone
self.input_conv = spconv.SparseSequential(
spconv.SubMConv3d(input_c,
m,
kernel_size=3,
padding=1,
bias=False,
indice_key='subm1'))
self.unet = UBlock([m, 2 * m, 3 * m, 4 * m, 5 * m, 6 * m, 7 * m],
norm_fn,
block_reps,
block,
indice_key_id=1)
self.output_layer = spconv.SparseSequential(
norm_fn(m) #,
#nn.ReLU()
)
#### semantic segmentation
self.linear = nn.Linear(m, classes) # bias(default): True
self.feat = torch.Tensor(10, m).cuda()
self.start = 1
'''#### offset
def __init__(self, nPlanes, norm_fn, block_reps, block, indice_key_id=1):
super().__init__()
self.nPlanes = nPlanes
blocks = {
'block{}'.format(i):
block(nPlanes[0],
nPlanes[0],
norm_fn,
indice_key='subm{}'.format(indice_key_id))
for i in range(block_reps)
}
blocks = OrderedDict(blocks)
self.blocks = spconv.SparseSequential(blocks)
if len(nPlanes) > 1:
self.conv = spconv.SparseSequential(
norm_fn(nPlanes[0]), nn.ReLU(),
spconv.SparseConv3d(
nPlanes[0],
nPlanes[1],
kernel_size=2,
stride=2,
bias=False,
indice_key='spconv{}'.format(indice_key_id)))
self.u = UBlock(nPlanes[1:],
norm_fn,
block_reps,
block,
indice_key_id=indice_key_id + 1)
if nPlanes[0] == 32:
self.deconv = spconv.SparseSequential(
norm_fn(nPlanes[1]), nn.ReLU(),
spconv.SparseInverseConv3d(
nPlanes[1],
nPlanes[0],
kernel_size=2,
bias=False,
indice_key='spconv{}'.format(indice_key_id)))
else:
self.deconv = spconv.SparseSequential(
norm_fn(nPlanes[1]), nn.ReLU(),
spconv.SparseInverseConv3d(
nPlanes[1],
nPlanes[0],
kernel_size=2,
bias=False,
indice_key='spconv{}'.format(indice_key_id)))
blocks_tail = {}
for i in range(block_reps):
blocks_tail['block{}'.format(i)] = block(
nPlanes[0] * (2 - i),
nPlanes[0],
norm_fn,
indice_key='subm{}'.format(indice_key_id))
blocks_tail = OrderedDict(blocks_tail)
self.blocks_tail = spconv.SparseSequential(blocks_tail)
def __init__(self, cfg):
super().__init__(cfg)
self.occupancy_cluster = cfg.occupancy_cluster
self.input_conv = spconv.SparseSequential(
spconv.SubMConv3d(self.input_c,
self.m,
kernel_size=3,
padding=1,
bias=False,
indice_key='subm1'))
self.unet = UBlock([
self.m, 2 * self.m, 3 * self.m, 4 * self.m, 5 * self.m, 6 * self.m,
7 * self.m
],
self.norm_fn,
self.block_reps,
self.block,
indice_key_id=1,
backbone=True,
UNet_Transformer=cfg.UNet_Transformer)
self.output_layer = spconv.SparseSequential(self.norm_fn(self.m),
nn.ReLU())
self.point_offset = nn.Sequential(
nn.Linear(self.m, self.m, bias=True),
self.norm_fn(self.m),
nn.ReLU(),
nn.Linear(self.m, 3, bias=True),
)
self.point_semantic = nn.Sequential(
nn.Linear(self.m, self.m, bias=True),
self.norm_fn(self.m),
nn.ReLU(),
nn.Linear(self.m, self.classes, bias=True),
)
self.point_occupancy = nn.Sequential(
nn.Linear(self.m, self.m, bias=True),
self.norm_fn(self.m),
nn.ReLU(),
nn.Linear(self.m, 1, bias=True),
)
#### score branch
self.score_unet = UBlock([self.m, 2 * self.m],
self.norm_fn,
2,
self.block,
indice_key_id=1,
backbone=False)
self.score_outputlayer = spconv.SparseSequential(
self.norm_fn(self.m), nn.ReLU())
self.score_linear = nn.Linear(self.m, 1)
self.apply(self.set_bn_init)
self.module_map = {
'input_conv': self.input_conv,
'unet': self.unet,
'output_layer': self.output_layer,
'score_unet': self.score_unet,
'score_outputlayer': self.score_outputlayer,
'score_linear': self.score_linear,
}
self.local_pretrained_model_parameter()
def __init__(self,
num_input_features=128,
norm_cfg=None,
name="RCNNSpMiddleFHD",
**kwargs):
super(RCNNSpMiddleFHD, self).__init__()
self.name = name
self.dcn = None
self.zero_init_residual = False
if norm_cfg is None:
norm_cfg = dict(type="BN1d", eps=1e-3, momentum=0.01)
self.middle_conv = spconv.SparseSequential(
SubMConv3d(num_input_features,
16,
3,
bias=False,
indice_key="subm0"),
build_norm_layer(norm_cfg, 16)[1],
nn.ReLU(),
SubMConv3d(16, 16, 3, bias=False, indice_key="subm0"),
build_norm_layer(norm_cfg, 16)[1],
nn.ReLU(),
SparseConv3d(16, 32, 3, 2, padding=1,
bias=False), # [32, 80, 41] -> [16, 40, 21]
build_norm_layer(norm_cfg, 32)[1],
nn.ReLU(),
SubMConv3d(32, 32, 3, bias=False, indice_key="subm1"),
build_norm_layer(norm_cfg, 32)[1],
nn.ReLU(),
# SubMConv3d(32, 32, 3, bias=False, indice_key="subm1"),
# build_norm_layer(norm_cfg, 32)[1],
# nn.ReLU(),
SparseConv3d(32, 64, 3, 2, bias=False,
padding=1), # [16, 40, 21] -> [8, 20, 11]
build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SubMConv3d(64, 64, 3, bias=False, indice_key="subm2"),
build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
# SubMConv3d(64, 64, 3, bias=False, indice_key="subm2"),
# build_norm_layer(norm_cfg, 64)[1],
# nn.ReLU(),
# SubMConv3d(64, 64, 3, bias=False, indice_key="subm2"),
# build_norm_layer(norm_cfg, 64)[1],
# nn.ReLU(),
SparseConv3d(64, 64, 3, 2, bias=False,
padding=[1, 1, 0]), # [8, 20, 11] -> [4, 10, 5]
build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
SubMConv3d(64, 64, 3, bias=False, indice_key="subm3"),
build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
# SubMConv3d(64, 64, 3, bias=False, indice_key="subm3"),
# build_norm_layer(norm_cfg, 64)[1],
# nn.ReLU(),
# SubMConv3d(64, 64, 3, bias=False, indice_key="subm3"),
# build_norm_layer(norm_cfg, 64)[1],
# nn.ReLU(),
SparseConv3d(64, 64, (1, 1, 3), (1, 1, 2),
bias=False), # [4, 10, 5] -> [4, 10, 2]
build_norm_layer(norm_cfg, 64)[1],
nn.ReLU(),
)
def __init__(self,
output_shape,
num_input_features=128,
num_filters_down1=[64],
num_filters_down2=[64, 64],
use_norm=True,
num_class=2,
layer_nums=[3, 5, 5],
layer_strides=[2, 2, 2],
num_filters=[128, 128, 256],
upsample_strides=[1, 2, 4],
num_upsample_filters=[256, 256, 256],
num_anchor_per_loc=2,
encode_background_as_zeros=True,
use_direction_classifier=True,
use_groupnorm=False,
num_groups=32,
use_bev=False,
box_code_size=7,
use_rc_net=False,
name='sparse_rpn'):
super(SparseRPN, self).__init__()
self._num_anchor_per_loc = num_anchor_per_loc
self._use_direction_classifier = use_direction_classifier
self.name = name
if use_norm:
BatchNorm2d = change_default_args(eps=1e-3,
momentum=0.01)(nn.BatchNorm2d)
BatchNorm1d = change_default_args(eps=1e-3,
momentum=0.01)(nn.BatchNorm1d)
Conv2d = change_default_args(bias=False)(nn.Conv2d)
SpConv3d = change_default_args(bias=False)(spconv.SparseConv3d)
SubMConv3d = change_default_args(bias=False)(spconv.SubMConv3d)
ConvTranspose2d = change_default_args(bias=False)(
nn.ConvTranspose2d)
else:
BatchNorm2d = Empty
BatchNorm1d = Empty
Conv2d = change_default_args(bias=True)(nn.Conv2d)
SpConv3d = change_default_args(bias=True)(spconv.SparseConv3d)
SubMConv3d = change_default_args(bias=True)(spconv.SubMConv3d)
ConvTranspose2d = change_default_args(bias=True)(
nn.ConvTranspose2d)
sparse_shape = np.array(output_shape[1:4]) + [1, 0, 0]
# sparse_shape[0] = 11
print(sparse_shape)
self.sparse_shape = sparse_shape
self.voxel_output_shape = output_shape
# [11, 400, 352]
self.block1 = spconv.SparseSequential(
SpConv3d(num_input_features,
num_filters[0],
3,
stride=[2, layer_strides[0], layer_strides[0]],
padding=[0, 1, 1]), BatchNorm1d(num_filters[0]),
nn.ReLU())
# [5, 200, 176]
for i in range(layer_nums[0]):
self.block1.add(
SubMConv3d(num_filters[0],
num_filters[0],
3,
padding=1,
indice_key="subm0"))
self.block1.add(BatchNorm1d(num_filters[0]))
self.block1.add(nn.ReLU())
self.deconv1 = spconv.SparseSequential(
SpConv3d(num_filters[0],
num_filters[0], (3, 1, 1),
stride=(2, 1, 1)), BatchNorm1d(num_filters[0]), nn.ReLU(),
SpConv3d(num_filters[0],
num_upsample_filters[0], (2, 1, 1),
stride=1), BatchNorm1d(num_upsample_filters[0]),
nn.ReLU(), spconv.ToDense(), Squeeze()) # [1, 200, 176]
# [5, 200, 176]
self.block2 = spconv.SparseSequential(
SpConv3d(num_filters[0],
num_filters[1],
3,
stride=[2, layer_strides[1], layer_strides[1]],
padding=[0, 1, 1]), BatchNorm1d(num_filters[1]),
nn.ReLU())
for i in range(layer_nums[1]):
self.block2.add(
SubMConv3d(num_filters[1],
num_filters[1],
3,
padding=1,
indice_key="subm1"))
self.block2.add(BatchNorm1d(num_filters[1]))
self.block2.add(nn.ReLU())
# [2, 100, 88]
self.deconv2 = spconv.SparseSequential(
SpConv3d(num_filters[1], num_filters[1], (2, 1, 1), stride=1),
BatchNorm1d(num_filters[1]), nn.ReLU(), spconv.ToDense(),
Squeeze(),
ConvTranspose2d(num_filters[1],
num_upsample_filters[1],
upsample_strides[1],
stride=upsample_strides[1]),
BatchNorm2d(num_upsample_filters[1]), nn.ReLU()) # [1, 200, 176]
self.block3 = spconv.SparseSequential(
SpConv3d(num_filters[1],
num_filters[2], [2, 3, 3],
stride=[1, layer_strides[2], layer_strides[2]],
padding=[0, 1, 1]), BatchNorm1d(num_filters[2]),
nn.ReLU())
for i in range(layer_nums[2]):
self.block3.add(
SubMConv3d(num_filters[2],
num_filters[2],
3,
padding=1,
indice_key="subm2"))
self.block3.add(BatchNorm1d(num_filters[2]))
self.block3.add(nn.ReLU())
self.deconv3 = Sequential(
spconv.ToDense(),
Squeeze(),
ConvTranspose2d(num_filters[2],
num_upsample_filters[2],
upsample_strides[2],
stride=upsample_strides[2]),
BatchNorm2d(num_upsample_filters[2]),
nn.ReLU(),
) # [1, 200, 176]
self.post = Sequential(
Conv2d(sum(num_upsample_filters), 128, 3, stride=1, padding=1),
BatchNorm2d(128),
nn.ReLU(),
Conv2d(128, 64, 3, stride=1, padding=1),
BatchNorm2d(64),
nn.ReLU(),
) # [1, 200, 176]
if encode_background_as_zeros:
num_cls = num_anchor_per_loc * num_class
else:
num_cls = num_anchor_per_loc * (num_class + 1)
'''self.conv_cls = nn.Conv2d(sum(num_upsample_filters), num_cls, 1)
self.conv_box = nn.Conv2d(
sum(num_upsample_filters), num_anchor_per_loc * box_code_size, 1)
if use_direction_classifier:
self.conv_dir_cls = nn.Conv2d(
sum(num_upsample_filters), num_anchor_per_loc * 2, 1)
'''
self.conv_cls = nn.Conv2d(64, num_cls, 1)
self.conv_box = nn.Conv2d(64, num_anchor_per_loc * box_code_size, 1)
if use_direction_classifier:
self.conv_dir_cls = nn.Conv2d(64, num_anchor_per_loc * 2, 1)
def __init__(self,
output_shape,
use_norm=True,
num_input_features=128,
num_filters_down1=[64],
num_filters_down2=[64, 64],
name='SpMiddle2K'):
super(SpMiddle2K, self).__init__()
self.name = name
if use_norm:
BatchNorm2d = change_default_args(eps=1e-3,
momentum=0.01)(nn.BatchNorm2d)
BatchNorm1d = change_default_args(eps=1e-3,
momentum=0.01)(nn.BatchNorm1d)
Conv2d = change_default_args(bias=False)(nn.Conv2d)
SpConv3d = change_default_args(bias=False)(spconv.SparseConv3d)
SubMConv3d = change_default_args(bias=False)(spconv.SubMConv3d)
ConvTranspose2d = change_default_args(bias=False)(
nn.ConvTranspose2d)
else:
BatchNorm2d = Empty
BatchNorm1d = Empty
Conv2d = change_default_args(bias=True)(nn.Conv2d)
SpConv3d = change_default_args(bias=True)(spconv.SparseConv3d)
SubMConv3d = change_default_args(bias=True)(spconv.SubMConv3d)
ConvTranspose2d = change_default_args(bias=True)(
nn.ConvTranspose2d)
sparse_shape = np.array(output_shape[1:4]) + [1, 0, 0]
# sparse_shape[0] = 11
print(sparse_shape)
self.sparse_shape = sparse_shape
self.voxel_output_shape = output_shape
# input: # [1600, 1200, 41]
self.middle_conv = spconv.SparseSequential(
SubMConv3d(
num_input_features, 8, 3,
indice_key="subm0"), # [3200, 2400, 81] -> [1600, 1200, 41]
BatchNorm1d(8),
nn.ReLU(),
SubMConv3d(8, 8, 3, indice_key="subm0"),
BatchNorm1d(8),
nn.ReLU(),
SpConv3d(8, 16, 3, 2,
padding=1), # [1600, 1200, 41] -> [800, 600, 21]
BatchNorm1d(16),
nn.ReLU(),
SubMConv3d(16, 16, 3, indice_key="subm1"),
BatchNorm1d(16),
nn.ReLU(),
SubMConv3d(16, 16, 3, indice_key="subm1"),
BatchNorm1d(16),
nn.ReLU(),
SpConv3d(16, 32, 3, 2,
padding=1), # [1600, 1200, 41] -> [800, 600, 21]
BatchNorm1d(32),
nn.ReLU(),
SubMConv3d(32, 32, 3, indice_key="subm2"),
BatchNorm1d(32),
nn.ReLU(),
SubMConv3d(32, 32, 3, indice_key="subm2"),
BatchNorm1d(32),
nn.ReLU(),
SpConv3d(32, 64, 3, 2,
padding=1), # [800, 600, 21] -> [400, 300, 11]
BatchNorm1d(64),
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm3"),
BatchNorm1d(64),
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm3"),
BatchNorm1d(64),
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm3"),
BatchNorm1d(64),
nn.ReLU(),
SpConv3d(64, 64, 3, 2,
padding=[0, 1, 1]), # [400, 300, 11] -> [200, 150, 5]
BatchNorm1d(64),
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm4"),
BatchNorm1d(64),
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm4"),
BatchNorm1d(64),
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm4"),
BatchNorm1d(64),
nn.ReLU(),
SpConv3d(64, 64, (3, 1, 1),
(2, 1, 1)), # [200, 150, 5] -> [200, 150, 2]
BatchNorm1d(64),
nn.ReLU(),
)
self.max_batch_size = 3
self.grid = torch.full([self.max_batch_size, *sparse_shape],
-1,
dtype=torch.int32).cuda()
def make_conv(in_channels,
out_channels,
kernel_size,
stride=1,
padding=0,
dilation=1,
indice_key="subm0"):
if use_weight_std:
assert use_submconv and not use_deconv, "WS are not added to others spconv layers yet"
if use_submconv:
if use_deconv:
conv_func = spconv.SparseConvTranspose2d
# conv_func = spconv.SparseInverseConv2d
else:
if use_weight_std:
conv_func = SubMConv2dWS
else:
conv_func = spconv.SubMConv2d
else:
if use_deconv:
conv_func = spconv.SparseConvTranspose2d
else:
conv_func = spconv.SparseConv2d
if use_sep:
assert in_channels == out_channels
groups = in_channels
else:
groups = 1
try:
conv = conv_func(in_channels,
out_channels,
kernel_size=kernel_size,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
bias=(norm is None),
indice_key=indice_key,
algo=spconv.ops.ConvAlgo.Native)
except:
conv = conv_func(in_channels,
out_channels,
kernel_size=kernel_size,
bias=(norm is None),
indice_key=indice_key,
algo=spconv.ops.ConvAlgo.Native)
# Caffe2 implementation uses XavierFill, which in fact
# corresponds to kaiming_uniform_ in PyTorch
nn.init.kaiming_uniform_(conv.weight, a=1)
if norm is None:
nn.init.constant_(conv.bias, 0)
module = [
conv,
]
if norm is not None and len(norm) > 0:
if norm == "GN":
raise NotImplementedError
print("GN")
norm_module = nn.GroupNorm(32, out_channels)
# elif norm == "SparseGN":
# # raise NotImplementedError
# norm_module = SparseGN(32, out_channels)
elif norm == "BN":
norm_module = nn.BatchNorm1d(out_channels)
else:
raise NotImplementedError
module.append(norm_module)
if activation is not None:
module.append(nn.ReLU(inplace=True))
if len(module) > 1:
return spconv.SparseSequential(*module)
return conv
def __init__(self,
output_shape,
use_norm=True,
num_input_features=128,
num_filters_down1=[64],
num_filters_down2=[64, 64],
name='SpMiddleD4HD'):
super(SpMiddleD4HD, self).__init__()
self.name = name
if use_norm:
BatchNorm2d = change_default_args(eps=1e-3,
momentum=0.01)(nn.BatchNorm2d)
BatchNorm1d = change_default_args(eps=1e-3,
momentum=0.01)(nn.BatchNorm1d)
Conv2d = change_default_args(bias=False)(nn.Conv2d)
SpConv3d = change_default_args(bias=False)(spconv.SparseConv3d)
SubMConv3d = change_default_args(bias=False)(spconv.SubMConv3d)
ConvTranspose2d = change_default_args(bias=False)(
nn.ConvTranspose2d)
else:
BatchNorm2d = Empty
BatchNorm1d = Empty
Conv2d = change_default_args(bias=True)(nn.Conv2d)
SpConv3d = change_default_args(bias=True)(spconv.SparseConv3d)
SubMConv3d = change_default_args(bias=True)(spconv.SubMConv3d)
ConvTranspose2d = change_default_args(bias=True)(
nn.ConvTranspose2d)
sparse_shape = np.array(output_shape[1:4]) + [1, 0, 0]
# sparse_shape[0] = 11
print(sparse_shape)
self.sparse_shape = sparse_shape
self.voxel_output_shape = output_shape
# num_input_features = 4
self.middle_conv = spconv.SparseSequential(
SubMConv3d(num_input_features, 32, 3, indice_key="subm0"),
BatchNorm1d(32),
nn.ReLU(),
SubMConv3d(32, 32, 3, indice_key="subm0"),
BatchNorm1d(32),
nn.ReLU(),
SpConv3d(32, 64, 3, 2,
padding=1), # [800, 600, 21] -> [400, 300, 11]
BatchNorm1d(64),
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm1"),
BatchNorm1d(64),
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm1"),
BatchNorm1d(64),
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm1"),
BatchNorm1d(64),
nn.ReLU(),
SpConv3d(64, 64, 3, 2,
padding=[0, 1, 1]), # [400, 300, 11] -> [200, 150, 5]
BatchNorm1d(64),
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm2"),
BatchNorm1d(64),
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm2"),
BatchNorm1d(64),
nn.ReLU(),
SubMConv3d(64, 64, 3, indice_key="subm2"),
BatchNorm1d(64),
nn.ReLU(),
SpConv3d(64, 64, (3, 1, 1),
(2, 1, 1)), # [200, 150, 5] -> [200, 150, 2]
BatchNorm1d(64),
nn.ReLU(),
)
def __init__(self, cfg):
super().__init__()
input_c = cfg.input_channel
m = cfg.m
classes = cfg.classes
block_reps = cfg.block_reps
block_residual = cfg.block_residual
self.cluster_radius = cfg.cluster_radius
self.cluster_meanActive = cfg.cluster_meanActive
self.cluster_shift_meanActive = cfg.cluster_shift_meanActive
self.cluster_npoint_thre = cfg.cluster_npoint_thre
self.score_scale = cfg.score_scale
self.score_fullscale = cfg.score_fullscale
self.mode = cfg.score_mode
self.prepare_epochs = cfg.prepare_epochs
self.pretrain_path = cfg.pretrain_path
self.pretrain_module = cfg.pretrain_module
self.fix_module = cfg.fix_module
norm_fn = functools.partial(nn.BatchNorm1d, eps=1e-4, momentum=0.1)
if block_residual:
block = ResidualBlock
else:
block = VGGBlock
if cfg.use_coords:
input_c += 3
#### backbone
self.input_conv = spconv.SparseSequential(
spconv.SubMConv3d(input_c,
m,
kernel_size=3,
padding=1,
bias=False,
indice_key='subm1'))
self.unet = UBlock([m, 2 * m, 3 * m, 4 * m, 5 * m, 6 * m, 7 * m],
norm_fn,
block_reps,
block,
indice_key_id=1)
self.output_layer = spconv.SparseSequential(norm_fn(m), nn.ReLU())
#### semantic segmentation
self.linear = nn.Linear(m, classes) # bias(default): True
#### offset
self.offset = nn.Sequential(nn.Linear(m, m, bias=True), norm_fn(m),
nn.ReLU())
self.offset_linear = nn.Linear(m, 3, bias=True)
#### score branch
self.score_unet = UBlock([m, 2 * m],
norm_fn,
2,
block,
indice_key_id=1)
self.score_outputlayer = spconv.SparseSequential(norm_fn(m), nn.ReLU())
self.score_linear = nn.Linear(m, 1)
self.apply(self.set_bn_init)
#### fix parameter
module_map = {
'input_conv': self.input_conv,
'unet': self.unet,
'output_layer': self.output_layer,
'linear': self.linear,
'offset': self.offset,
'offset_linear': self.offset_linear,
'score_unet': self.score_unet,
'score_outputlayer': self.score_outputlayer,
'score_linear': self.score_linear
}
for m in self.fix_module:
mod = module_map[m]
for param in mod.parameters():
param.requires_grad = False
#### load pretrain weights --> Error: https://github.com/CSAILVision/places365/issues/25
if self.pretrain_path is not None:
pretrain_dict = torch.load(self.pretrain_path)
for m in self.pretrain_module:
print("Load pretrained " + m +
": %d/%d" % utils.load_model_param(
module_map[m], pretrain_dict, prefix=m))
def __init__(self, model_cfg, input_channels, grid_size, voxel_size,
point_cloud_range, **kwargs):
super().__init__()
self.model_cfg = model_cfg
self.sparse_shape = grid_size[::-1] + [1, 0, 0]
self.voxel_size = voxel_size
self.point_cloud_range = point_cloud_range
norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)
self.conv_input = spconv.SparseSequential(
spconv.SubMConv3d(input_channels,
16,
3,
padding=1,
bias=False,
indice_key='subm1'),
norm_fn(16),
nn.ReLU(),
)
block = post_act_block
self.conv1 = spconv.SparseSequential(
block(16, 16, 3, norm_fn=norm_fn, padding=1, indice_key='subm1'), )
self.conv2 = spconv.SparseSequential(
# [1600, 1408, 41] <- [800, 704, 21]
block(16,
32,
3,
norm_fn=norm_fn,
stride=2,
padding=1,
indice_key='spconv2',
conv_type='spconv'),
block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
)
self.conv3 = spconv.SparseSequential(
# [800, 704, 21] <- [400, 352, 11]
block(32,
64,
3,
norm_fn=norm_fn,
stride=2,
padding=1,
indice_key='spconv3',
conv_type='spconv'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
)
self.conv4 = spconv.SparseSequential(
# [400, 352, 11] <- [200, 176, 5]
block(64,
64,
3,
norm_fn=norm_fn,
stride=2,
padding=(0, 1, 1),
indice_key='spconv4',
conv_type='spconv'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
)
if self.model_cfg.get('RETURN_ENCODED_TENSOR', True):
last_pad = self.model_cfg.get('last_pad', 0)
self.conv_out = spconv.SparseSequential(
# [200, 150, 5] -> [200, 150, 2]
spconv.SparseConv3d(64,
128, (3, 1, 1),
stride=(2, 1, 1),
padding=last_pad,
bias=False,
indice_key='spconv_down2'),
norm_fn(128),
nn.ReLU(),
)
else:
self.conv_out = None
# decoder
# [400, 352, 11] <- [200, 176, 5]
self.conv_up_t4 = SparseBasicBlock(64,
64,
indice_key='subm4',
norm_fn=norm_fn)
self.conv_up_m4 = block(128,
64,
3,
norm_fn=norm_fn,
padding=1,
indice_key='subm4')
self.inv_conv4 = block(64,
64,
3,
norm_fn=norm_fn,
indice_key='spconv4',
conv_type='inverseconv')
# [800, 704, 21] <- [400, 352, 11]
self.conv_up_t3 = SparseBasicBlock(64,
64,
indice_key='subm3',
norm_fn=norm_fn)
self.conv_up_m3 = block(128,
64,
3,
norm_fn=norm_fn,
padding=1,
indice_key='subm3')
self.inv_conv3 = block(64,
32,
3,
norm_fn=norm_fn,
indice_key='spconv3',
conv_type='inverseconv')
# [1600, 1408, 41] <- [800, 704, 21]
self.conv_up_t2 = SparseBasicBlock(32,
32,
indice_key='subm2',
norm_fn=norm_fn)
self.conv_up_m2 = block(64, 32, 3, norm_fn=norm_fn, indice_key='subm2')
self.inv_conv2 = block(32,
16,
3,
norm_fn=norm_fn,
indice_key='spconv2',
conv_type='inverseconv')
# [1600, 1408, 41] <- [1600, 1408, 41]
self.conv_up_t1 = SparseBasicBlock(16,
16,
indice_key='subm1',
norm_fn=norm_fn)
self.conv_up_m1 = block(32, 16, 3, norm_fn=norm_fn, indice_key='subm1')
self.conv5 = spconv.SparseSequential(
block(16, 16, 3, norm_fn=norm_fn, padding=1, indice_key='subm1'))
self.num_point_features = 16
def __init__(self, cfg):
super().__init__(cfg)
'''backbone network'''
self.input_conv = spconv.SparseSequential(
spconv.SubMConv3d(self.input_c,
self.m,
kernel_size=3,
padding=1,
bias=False,
indice_key='subm1'))
self.unet = UBlock([
self.m, 2 * self.m, 3 * self.m, 4 * self.m, 5 * self.m, 6 * self.m,
7 * self.m
],
self.norm_fn,
self.block_reps,
self.block,
indice_key_id=1,
backbone=True,
UNet_Transformer=cfg.UNet_Transformer)
self.output_layer = spconv.SparseSequential(self.norm_fn(self.m),
nn.ReLU())
self.point_offset = nn.Sequential(
nn.Linear(self.m, self.m, bias=True),
self.norm_fn(self.m),
nn.ReLU(),
nn.Linear(self.m, 3, bias=True),
)
self.point_semantic = nn.Sequential(
nn.Linear(self.m, self.m, bias=True),
self.norm_fn(self.m),
nn.ReLU(),
nn.Linear(self.m, self.classes, bias=True),
)
#### score branch
self.score_unet = UBlock([self.m, 2 * self.m],
self.norm_fn,
2,
self.block,
indice_key_id=1,
backbone=False)
self.score_outputlayer = spconv.SparseSequential(
self.norm_fn(self.m), nn.ReLU())
self.score_linear = nn.Linear(self.m, 1)
if self.point_xyz_reconstruction_loss['activate']:
self.point_reconstruction_coords = nn.Sequential(
nn.Linear(self.m, self.m, bias=True),
self.norm_fn(self.m),
nn.ReLU(),
nn.Linear(self.m, 3, bias=True),
)
if self.instance_classifier['activate']:
self.point_instance_classifier = nn.Sequential(
nn.Linear(self.m, self.m, bias=True),
self.norm_fn(self.m),
nn.ReLU(),
nn.Linear(self.m,
self.instance_classifier['instance_num'],
bias=True),
)
if self.voxel_center_prediction['activate']:
self.voxel_center_pred = nn.Sequential(
nn.Linear(self.m, self.m, bias=True),
self.norm_fn(self.m),
nn.ReLU(),
nn.Linear(self.m, 1, bias=True),
)
self.voxel_center_offset = nn.Sequential(
nn.Linear(self.m, self.m, bias=True),
self.norm_fn(self.m),
nn.ReLU(),
nn.Linear(self.m, 3, bias=True),
)
self.voxel_center_semantic = nn.Sequential(
nn.Linear(self.m, self.m, bias=True),
self.norm_fn(self.m),
nn.ReLU(),
nn.Linear(self.m, self.classes, bias=True),
)
self.apply(self.set_bn_init)
self.module_map = {
'input_conv': self.input_conv,
'unet': self.unet,
'output_layer': self.output_layer,
'score_unet': self.score_unet,
'score_outputlayer': self.score_outputlayer,
'score_linear': self.score_linear,
'point_offset': self.point_offset,
'point_semantic': self.point_semantic
}
self.local_pretrained_model_parameter()