def v(depth, nPlanes):
m = scn.Sequential()
if depth == 1:
for _ in range(reps):
res(m, nPlanes, nPlanes, dropout_p)
else:
m = scn.Sequential()
for _ in range(reps):
res(m, nPlanes, nPlanes, dropout_p)
if dropout_width:
m.add(scn.ConcatTable().add(scn.Identity()).add(
scn.Sequential().add(scn.BatchNormReLU(nPlanes)).add(
#In place of Maxpooling
scn.Convolution(
dimension, nPlanes, nPlanes, 2, 2,
False)).add(scn.Dropout(dropout_p)).add(
v(depth - 1, nPlanes)).add(
scn.BatchNormReLU(nPlanes)).add(
scn.Deconvolution(
dimension, nPlanes, nPlanes, 2, 2,
False))))
else:
m.add(scn.ConcatTable().add(scn.Identity()).add(
scn.Sequential().add(scn.BatchNormReLU(nPlanes)).add(
scn.Convolution(dimension, nPlanes, nPlanes, 2, 2,
False)).add(v(depth - 1, nPlanes)).add(
scn.BatchNormReLU(nPlanes)).add(
scn.Deconvolution(
dimension, nPlanes,
nPlanes, 2, 2, False))))
m.add(scn.JoinTable())
for i in range(reps):
res(m, 2 * nPlanes if i == 0 else nPlanes, nPlanes, dropout_p)
return m
def __init__(self, cfg, name='yresnet_decoder'):
super(YResNetDecoder, self).__init__(cfg, name='network_base')
self.model_config = cfg[name]
self.reps = self.model_config.get('reps',
2) # Conv block repetition factor
self.kernel_size = self.model_config.get('kernel_size', 2)
self.num_strides = self.model_config.get('num_strides', 5)
self.num_filters = self.model_config.get('filters', 16)
self.nPlanes = [
i * self.num_filters for i in range(1, self.num_strides + 1)
]
self.downsample = [self.kernel_size, 2] # [filter size, filter stride]
self.concat = scn.JoinTable()
self.add = scn.AddTable()
dropout_prob = self.model_config.get('dropout_prob', 0.5)
self.encoder_num_filters = self.model_config.get(
'encoder_num_filters', None)
if self.encoder_num_filters is None:
self.encoder_num_filters = self.num_filters
self.encoder_nPlanes = [
i * self.encoder_num_filters
for i in range(1, self.num_strides + 1)
]
# Define Sparse YResNet Decoder.
self.decoding_block = scn.Sequential()
self.decoding_conv = scn.Sequential()
for idx, i in enumerate(list(range(self.num_strides - 2, -1, -1))):
if idx == 0:
m = scn.Sequential().add(
scn.BatchNormLeakyReLU(self.encoder_nPlanes[i + 1],
leakiness=self.leakiness)).add(
scn.Deconvolution(
self.dimension,
self.encoder_nPlanes[i + 1],
self.nPlanes[i],
self.downsample[0],
self.downsample[1],
self.allow_bias))
else:
m = scn.Sequential().add(
scn.BatchNormLeakyReLU(
self.nPlanes[i + 1], leakiness=self.leakiness)).add(
scn.Deconvolution(self.dimension,
self.nPlanes[i + 1],
self.nPlanes[i],
self.downsample[0],
self.downsample[1],
self.allow_bias)).add(
scn.Dropout(p=dropout_prob))
self.decoding_conv.add(m)
m = scn.Sequential()
for j in range(self.reps):
self._resnet_block(m, self.nPlanes[i] + (self.encoder_nPlanes[i] \
if j == 0 else 0), self.nPlanes[i])
self.decoding_block.add(m)
def U(nPlanes, n_input_planes=-1): #Recursive function
m = scn.Sequential()
for i in range(reps):
block(m, n_input_planes if n_input_planes != -1 else nPlanes[0],
nPlanes[0])
n_input_planes = -1
if len(nPlanes) > 1:
m.add(scn.ConcatTable().add(scn.Identity()).add(
scn.Sequential().add(
scn.BatchNormLeakyReLU(
nPlanes[0], leakiness=leakiness)).add(
scn.Convolution(dimension, nPlanes[0], nPlanes[1],
downsample[0], downsample[1],
False)).add(U(nPlanes[1:])).add(
scn.BatchNormLeakyReLU(
nPlanes[1],
leakiness=leakiness)).add(
scn.Deconvolution(
dimension,
nPlanes[1],
nPlanes[0],
downsample[0],
downsample[1],
False))))
m.add(scn.JoinTable())
for i in range(reps):
block(m, nPlanes[0] * (2 if i == 0 else 1), nPlanes[0])
return m
def get_upsampler(
num_dims, sparse, input_channels, output_channels,
stride=2, bias=True):
if isinstance(stride, int):
inverse_stride = np.full(num_dims, stride)
else:
inverse_stride = stride
assert len(stride) == num_dims
stride_tuple = tuple(inverse_stride)
if sparse:
layer = scn.Deconvolution(
num_dims, input_channels, output_channels,
filter_size=stride_tuple, filter_stride=stride_tuple, bias=bias)
else:
conv_class = get_dense_transposed_conv_class(num_dims)
layer = conv_class(
input_channels, output_channels, padding=0, bias=bias,
kernel_size=stride_tuple, stride=stride_tuple)
if (1 == inverse_stride).all():
stride = inverse_stride
else:
stride = 1 / inverse_stride
return sparse, stride, output_channels, layer
def __init__(self, *, inplanes, params):
nn.Module.__init__(self)
# The deepest block applies convolutions that act on all three planes together
# First we apply a convolution to map all three planes into 1 plane (of the same spatial size)
self.merger = scn.Convolution(dimension=3,
nIn=inplanes,
nOut=params.bottleneck_deepest,
filter_size=[3, 1, 1],
filter_stride=[1, 1, 1],
bias=params.use_bias)
self.blocks = SparseBlockSeries(inplanes=params.bottleneck_deepest,
n_blocks=params.blocks_deepest_layer,
n_planes=1,
params=params)
self.splitter = scn.Deconvolution(dimension=3,
nIn=params.bottleneck_deepest,
nOut=inplanes,
filter_size=[3, 1, 1],
filter_stride=[1, 1, 1],
bias=params.use_bias)
def __init__(self, inplanes, kernel, stride, bias=False, dim=3):
torch.nn.Module.__init__(self)
outplanes = int(inplanes / 2)
#f1
self.bnr1 = scn.BatchNormReLU(inplanes)
self.deconv1 = scn.Deconvolution(dim, inplanes, outplanes, kernel,
stride, bias)
self.bnr2 = scn.BatchNormReLU(outplanes)
self.subconv = scn.SubmanifoldConvolution(dim, outplanes, outplanes,
kernel, bias)
#f2
self.deconv2 = scn.Deconvolution(dim, inplanes, outplanes, kernel,
stride, bias)
self.add = scn.AddTable()
def decoder(self,
in_channels,
out_channels,
filter_size,
filter_stride=1,
bias=True):
layer = scn.Sequential(
scn.Deconvolution(3, in_channels, out_channels, filter_size,
filter_stride, bias), scn.ReLU())
return layer
def iter_unet(self, n_input_planes):
# different from scn implementation, which is a recursive function
enc_convs = scn.Sequential()
dec_convs = scn.Sequential()
for n_planes_in, n_planes_out in zip(self.n_planes[:-1],
self.n_planes[1:]):
# encode
conv1x1 = scn.Sequential()
for i in range(self.block_reps):
conv1x1.add(
self.block(
n_input_planes
if n_input_planes != -1 else n_planes_in, n_planes_in))
n_input_planes = -1
conv = scn.Sequential()
conv.add(
scn.BatchNormLeakyReLU(n_planes_in, leakiness=self.leakiness))
conv.add(
scn.Convolution(self.dimension, n_planes_in, n_planes_out,
self.downsample[0], self.downsample[1], False))
enc_conv = scn.Sequential()
enc_conv.add(conv1x1)
enc_conv.add(conv)
enc_convs.add(enc_conv)
# decode(corresponding stage of encode; symmetric with U)
b_join = scn.Sequential() # before_join
b_join.add(
scn.BatchNormLeakyReLU(n_planes_out, leakiness=self.leakiness))
b_join.add(
scn.Deconvolution(self.dimension, n_planes_out, n_planes_in,
self.downsample[0], self.downsample[1],
False))
join_table = scn.JoinTable()
a_join = scn.Sequential() # after_join
for i in range(self.block_reps):
a_join.add(
self.block(n_planes_in * (2 if i == 0 else 1),
n_planes_in))
dec_conv = scn.Sequential()
dec_conv.add(b_join)
dec_conv.add(join_table)
dec_conv.add(a_join)
dec_convs.add(dec_conv)
middle_conv = scn.Sequential()
for i in range(self.block_reps):
middle_conv.add(
self.block(
n_input_planes if n_input_planes != -1 else
self.n_planes[-1], self.n_planes[-1]))
n_input_planes = -1
return enc_convs, middle_conv, dec_convs
def __init__(self, *, inplanes, outplanes, nplanes=1, params):
nn.Module.__init__(self)
self.conv = scn.Deconvolution(dimension=3,
nIn=inplanes,
nOut=outplanes,
filter_size=[nplanes, 2, 2],
filter_stride=[1, 2, 2],
bias=params.use_bias)
self.do_batch_norm = False
if params.batch_norm:
self.do_batch_norm = True
self.bn = scn.BatchNormalization(outplanes)
self.relu = scn.ReLU()
def make_decoder_layer(self,
ilayer,
ninputchs,
noutputchs,
nreps,
leakiness=0.01,
downsample=[2, 2],
islast=False):
"""
defines two layers:
1) the deconv layer pre-concat
2) residual blocks post-concat
inputs
------
ilayer [int]: layer ID
ninputchs [int]: number of features going into layer
noutputchs [int]: number of features output by layer
nreps [int]: number of times residual modules should repeat
leakiness [float]: leakiness of LeakyReLU activiation functions
downsample [list of ints size 2]: upsampling factor in weight and height
islast [bool]: last decoder layer does not have skip connection
"""
# resnet block
decode_blocks = create_resnet_layer(nreps,
ninputchs,
2 * noutputchs,
downsample=downsample)
# deconv
decode_blocks.add(
scn.BatchNormLeakyReLU(2 * noutputchs, leakiness=leakiness))
decode_blocks.add(
scn.Deconvolution(self.dimension, 2 * noutputchs, noutputchs,
downsample[0], downsample[1], False))
setattr(self, "deconv%d" % (ilayer), decode_blocks)
if self._verbose:
print "DecoderLayer[", ilayer, "] inputchs[", ninputchs,
print " -> resout[", 2 * noutputchs, "] -> deconv output[", noutputchs, "]"
if not islast:
# joiner for skip connections
joiner = scn.JoinTable()
setattr(self, "skipjoin%d" % (ilayer), joiner)
else:
joiner = None
return decode_blocks, joiner
def up(m, nPlane_in, nPlane_uped, scale):
#print(f'up, scale={scale}, feature={nPlane_in}->{nPlane_uped}, kernel={self.down_kernels[scale]}, stride={self.down_strides[scale]}')
m.add(
scn.BatchNormLeakyReLU(
nPlane_in,
momentum=bn_momentum,
leakiness=leakiness,
track_running_stats=track_running_stats)).add(
scn.Deconvolution(dimension, nPlane_in, nPlane_uped,
self.down_kernels[scale],
self.down_strides[scale], False))
operation = {
'kernel': self.down_kernels[scale],
'stride': self.down_strides[scale]
}
return operation
def __init__(self, num_layers, ndim, shape, in_channels, out_channels,
kernel_size, stride):
super().__init__()
self.scn_input = scn.InputLayer(ndim, shape, mode=0)
self.net = nn.Sequential(
scn.Convolution(ndim,
in_channels,
out_channels,
kernel_size,
stride,
bias=False),
scn.Deconvolution(ndim,
out_channels,
in_channels,
kernel_size,
stride,
bias=False),
scn.SparseToDense(ndim, in_channels),
)
def make_decoder_layer(self,
ilayer,
ninputchs,
noutputchs,
nreps,
leakiness=0.01,
downsample=[2, 2],
islast=False):
"""
defines two layers:
1) the deconv layer pre-concat
2) residual blocks post-concat
inputs
------
ninputchs: number of features going into layer
noutputchs: number of features output by layer
"""
# resnet block
decode_blocks = create_resnet_layer(nreps,
ninputchs,
2 * noutputchs,
downsample=downsample)
# deconv
decode_blocks.add(
scn.BatchNormLeakyReLU(2 * noutputchs, leakiness=leakiness))
decode_blocks.add(
scn.Deconvolution(self.dimension, 2 * noutputchs, noutputchs,
downsample[0], downsample[1], False))
setattr(self, "deconv%d" % (ilayer), decode_blocks)
if self._verbose:
print "DecoderLayer[", ilayer, "] inputchs[", ninputchs,
print " -> resout[", 2 * noutputchs, "] -> deconv output[", noutputchs, "]"
if not islast:
# joiner for skip connections
joiner = scn.JoinTable()
setattr(self, "skipjoin%d" % (ilayer), joiner)
else:
joiner = None
return decode_blocks, joiner
def resnet_block(dimension,
n_in,
n_out,
kernel,
leakiness=0,
computation='convolution'):
'''Build and return ResNet block
'''
assert computation in [
'submanifoldconvolution', 'convolution', 'fullconvolution',
'deconvolution'
]
if computation == 'convolution':
computation = lambda n_in, n_out: scn.Convolution(
dimension, n_in, n_out, kernel[0], kernel[1], False)
elif computation == 'submanifoldconvolution':
assert type(
kernel
) == int, f"`kernel` must be int, {type(kernel)} was provided"
computation = lambda n_in, n_out: scn.SubmanifoldConvolution(
dimension, n_in, n_out, kernel, False)
elif computation == 'deconvolution':
assert type(
kernel
) == int, f"`kernel` must be int, {type(kernel)} was provided"
computation = lambda n_in, n_out: scn.Deconvolution(
dimension, n_in, n_out, kernel, kernel, False)
else:
computation = lambda n_in, n_out: scn.FullConvolution(
dimension, n_in, n_out, kernel[0], kernel[1], False)
block = scn.Sequential()
block.add(scn.ConcatTable(
).add(scn.NetworkInNetwork(n_in, n_out, False)).add(scn.Sequential().add(
# scn.BatchNormLeakyReLU(n_in, leakiness=leakiness)
scn.LeakyReLU(leakiness)).add(computation(n_in, n_out)).add(
# scn.BatchNormLeakyReLU(n_out, leakiness=leakiness)
scn.LeakyReLU(leakiness)).add(computation(n_out, n_out)))).add(
scn.AddTable())
return block
def __init__(self,
dimension,
nPlanes_down,
nPlanes_up,
reps,
type,
predict=False,
nClasses=None,
extract=False,
kernel_size=None,
group=1,
sgc_config=None):
super(up, self).__init__()
self.dimension = dimension
self.type = type
self.predict = predict
self.extract = extract
self.kernel_size = kernel_size
self.group = group
self.sgc_config = sgc_config
if type == 'c':
self.upsample = scn.Sequential().add(
scn.BatchNormReLU(nPlanes_down)).add(
scn.DenseDeconvolution(dimension, nPlanes_down, nPlanes_up,
2, 2, False))
elif type == 'v':
self.upsample = scn.Sequential().add(
scn.BatchNormReLU(nPlanes_down)).add(
scn.Deconvolution(dimension, nPlanes_down, nPlanes_up, 2,
2, False))
self.conv = scn.Sequential()
for i in range(reps):
res(self.conv, dimension, nPlanes_up, nPlanes_up)
if predict:
self.Linear = scn.Sequential().add(
scn.BatchNormReLU(nPlanes_up)).add(
scn.Linear(dimension, nPlanes_up, nClasses))
def U(nPlanes): #Recursive function
m = scn.Sequential()
if len(nPlanes) == 1:
for _ in range(reps):
block(m, nPlanes[0], nPlanes[0])
else:
m = scn.Sequential()
for _ in range(reps):
block(m, nPlanes[0], nPlanes[0])
m.add(scn.ConcatTable().add(scn.Identity()).add(
scn.Sequential().add(scn.BatchNormReLU(nPlanes[0])).add(
scn.Convolution(dimension, nPlanes[0], nPlanes[1],
downsample[0], downsample[1],
False)).add(U(nPlanes[1:])).add(
scn.BatchNormReLU(nPlanes[1])).add(
scn.Deconvolution(
dimension, nPlanes[1],
nPlanes[0], downsample[0],
downsample[1], False))))
m.add(scn.JoinTable())
for i in range(reps):
block(m, nPlanes[0] * (2 if i == 0 else 1), nPlanes[0])
return m
def __init__(self):
super(CRinGeNetSparse, self).__init__()
self._mlp_pid = torch.nn.Sequential(torch.nn.Linear(3, 512),
torch.nn.ReLU(),
torch.nn.BatchNorm1d(512),
torch.nn.Linear(512, 512),
torch.nn.ReLU(),
torch.nn.BatchNorm1d(512))
self._mlp_pos = torch.nn.Sequential(torch.nn.Linear(3, 512),
torch.nn.ReLU(),
torch.nn.BatchNorm1d(512),
torch.nn.Linear(512, 512),
torch.nn.ReLU(),
torch.nn.BatchNorm1d(512))
self._mlp_dir = torch.nn.Sequential(torch.nn.Linear(3, 512),
torch.nn.ReLU(),
torch.nn.BatchNorm1d(512),
torch.nn.Linear(512, 512),
torch.nn.ReLU(),
torch.nn.BatchNorm1d(512))
self._mlp_E = torch.nn.Sequential(torch.nn.Linear(1, 512),
torch.nn.ReLU(),
torch.nn.BatchNorm1d(512),
torch.nn.Linear(512, 512),
torch.nn.ReLU(),
torch.nn.BatchNorm1d(512))
self._mlp = torch.nn.Sequential(torch.nn.Linear(2048, 1024),
torch.nn.ReLU(),
torch.nn.BatchNorm1d(1024),
torch.nn.Linear(1024, 1024),
torch.nn.ReLU(),
torch.nn.BatchNorm1d(1024),
torch.nn.Linear(1024, 14784),
torch.nn.ReLU(),
torch.nn.BatchNorm1d(14784))
self._upconvs = torch.nn.Sequential(
# torch.nn.ConvTranspose2d(64, 64, 4, 2), torch.nn.ReLU(), torch.nn.BatchNorm2d(64),
scn.Deconvolution(2, 64, 64, 4, 2, 0.),
torch.nn.ReLU(),
torch.nn.BatchNorm2d(64),
# Should be submanifold convo?!
scn.Convolution(2, 64, 64, 3, 1, 0.),
torch.nn.ReLU(),
torch.nn.BatchNorm2d(64),
# torch.nn.Conv2d(64, 64, 3), torch.nn.ReLU(),torch.nn.BatchNorm2d(64),
# torch.nn.ConvTranspose2d(64, 32, 4, 2), torch.nn.ReLU(), torch.nn.BatchNorm2d(32),
scn.Deconvolution(2, 64, 32, 4, 2, 0.),
torch.nn.ReLU(),
torch.nn.BatchNorm2d(64),
scn.Convolution(2, 32, 32, 3, 1, 0.),
torch.nn.ReLU(),
torch.nn.BatchNorm2d(64),
# torch.nn.Conv2d(32, 32, 3), torch.nn.ReLU(), torch.nn.BatchNorm2d(32),
scn.Deconvolution(2, 32, 32, 4, 2, 0.),
torch.nn.ReLU(),
torch.nn.BatchNorm2d(32),
# torch.nn.ConvTranspose2d(32, 32, 4, 2), torch.nn.ReLU(), torch.nn.BatchNorm2d(32),
scn.Convolution(2, 32, 1, 3, 1, 0.),
torch.nn.ReLU()
# torch.nn.Conv2d(32, 1, 3)
)
def __init__(self,
output_shape,
use_norm=True,
num_filters_down1=[32, 64, 96, 128],
num_filters_down2=[32, 64, 96, 128],
name='tDBN_bv_2'):
super(tDBN_bv_2, self).__init__()
self.name = name
if use_norm:
BatchNorm1d = change_default_args(eps=1e-3,
momentum=0.01)(nn.BatchNorm1d)
Linear = change_default_args(bias=False)(nn.Linear)
else:
BatchNorm1d = Empty
Linear = change_default_args(bias=True)(nn.Linear)
sparse_shape = np.array(output_shape[1:4]) # + [1, 0, 0]
# sparse_shape[0] = 11
# print(sparse_shape)
self.scn_input = scn.InputLayer(3, sparse_shape.tolist())
self.voxel_output_shape = output_shape
To_use_bias = False
residual_use = True # using residual block or not
dimension = 3
reps = 2
dimension = 3
leakiness = 0
input_filters_layers = num_filters_down1[:
4] # feature channels in the raw data.
num_filter_fpn = num_filters_down1[
3:] # [ 64, 128, 256, 512] #dimension of feature maps, num_filter_fpn[3] == dimension_feature_map[3]
dimension_feature_map = num_filters_down2 # [ 64, 128, 256, 512] # dimensions of output into 2D feature map
dimension_kernel_size = [15, 7, 3, 1]
filters_input_pairs = [[
input_filters_layers[i], input_filters_layers[i + 1]
] for i in range(len(input_filters_layers) - 1)]
m = None
m = scn.Sequential()
# -----------------------------------------------------------------
## block1 and feature map 0, convert from voxel into 3D tensor
# -----------------------------------------------------------------
for i, o in [[1, input_filters_layers[0]]]:
m.add(scn.SubmanifoldConvolution(3, i, o, 3, False))
for i, o in filters_input_pairs: # , [num_filter_fpn[0], num_filter_fpn[0]]]:
for _ in range(reps):
self.block(m, i, i, residual_blocks=residual_use)
m.add(scn.BatchNormLeakyReLU(i, leakiness=leakiness)).add(
scn.Convolution(dimension, i, o, 3, 2, False))
self.block_input = m
middle_layers = []
m = None
m = scn.Sequential()
for _ in range(reps):
self.block(m,
num_filter_fpn[0],
num_filter_fpn[0],
residual_blocks=residual_use)
self.x0_in = m
for k in range(1, 4):
m = None
m = scn.Sequential()
# downsample
m.add(
scn.BatchNormLeakyReLU(
num_filter_fpn[k - 1], leakiness=leakiness)).add(
scn.Convolution(dimension, num_filter_fpn[k - 1],
num_filter_fpn[k], 3, 2, False))
# cnn
for _ in range(reps):
if k == 4:
self.block(m,
num_filter_fpn[k],
num_filter_fpn[k],
dimension=2,
residual_blocks=residual_use
) ## it has be compressed into 2 dimensions
else:
self.block(m,
num_filter_fpn[k],
num_filter_fpn[k],
dimension=3,
residual_blocks=residual_use)
if k == 1:
self.x1_in = m
if k == 2:
self.x2_in = m
if k == 3:
self.x3_in = m
#self.feature_map3 = Sequential(*middle_layers) # XXX
self.feature_map3 = scn.Sequential(
scn.BatchNormLeakyReLU(num_filter_fpn[3],
leakiness=leakiness)).add(
scn.SparseToDense(3, num_filter_fpn[3])
) ## last one is the 2D instead of 3D
for k in range(2, -1, -1):
m = None
m = scn.Sequential()
# upsample
m.add(
scn.BatchNormLeakyReLU(
num_filter_fpn[k + 1], leakiness=leakiness)).add(
scn.Deconvolution(dimension, num_filter_fpn[k + 1],
num_filter_fpn[k], 3, 2, False))
if k == 2:
self.upsample32 = m
if k == 1:
self.upsample21 = m
if k == 0:
self.upsample10 = m
m = None
m = scn.Sequential()
m.add(scn.JoinTable())
for i in range(reps):
self.block(m,
num_filter_fpn[k] * (2 if i == 0 else 1),
num_filter_fpn[k],
residual_blocks=residual_use)
if k == 2:
self.concate2 = m
if k == 1:
self.concate1 = m
if k == 0:
self.concate0 = m
m = None
m = scn.Sequential()
m.add(
scn.BatchNormLeakyReLU(
num_filter_fpn[k], leakiness=leakiness)).add(
scn.Convolution(
3,
num_filter_fpn[k],
dimension_feature_map[k],
(dimension_kernel_size[k], 1, 1), (1, 1, 1),
bias=False)).add(
scn.BatchNormReLU(
dimension_feature_map[k],
eps=1e-3,
momentum=0.99)).add(
scn.SparseToDense(
3, dimension_feature_map[k]))
if k == 2:
self.feature_map2 = m
if k == 1:
self.feature_map1 = m
if k == 0:
self.feature_map0 = m
def __init__(self, cfg, name="uresnet_lonely"):
super(UResNet, self).__init__()
import sparseconvnet as scn
if 'modules' in cfg:
self.model_config = cfg['modules'][name]
else:
self.model_config = cfg
# Whether to compute ghost mask separately or not
self._dimension = self.model_config.get('data_dim', 3)
reps = self.model_config.get('reps', 2) # Conv block repetition factor
kernel_size = self.model_config.get('kernel_size', 2)
num_strides = self.model_config.get('num_strides', 5)
m = self.model_config.get('filters', 16) # Unet number of features
nInputFeatures = self.model_config.get('features', 1)
spatial_size = self.model_config.get('spatial_size', 512)
leakiness = self.model_config.get('leak', 0.0)
nPlanes = [i * m for i in range(1, num_strides + 1)
] # UNet number of features per level
print("nPlanes: ", nPlanes)
downsample = [kernel_size, 2] # [filter size, filter stride]
self.last = None
def block(m, a, b): # ResNet style blocks
m.add(scn.ConcatTable().add(scn.Identity(
) if a == b else scn.NetworkInNetwork(a, b, False)).add(
scn.Sequential().add(
scn.BatchNormLeakyReLU(a, leakiness=leakiness)).add(
scn.SubmanifoldConvolution(
self._dimension, a, b, 3, False)).add(
scn.BatchNormLeakyReLU(
b, leakiness=leakiness)).add(
scn.SubmanifoldConvolution(
self._dimension, b, b, 3,
False)))).add(scn.AddTable())
self.input = scn.Sequential().add(
scn.InputLayer(self._dimension, spatial_size, mode=3)).add(
scn.SubmanifoldConvolution(self._dimension, nInputFeatures, m,
3, False)) # Kernel size 3, no bias
self.concat = scn.JoinTable()
# Encoding
self.bn = scn.BatchNormLeakyReLU(nPlanes[0], leakiness=leakiness)
self.encoding_block = scn.Sequential()
self.encoding_conv = scn.Sequential()
module = scn.Sequential()
for i in range(num_strides):
module = scn.Sequential()
for _ in range(reps):
block(module, nPlanes[i], nPlanes[i])
self.encoding_block.add(module)
module2 = scn.Sequential()
if i < num_strides - 1:
module2.add(
scn.BatchNormLeakyReLU(
nPlanes[i], leakiness=leakiness)).add(
scn.Convolution(self._dimension, nPlanes[i],
nPlanes[i + 1], downsample[0],
downsample[1], False))
self.encoding_conv.add(module2)
self.encoding = module
# Decoding
self.decoding_conv, self.decoding_blocks = scn.Sequential(
), scn.Sequential()
for i in range(num_strides - 2, -1, -1):
module1 = scn.Sequential().add(
scn.BatchNormLeakyReLU(nPlanes[i + 1],
leakiness=leakiness)).add(
scn.Deconvolution(
self._dimension, nPlanes[i + 1],
nPlanes[i], downsample[0],
downsample[1], False))
self.decoding_conv.add(module1)
module2 = scn.Sequential()
for j in range(reps):
block(module2, nPlanes[i] * (2 if j == 0 else 1), nPlanes[i])
self.decoding_blocks.add(module2)
self.output = scn.Sequential().add(scn.BatchNormReLU(m)).add(
scn.OutputLayer(self._dimension))
def __init__(self, cfg, name="uresnet_clustering"):
super(UResNet, self).__init__()
import sparseconvnet as scn
self._model_config = cfg['modules'][name]
# Whether to compute ghost mask separately or not
self._ghost = self._model_config.get('ghost', False)
self._dimension = self._model_config.get('data_dim', 3)
reps = self._model_config.get('reps',
2) # Conv block repetition factor
kernel_size = self._model_config.get('kernel_size', 2)
num_strides = self._model_config.get('num_strides', 5)
m = self._model_config.get('filters', 16) # Unet number of features
nInputFeatures = self._model_config.get('features', 1)
spatial_size = self._model_config.get('spatial_size', 512)
num_classes = self._model_config.get('num_classes', 5)
self._N = self._model_config.get('num_cluster_conv', 0)
self._simpleN = self._model_config.get('simple_conv', True)
self._add_coordinates = self._model_config.get('cluster_add_coords',
False)
self._density_estimate = self._model_config.get(
'density_estimate', False)
nPlanes = [i * m for i in range(1, num_strides + 1)
] # UNet number of features per level
downsample = [kernel_size, 2] # [filter size, filter stride]
self.last = None
leakiness = 0
def block(m, a, b): # ResNet style blocks
m.add(scn.ConcatTable().add(scn.Identity(
) if a == b else scn.NetworkInNetwork(a, b, False)).add(
scn.Sequential().add(
scn.BatchNormLeakyReLU(a, leakiness=leakiness)).add(
scn.SubmanifoldConvolution(
self._dimension, a, b, 3, False)).add(
scn.BatchNormLeakyReLU(
b, leakiness=leakiness)).add(
scn.SubmanifoldConvolution(
self._dimension, b, b, 3,
False)))).add(scn.AddTable())
self.input = scn.Sequential().add(
scn.InputLayer(self._dimension, spatial_size, mode=3)).add(
scn.SubmanifoldConvolution(self._dimension, nInputFeatures, m,
3, False)) # Kernel size 3, no bias
self.concat = scn.JoinTable()
# Encoding
self.bn = scn.BatchNormLeakyReLU(nPlanes[0], leakiness=leakiness)
self.encoding_block = scn.Sequential()
self.encoding_conv = scn.Sequential()
module = scn.Sequential()
for i in range(num_strides):
module = scn.Sequential()
for _ in range(reps):
block(module, nPlanes[i], nPlanes[i])
self.encoding_block.add(module)
module2 = scn.Sequential()
if i < num_strides - 1:
module2.add(
scn.BatchNormLeakyReLU(
nPlanes[i], leakiness=leakiness)).add(
scn.Convolution(self._dimension, nPlanes[i],
nPlanes[i + 1], downsample[0],
downsample[1], False))
self.encoding_conv.add(module2)
self.encoding = module
# Decoding
self.decoding_conv, self.decoding_blocks = scn.Sequential(
), scn.Sequential()
for i in range(num_strides - 2, -1, -1):
inFeatures = nPlanes[i + 1] * (2 if
(self._N > 0
and i < num_strides - 2) else 1)
module1 = scn.Sequential().add(
scn.BatchNormLeakyReLU(inFeatures, leakiness=leakiness)).add(
scn.Deconvolution(self._dimension, inFeatures, nPlanes[i],
downsample[0], downsample[1], False))
self.decoding_conv.add(module1)
module2 = scn.Sequential()
for j in range(reps):
block(module2, nPlanes[i] * (2 if j == 0 else 1), nPlanes[i])
self.decoding_blocks.add(module2)
# Clustering convolutions
if self._N > 0:
self.clustering_conv = scn.Sequential()
for i in range(num_strides - 2, -1, -1):
conv = scn.Sequential()
for j in range(self._N):
if self._simpleN:
conv.add(
scn.SubmanifoldConvolution(
self._dimension, nPlanes[i] +
(4 if j == 0 and self._add_coordinates else 0),
nPlanes[i], 3, False))
conv.add(
scn.BatchNormLeakyReLU(nPlanes[i],
leakiness=leakiness))
else:
block(
conv, nPlanes[i] +
(4 if j == 0 and self._add_coordinates else 0),
nPlanes[i])
self.clustering_conv.add(conv)
outFeatures = m * (2 if self._N > 0 else 1)
self.output = scn.Sequential().add(scn.BatchNormReLU(outFeatures)).add(
scn.OutputLayer(self._dimension))
self.linear = torch.nn.Linear(outFeatures, num_classes)
if self._density_estimate:
self._density_layer = []
for i in range(num_strides - 2, -1, -1):
self._density_layer.append(torch.nn.Linear(nPlanes[i], 2))
self._density_layer = torch.nn.Sequential(*self._density_layer)
def __init__(self, dimension, n_in, n_out, filter_size, filter_stride, bias):
super().__init__()
self.convolution = scn.Convolution(dimension, n_in, n_out, filter_size, filter_stride, bias)
self.deconvolution = scn.Deconvolution(dimension, n_in, n_out, filter_size, filter_stride, bias)
def SparseVggNet(dimension, nInputPlanes, layers):
"""
VGG style nets
Use submanifold convolutions
Also implements 'Plus'-augmented nets
"""
nPlanes = nInputPlanes
m = scn.Sequential()
for x in layers:
if x == 'MP':
m.add(scn.MaxPooling(dimension, 3, 2))
elif x[0] == 'MP':
m.add(scn.MaxPooling(dimension, x[1], x[2]))
elif x == 'C3/2':
m.add(scn.Convolution(dimension, nPlanes, nPlanes, 3, 2, False))
m.add(scn.BatchNormReLU(nPlanes))
elif x[0] == 'C3/2':
m.add(scn.Convolution(dimension, nPlanes, x[1], 3, 2, False))
nPlanes = x[1]
m.add(scn.BatchNormReLU(nPlanes))
elif x[0] == 'C' and len(x) == 2:
m.add(
scn.SubmanifoldConvolution(dimension, nPlanes, x[1], 3, False))
nPlanes = x[1]
m.add(scn.BatchNormReLU(nPlanes))
elif x[0] == 'C' and len(x) == 3:
m.add(scn.ConcatTable().add(
scn.SubmanifoldConvolution(
dimension, nPlanes, x[1], 3,
False)).add(scn.Sequential().add(
scn.Convolution(
dimension, nPlanes, x[2], 3, 2,
False)).add(scn.BatchNormReLU(x[2])).add(
scn.SubmanifoldConvolution(
dimension, x[2], x[2], 3,
False)).add(scn.BatchNormReLU(x[2])).add(
scn.Deconvolution(
dimension, x[2], x[2], 3, 2,
False)))).add(scn.JoinTable())
nPlanes = x[1] + x[2]
m.add(scn.BatchNormReLU(nPlanes))
elif x[0] == 'C' and len(x) == 4:
m.add(scn.ConcatTable().add(
scn.SubmanifoldConvolution(
dimension, nPlanes, x[1], 3,
False)).add(scn.Sequential().add(
scn.Convolution(
dimension, nPlanes, x[2], 3, 2,
False)).add(scn.BatchNormReLU(x[2])).add(
scn.SubmanifoldConvolution(
dimension, x[2], x[2], 3,
False)).add(scn.BatchNormReLU(x[2])).add(
scn.Deconvolution(
dimension, x[2], x[2], 3, 2,
False))).
add(scn.Sequential().add(
scn.Convolution(
dimension, nPlanes, x[3],
3, 2, False)).add(scn.BatchNormReLU(x[3])).add(
scn.SubmanifoldConvolution(
dimension, x[3], x[3], 3,
False)).add(scn.BatchNormReLU(x[3])).add(
scn.Convolution(dimension, x[3], x[3], 3,
2, False)).
add(scn.BatchNormReLU(x[3])).add(
scn.SubmanifoldConvolution(
dimension, x[3], x[3],
3, False)).add(scn.BatchNormReLU(x[3])).add(
scn.Deconvolution(
dimension, x[3], x[3], 3, 2,
False)).add(scn.BatchNormReLU(x[3])).add(
scn.SubmanifoldConvolution(
dimension, x[3], x[3], 3,
False)).add(
scn.BatchNormReLU(x[3])).add(
scn.Deconvolution(
dimension, x[3],
x[3], 3, 2,
False)))).add(
scn.JoinTable())
nPlanes = x[1] + x[2] + x[3]
m.add(scn.BatchNormReLU(nPlanes))
elif x[0] == 'C' and len(x) == 5:
m.add(scn.ConcatTable().add(
scn.SubmanifoldConvolution(
dimension, nPlanes, x[1], 3,
False)).add(scn.Sequential().add(
scn.Convolution(
dimension, nPlanes, x[2], 3, 2,
False)).add(scn.BatchNormReLU(x[2])).add(
scn.SubmanifoldConvolution(
dimension, x[2], x[2], 3,
False)).add(scn.BatchNormReLU(x[2])).add(
scn.Deconvolution(
dimension, x[2], x[2], 3, 2,
False))).
add(scn.Sequential().add(
scn.Convolution(
dimension, nPlanes, x[3],
3, 2, False)).add(scn.BatchNormReLU(x[3])).add(
scn.SubmanifoldConvolution(
dimension, x[3], x[3], 3,
False)).add(scn.BatchNormReLU(x[3])).add(
scn.Convolution(dimension, x[3], x[3], 3,
2, False)).
add(scn.BatchNormReLU(x[3])).add(
scn.SubmanifoldConvolution(
dimension, x[3], x[3],
3, False)).add(scn.BatchNormReLU(x[3])).add(
scn.Deconvolution(
dimension, x[3], x[3], 3, 2,
False)).add(scn.BatchNormReLU(x[3])).add(
scn.SubmanifoldConvolution(
dimension, x[3], x[3], 3,
False)).add(
scn.BatchNormReLU(x[3])).add(
scn.Deconvolution(
dimension, x[3],
x[3], 3, 2, False))).
add(scn.Sequential().add(
scn.Convolution(
dimension, nPlanes,
x[4], 3, 2,
False)).add(scn.BatchNormReLU(x[4])).add(
scn.SubmanifoldConvolution(
dimension, x[4], x[4], 3,
False)).add(scn.BatchNormReLU(x[4])).add(
scn.Convolution(
dimension, x[4], x[4], 3,
2, False)).add(
scn.BatchNormReLU(x[4])).add(
scn.SubmanifoldConvolution(
dimension, x[4], x[4], 3,
False)).add(
scn.BatchNormReLU(
x[4])).
add(scn.Convolution(
dimension, x[4], x[4], 3, 2,
False)).add(scn.BatchNormReLU(x[4])).add(
scn.SubmanifoldConvolution(
dimension, x[4], x[4], 3,
False)).add(scn.BatchNormReLU(x[4])).add(
scn.Deconvolution(
dimension, x[4], x[4], 3,
2, False)).add(
scn.BatchNormReLU(x[4])).add(
scn.SubmanifoldConvolution(
dimension, x[4], x[4], 3,
False)).
add(scn.BatchNormReLU(x[4])).add(
scn.Deconvolution(
dimension, x[4], x[4], 3,
2, False)).add(scn.BatchNormReLU(x[4])).add(
scn.SubmanifoldConvolution(
dimension, x[4], x[4], 3,
False)).add(scn.BatchNormReLU(x[4])).add(
scn.Deconvolution(
dimension, x[4], x[4], 3, 2,
False)))).add(scn.JoinTable())
nPlanes = x[1] + x[2] + x[3] + x[4]
m.add(scn.BatchNormReLU(nPlanes))
return m
def __init__(self, cfg):
super(UResNet, self).__init__()
import sparseconvnet as scn
model_config = cfg['modules']['uresnet_lonely']
self._model_config = model_config
dimension = model_config['data_dim']
reps = 2 # Conv block repetition factor
kernel_size = 2 # Use input_spatial_size method for other values?
m = model_config['filters'] # Unet number of features
nPlanes = [i * m for i in range(1, model_config['num_strides'] + 1)
] # UNet number of features per level
# nPlanes = [(2**i) * m for i in range(1, num_strides+1)] # UNet number of features per level
nInputFeatures = 1
downsample = [kernel_size,
2] # downsample = [filter size, filter stride]
self.last = None
leakiness = 0
def block(m, a, b):
# ResNet style blocks
m.add(scn.ConcatTable().add(scn.Identity(
) if a == b else scn.NetworkInNetwork(a, b, False)).add(
scn.Sequential().add(
scn.BatchNormLeakyReLU(a, leakiness=leakiness)).add(
scn.SubmanifoldConvolution(
dimension, a, b, 3, False)).add(
scn.BatchNormLeakyReLU(
b, leakiness=leakiness)).add(
scn.SubmanifoldConvolution(
dimension, b, b, 3,
False)))).add(scn.AddTable())
self.input = scn.Sequential().add(
scn.InputLayer(dimension, model_config['spatial_size'],
mode=3)).add(
scn.SubmanifoldConvolution(
dimension, nInputFeatures, m, 3,
False)) # Kernel size 3, no bias
self.concat = scn.JoinTable()
# Encoding
self.bn = scn.BatchNormLeakyReLU(nPlanes[0], leakiness=leakiness)
# self.encoding = []
self.encoding_block = scn.Sequential()
self.encoding_conv = scn.Sequential()
module = scn.Sequential()
for i in range(model_config['num_strides']):
module = scn.Sequential()
for _ in range(reps):
block(module, nPlanes[i], nPlanes[i])
self.encoding_block.add(module)
module2 = scn.Sequential()
if i < model_config['num_strides'] - 1:
module2.add(
scn.BatchNormLeakyReLU(
nPlanes[i], leakiness=leakiness)).add(
scn.Convolution(dimension, nPlanes[i],
nPlanes[i + 1], downsample[0],
downsample[1], False))
# self.encoding.append(module)
self.encoding_conv.add(module2)
self.encoding = module
# Decoding
self.decoding_conv, self.decoding_blocks = scn.Sequential(
), scn.Sequential()
for i in range(model_config['num_strides'] - 2, -1, -1):
module1 = scn.Sequential().add(
scn.BatchNormLeakyReLU(nPlanes[i + 1],
leakiness=leakiness)).add(
scn.Deconvolution(
dimension, nPlanes[i + 1],
nPlanes[i], downsample[0],
downsample[1], False))
self.decoding_conv.add(module1)
module2 = scn.Sequential()
for j in range(reps):
block(module2, nPlanes[i] * (2 if j == 0 else 1), nPlanes[i])
self.decoding_blocks.add(module2)
self.output = scn.Sequential().add(scn.BatchNormReLU(m)).add(
scn.OutputLayer(dimension))
self.linear = torch.nn.Linear(m, model_config['num_classes'])
def __init__(self, cfg):
super(PPNUResNet, self).__init__()
import sparseconvnet as scn
self._model_config = cfg['modules']['uresnet_ppn_type']
self._dimension = self._model_config.get('data_dim', 3)
nInputFeatures = self._model_config.get('features', 1)
spatial_size = self._model_config.get('spatial_size', 512)
num_classes = self._model_config.get('num_classes', 5)
m = self._model_config.get('filters', 16) # Unet number of features
num_strides = self._model_config.get('num_strides', 5)
reps = 2 # Conv block repetition factor
kernel_size = 2 # Use input_spatial_size method for other values?
nPlanes = [i * m for i in range(1, num_strides + 1)
] # UNet number of features per level
# nPlanes = [(2**i) * m for i in range(1, num_strides+1)] # UNet number of features per level
downsample = [kernel_size,
2] # downsample = [filter size, filter stride]
self.last = None
leakiness = 0
def block(m, a, b):
# ResNet style blocks
m.add(scn.ConcatTable().add(scn.Identity(
) if a == b else scn.NetworkInNetwork(a, b, False)).add(
scn.Sequential().add(
scn.BatchNormLeakyReLU(a, leakiness=leakiness)).add(
scn.SubmanifoldConvolution(
self._dimension, a, b, 3, False)).add(
scn.BatchNormLeakyReLU(
b, leakiness=leakiness)).add(
scn.SubmanifoldConvolution(
self._dimension, b, b, 3,
False)))).add(scn.AddTable())
self.input = scn.Sequential().add(
scn.InputLayer(self._dimension, spatial_size, mode=3)).add(
scn.SubmanifoldConvolution(self._dimension, nInputFeatures, m,
3, False)) # Kernel size 3, no bias
self.concat = scn.JoinTable()
# Encoding
self.bn = scn.BatchNormLeakyReLU(nPlanes[0], leakiness=leakiness)
# self.encoding = []
self.encoding_block = scn.Sequential()
self.encoding_conv = scn.Sequential()
module = scn.Sequential()
for i in range(num_strides):
module = scn.Sequential()
for _ in range(reps):
block(module, nPlanes[i], nPlanes[i])
self.encoding_block.add(module)
module2 = scn.Sequential()
if i < num_strides - 1:
module2.add(
scn.BatchNormLeakyReLU(
nPlanes[i], leakiness=leakiness)).add(
scn.Convolution(self._dimension, nPlanes[i],
nPlanes[i + 1], downsample[0],
downsample[1], False))
# self.encoding.append(module)
self.encoding_conv.add(module2)
self.encoding = module
# Decoding
self.decoding_conv, self.decoding_blocks = scn.Sequential(
), scn.Sequential()
for i in range(num_strides - 2, -1, -1):
module1 = scn.Sequential().add(
scn.BatchNormLeakyReLU(nPlanes[i + 1],
leakiness=leakiness)).add(
scn.Deconvolution(
self._dimension, nPlanes[i + 1],
nPlanes[i], downsample[0],
downsample[1], False))
self.decoding_conv.add(module1)
module2 = scn.Sequential()
for j in range(reps):
block(module2, nPlanes[i] * (2 if j == 0 else 1), nPlanes[i])
self.decoding_blocks.add(module2)
self.output = scn.Sequential().add(scn.BatchNormReLU(m)).add(
scn.OutputLayer(self._dimension))
self.linear = torch.nn.Linear(m, num_classes)
# PPN stuff
self.half_stride = int(num_strides / 2.0)
self.ppn1_conv = scn.SubmanifoldConvolution(self._dimension,
nPlanes[-1], nPlanes[-1],
3, False)
self.ppn1_scores = scn.SubmanifoldConvolution(self._dimension,
nPlanes[-1], 2, 3, False)
self.selection1 = Selection()
self.selection2 = Selection()
self.unpool1 = scn.Sequential()
for i in range(num_strides - self.half_stride - 1):
self.unpool1.add(
scn.UnPooling(self._dimension, downsample[0], downsample[1]))
self.unpool2 = scn.Sequential()
for i in range(self.half_stride):
self.unpool2.add(
scn.UnPooling(self._dimension, downsample[0], downsample[1]))
middle_filters = int(m * self.half_stride * (self.half_stride + 1) /
2.0)
self.ppn2_conv = scn.SubmanifoldConvolution(self._dimension,
middle_filters,
middle_filters, 3, False)
self.ppn2_scores = scn.SubmanifoldConvolution(self._dimension,
middle_filters, 2, 3,
False)
self.multiply1 = Multiply()
self.multiply2 = Multiply()
self.ppn3_conv = scn.SubmanifoldConvolution(self._dimension,
nPlanes[0], nPlanes[0], 3,
False)
self.ppn3_pixel_pred = scn.SubmanifoldConvolution(
self._dimension, nPlanes[0], self._dimension, 3, False)
self.ppn3_scores = scn.SubmanifoldConvolution(self._dimension,
nPlanes[0], 2, 3, False)
self.ppn3_type = scn.SubmanifoldConvolution(self._dimension,
nPlanes[0], num_classes, 3,
False)
self.add_labels1 = AddLabels()
self.add_labels2 = AddLabels()
def decoder(self, a, b):
return (scn.Sequential().add(scn.BatchNormLeakyReLU(a)).add(
scn.Deconvolution(self.dimension, a, b, 2, 2,
False)).add(self.decoder_block(b, b, 1, 1)))
