def __init__(self,
in_filters,
out_filters,
kernel_size=(3, 3, 3),
indice_key=None,
up_key=None):
super(UpBlock, self).__init__()
# self.drop_out = drop_out
self.trans_dilao = conv3x3(in_filters,
out_filters,
indice_key=indice_key + "new_up")
self.trans_act = nn.LeakyReLU()
self.trans_bn = nn.BatchNorm1d(out_filters)
self.conv1 = conv1x3(out_filters, out_filters, indice_key=indice_key)
self.act1 = nn.LeakyReLU()
self.bn1 = nn.BatchNorm1d(out_filters)
self.conv2 = conv3x1(out_filters, out_filters, indice_key=indice_key)
self.act2 = nn.LeakyReLU()
self.bn2 = nn.BatchNorm1d(out_filters)
self.conv3 = conv3x3(out_filters, out_filters, indice_key=indice_key)
self.act3 = nn.LeakyReLU()
self.bn3 = nn.BatchNorm1d(out_filters)
# self.dropout3 = nn.Dropout3d(p=dropout_rate)
self.up_subm = spconv.SparseInverseConv3d(out_filters,
out_filters,
kernel_size=3,
indice_key=up_key,
bias=False)
self.weight_initialization()
def post_act_block(self, in_channels, out_channels, kernel_size, indice_key, stride=1, padding=0,
conv_type='subm', norm_fn=None):
if conv_type == 'subm':
m = spconv.SparseSequential(
spconv.SubMConv3d(in_channels, out_channels, kernel_size, bias=False, indice_key=indice_key),
norm_fn(out_channels),
nn.ReLU(),
)
elif conv_type == 'spconv':
m = spconv.SparseSequential(
spconv.SparseConv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding,
bias=False, indice_key=indice_key),
norm_fn(out_channels),
nn.ReLU(),
)
elif conv_type == 'inverseconv':
m = spconv.SparseSequential(
spconv.SparseInverseConv3d(in_channels, out_channels, kernel_size,
indice_key=indice_key, bias=False),
norm_fn(out_channels),
nn.ReLU(),
)
else:
raise NotImplementedError
return m
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)
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, nPlanes, norm_fn, block_reps, block, indice_key_id=1, backbone=False, UNet_Transformer=None):
super().__init__()
self.nPlanes = nPlanes
self.backbone = backbone
self.UNet_Transformer = UNet_Transformer
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,
backbone=backbone, UNet_Transformer=self.UNet_Transformer)
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)
elif self.backbone and self.UNet_Transformer['activate']:
self.transformer_encoder = UNetTransformer(
d_model=self.nPlanes[-1],
nhead=self.UNet_Transformer['multi_heads'],
num_encoder_layers=self.UNet_Transformer['num_encoder_layers'],
dim_feedforward=self.UNet_Transformer['dim_feedforward'],
dropout=self.UNet_Transformer['dropout'],
)
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 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
