def __init__(self, class_num, full_scale, m=16, dimension=3):
nn.Module.__init__(self)
self.dimension = dimension
self.input = scn.InputLayer(dimension, full_scale, mode=4)
self.input_s = scn.InputLayer(dimension, full_scale, mode=4)
self.down_in = scn.SubmanifoldConvolution(dimension, 1, m, 3, False)
self.down_in_s = scn.SubmanifoldConvolution(dimension, 1, m, 3, False)
self.main_block1 = self.block(m, m, 2, 1)
self.main_block2 = self.block(m, 2 * m, 1, 2)
self.main_block3 = self.block(2 * m, 3 * m, 1, 2)
self.main_block4 = self.block(3 * m, 4 * m, 1, 2)
self.main_block5 = self.block(4 * m, 5 * m, 1, 2)
self.main_block6 = self.block(5 * m, 6 * m, 1, 2)
self.main_block7 = self.block(6 * m, 7 * m, 2, 2)
self.main_block8 = self.block(7 * m, 7 * m, 2, 1)
self.support_block1 = self.block(m, m, 2, 1)
self.support_block2 = self.block(m, 2 * m, 1, 2)
self.support_block3 = self.block(2 * m, 3 * m, 1, 2)
self.support_block4 = self.block(3 * m, 4 * m, 1, 2)
self.support_block5 = self.block(4 * m, 5 * m, 1, 2)
self.support_block6 = self.block(5 * m, 6 * m, 1, 2)
self.support_block7 = self.block(6 * m, 7 * m, 1, 2)
self.support_block8 = self.block(7 * m, 7 * m, 1, 1)
self.support_block2_tune = self.guide_tune(dimension, 2 * m, 2 * m, 1,
False)
self.support_block2_out = GlobalMeanAttentionPooling(dimension)
self.support_block3_tune = self.guide_tune(dimension, 4 * m, 4 * m, 1,
False)
self.support_block3_out = GlobalMeanAttentionPooling(dimension)
self.support_block4_tune = self.guide_tune(dimension, 7 * m, 7 * m, 1,
False)
self.support_block4_out = GlobalMeanAttentionPooling(dimension)
self.global_add2 = GlobalMaskLayer(dimension)
self.global_add3 = GlobalMaskLayer(dimension)
self.global_add4 = GlobalMaskLayer(dimension)
self.spatial_pick = DistMatchLayer_v4(dimension, 7 * m, topk=3)
# self.join_sub = scn.JoinTable()
self.tune_sub = self.guide_tune(dimension, 14 * m, 7 * m, 1, False)
self.deconv7 = self.decoder(7 * m, 6 * m)
self.join6 = scn.JoinTable()
self.deconv6 = self.decoder(12 * m, 5 * m)
self.deconv5 = self.decoder(5 * m, 4 * m)
self.join4 = scn.JoinTable()
self.deconv4 = self.decoder(8 * m, 3 * m)
self.deconv3 = self.decoder(3 * m, 2 * m)
self.join2 = scn.JoinTable()
self.deconv2 = self.decoder(4 * m, 2 * m)
# self.deconv1 = self.decoder(2 * m, m)
self.output = scn.OutputLayer(dimension)
self.linear = nn.Linear(2 * m, class_num)
def MultiscaleShapeContext(dimension,
n_features=1,
n_layers=3,
shape_context_size=3,
downsample_size=2,
downsample_stride=2,
bn=True):
m = sparseconvnet.Sequential()
if n_layers == 1:
m.add(
sparseconvnet.ShapeContext(dimension, n_features,
shape_context_size))
else:
m.add(sparseconvnet.ConcatTable().add(
sparseconvnet.ShapeContext(
dimension, n_features, shape_context_size)).add(
sparseconvnet.Sequential(
sparseconvnet.AveragePooling(dimension,
downsample_size,
downsample_stride),
MultiscaleShapeContext(dimension, n_features,
n_layers - 1,
shape_context_size,
downsample_size,
downsample_stride, False),
sparseconvnet.UnPooling(dimension, downsample_size,
downsample_stride)))).add(
sparseconvnet.JoinTable())
if bn:
m.add(
sparseconvnet.BatchNormalization(shape_context_size**dimension *
n_features * n_layers))
return 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 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 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):
super(CNN, self).__init__()
###############################
# Hardcoded settings
###############################
self._dimension = 3
reps = 2
kernel_size = 2
num_strides = 7
init_num_features = 8
nInputFeatures = 1
spatial_size = 128 #padding the rest for 169 PMTs
num_classes = 2 # good versus ghost
nPlanes = [(2**i) * init_num_features for i in range(0, num_strides)
] # every layer double the number of features
downsample = [kernel_size, 2]
leakiness = 0
#################################
# Input layer
#################################
self.input = scn.Sequential().add(
scn.InputLayer(self._dimension, spatial_size, mode=3)).add(
scn.SubmanifoldConvolution(self._dimension, nInputFeatures,
init_num_features, 3,
False)) # Kernel size 3, no bias
self.concat = scn.JoinTable()
#################################
# Encode layers
#################################\
self.encoding_conv = scn.Sequential()
for i in range(num_strides):
if i < 4: #hardcoded
self.encoding_conv.add(
scn.BatchNormLeakyReLU(
nPlanes[i], leakiness=leakiness)).add(
scn.Convolution(self._dimension, nPlanes[i],
nPlanes[i + 1], downsample[0],
downsample[1], False))
elif i < num_strides - 1:
self.encoding_conv.add(scn.MaxPooling(self._dimension, 2, 2))
self.output = scn.Sequential().add(
scn.SparseToDense(self._dimension, nPlanes[-1]))
###################################
# Final linear layer
###################################
self.deepest_layer_num_features = int(
nPlanes[-1] * np.power(spatial_size / (2**(num_strides - 1)), 3.))
self.classifier = torch.nn.Sequential(
torch.nn.ReLU(),
torch.nn.Linear(self.deepest_layer_num_features, 2),
)
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 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.UnPooling(dimension, downsample[0],
downsample[1]))))
m.add(scn.JoinTable())
return m
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 get_channels_combiner(
num_dims, sparse, input_channels_list, concat=True,
input_strideprod_list=None):
if input_strideprod_list is not None:
assert (
input_strideprod_list[0] == np.array(input_strideprod_list)).all()
stride = np.full(num_dims, 1)
if concat:
output_channels = sum(input_channels_list)
if sparse:
layer = scn.JoinTable()
else:
layer = DenseJoinTable()
else:
output_channels = input_channels_list[0]
assert (np.array(input_channels_list) == output_channels).all()
if sparse:
layer = scn.AddTable()
else:
layer = DenseAddTable()
return sparse, stride, output_channels, layer
def __init__(self, inputshape, reps, nin_features, nout_features, nplanes):
nn.Module.__init__(self)
"""
inputs
------
inputshape [list of int]: dimensions of the matrix or image
reps [int]: number of residual modules per layer (for both encoder and decoder)
nin_features [int]: number of features in the first convolutional layer
nout_features [int]: number of features that feed into the regression layer
nplanes [int]: the depth of the U-Net
"""
# set parameters
self.dimensions = 2 # not playing with 3D for now
# input shape: LongTensor, tuple, or list. Handled by InputLayer
# size of each spatial dimesion
self.inputshape = inputshape
if len(self.inputshape) != self.dimensions:
raise ValueError(
"expected inputshape to contain size of 2 dimensions only." +
"given %d values" % (len(self.inputshape)))
# mode variable: how to deal with repeated data
self.mode = 0
# nfeatures
self.nfeatures = nin_features
self.nout_features = nout_features
# plane structure
self.nPlanes = [self.nfeatures * 2**(n + 1) for n in xrange(nplanes)]
print self.nPlanes
# repetitions (per plane)
self.reps = reps
# residual blocks
self.residual_blocks = True
# need encoder for both source and target
# then cat tensor
# and produce one decoder for flow, another decoder for visibility
# model:
# input
self.src_inputlayer = scn.InputLayer(self.dimensions,
self.inputshape,
mode=self.mode)
self.tar1_inputlayer = scn.InputLayer(self.dimensions,
self.inputshape,
mode=self.mode)
self.tar2_inputlayer = scn.InputLayer(self.dimensions,
self.inputshape,
mode=self.mode)
# stem
self.src_stem = scn.SubmanifoldConvolution(self.dimensions, 1,
self.nfeatures, 3, False)
self.tar1_stem = scn.SubmanifoldConvolution(self.dimensions, 1,
self.nfeatures, 3, False)
self.tar2_stem = scn.SubmanifoldConvolution(self.dimensions, 1,
self.nfeatures, 3, False)
# encoders
self.source_encoder = SparseEncoder("src", self.reps, self.nfeatures,
self.nPlanes)
self.target1_encoder = SparseEncoder("tar1", self.reps, self.nfeatures,
self.nPlanes)
self.target2_encoder = SparseEncoder("tar2", self.reps, self.nfeatures,
self.nPlanes)
# concat
self.join_enclayers = []
for ilayer in xrange(len(self.nPlanes)):
self.join_enclayers.append(scn.JoinTable())
setattr(self, "join_enclayers%d" % (ilayer),
self.join_enclayers[ilayer])
# calculate decoder planes
self.decode_layers_inchs = []
self.decode_layers_outchs = []
for ilayer, enc_outchs in enumerate(reversed(self.nPlanes)):
self.decode_layers_inchs.append(4 *
enc_outchs if ilayer > 0 else 3 *
enc_outchs)
self.decode_layers_outchs.append(self.nPlanes[-(1 + ilayer)] / 2)
print "decode in chs: ", self.decode_layers_inchs
print "decode out chs: ", self.decode_layers_outchs
# decoders
self.flow1_decoder = SparseDecoder("flow1", self.reps,
self.decode_layers_inchs,
self.decode_layers_outchs)
self.flow2_decoder = SparseDecoder("flow2", self.reps,
self.decode_layers_inchs,
self.decode_layers_outchs)
# last deconv concat
self.flow1_concat = scn.JoinTable()
self.flow2_concat = scn.JoinTable()
# final feature set convolution
flow_resblock_inchs = 3 * self.nfeatures + self.decode_layers_outchs[-1]
self.flow1_resblock = create_resnet_layer(self.reps,
flow_resblock_inchs,
self.nout_features)
self.flow2_resblock = create_resnet_layer(self.reps,
flow_resblock_inchs,
self.nout_features)
# regression layer
self.flow1_out = scn.SubmanifoldConvolution(self.dimensions,
self.nout_features, 1, 1,
True)
self.flow2_out = scn.SubmanifoldConvolution(self.dimensions,
self.nout_features, 1, 1,
True)
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, 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, in_channel, n_classes, batchnorm=True, droput=0.0):
self.in_channel = in_channel
self.n_classes = n_classes
super(UNet3D, self).__init__()
self.conv1_1 = self.encoder(in_channel,
32,
bias=False,
batchnorm=batchnorm)
self.conv1_2 = self.encoder(32, 64, bias=False, batchnorm=batchnorm)
self.pool1 = scn.MaxPooling(3, 2, 2)
self.conv2_1 = self.encoder(64, 64, bias=False, batchnorm=batchnorm)
self.conv2_2 = self.encoder(64, 128, bias=False, batchnorm=batchnorm)
self.pool2 = scn.MaxPooling(3, 2, 2)
self.conv3_1 = self.encoder(128,
128,
bias=False,
filter_stride=1,
filter_size=3,
batchnorm=batchnorm)
self.conv3_2 = self.encoder(128,
256,
bias=False,
filter_stride=1,
filter_size=3,
batchnorm=batchnorm)
self.pool3 = scn.MaxPooling(3, 2, 2)
self.conv4_1 = self.encoder(256, 256, bias=False, batchnorm=batchnorm)
self.conv4_2 = self.encoder(256, 512, bias=False, batchnorm=batchnorm)
self.up5_1 = self.decoder(512,
512,
filter_size=2,
filter_stride=2,
bias=False)
self.up5_2 = scn.JoinTable()
self.conv5_1 = self.encoder(256 + 512,
256,
bias=False,
batchnorm=batchnorm)
self.conv5_2 = self.encoder(256, 256, bias=False, batchnorm=batchnorm)
self.up6_1 = self.decoder(256,
256,
filter_size=2,
filter_stride=2,
bias=False)
self.up6_2 = scn.JoinTable()
self.conv6_1 = self.encoder(128 + 256, 128, bias=False)
self.conv6_2 = self.encoder(128, 128, bias=False)
self.up7_1 = self.decoder(128,
128,
filter_size=2,
filter_stride=2,
bias=False)
self.up7_2 = scn.JoinTable()
self.conv7_1 = self.encoder(64 + 128,
64,
bias=False,
batchnorm=batchnorm)
self.conv7_2 = self.encoder(64, 64, bias=False, batchnorm=batchnorm)
self.conv8 = self.encoder(64,
n_classes,
filter_size=1,
bias=False,
batchnorm=False)
self.act = scn.Sigmoid()
self.log_level = 0
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,
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='unet_full'):
super().__init__(cfg, name)
self.model_config = cfg[name]
self.num_filters = self.model_config.get('filters', 16)
self.ghost = self.model_config.get('ghost', False)
self.seed_dim = self.model_config.get('seed_dim', 1)
self.sigma_dim = self.model_config.get('sigma_dim', 1)
self.embedding_dim = self.model_config.get('embedding_dim', 3)
self.num_classes = self.model_config.get('num_classes', 5)
self.num_gnn_features = self.model_config.get('num_gnn_features', 16)
self.inputKernel = self.model_config.get('input_kernel_size', 3)
self.coordConv = self.model_config.get('coordConv', False)
# Network Freezing Options
self.encoder_freeze = self.model_config.get('encoder_freeze', False)
self.ppn_freeze = self.model_config.get('ppn_freeze', False)
self.segmentation_freeze = self.model_config.get('segmentation_freeze', False)
self.embedding_freeze = self.model_config.get('embedding_freeze', False)
self.seediness_freeze = self.model_config.get('seediness_freeze', False)
# Input Layer Configurations and commonly used scn operations.
self.input = scn.Sequential().add(
scn.InputLayer(self.dimension, self.spatial_size, mode=3)).add(
scn.SubmanifoldConvolution(self.dimension, self.nInputFeatures, \
self.num_filters, self.inputKernel, self.allow_bias)) # Kernel size 3, no bias
self.concat = scn.JoinTable()
self.add = scn.AddTable()
# Backbone UResNet. Do NOT change namings!
self.encoder = UResNetEncoder(cfg, name='uresnet_encoder')
# self.seg_net = UResNetDecoder(cfg, name='segmentation_decoder')
self.seed_net = UResNetDecoder(cfg, name='seediness_decoder')
self.cluster_net = UResNetDecoder(cfg, name='embedding_decoder')
# Encoder-Decoder 1x1 Connections
encoder_planes = [i for i in self.encoder.nPlanes]
# seg_planes = [i for i in self.seg_net.nPlanes]
cluster_planes = [i for i in self.cluster_net.nPlanes]
seed_planes = [i for i in self.seed_net.nPlanes]
# print("Encoder Planes: ", encoder_planes)
# print("Seg Planes: ", seg_planes)
# print("Cluster Planes: ", cluster_planes)
# print("Seediness Planes: ", seed_planes)
self.skip_mode = self.model_config.get('skip_mode', 'default')
# self.seg_skip = scn.Sequential()
self.cluster_skip = scn.Sequential()
self.seed_skip = scn.Sequential()
# print(self.seg_skip)
# print(self.cluster_skip)
# print(self.seed_skip)
# Output Layers
self.output_cluster = scn.Sequential()
self._nin_block(self.output_cluster, self.cluster_net.num_filters, 4)
self.output_cluster.add(scn.OutputLayer(self.dimension))
self.output_seediness = scn.Sequential()
self._nin_block(self.output_seediness, self.seed_net.num_filters, 1)
self.output_seediness.add(scn.OutputLayer(self.dimension))
'''
self.output_segmentation = scn.Sequential()
self._nin_block(self.output_segmentation, self.seg_net.num_filters, self.num_classes)
self.output_segmentation.add(scn.OutputLayer(self.dimension))
'''
'''
self.output_gnn_features = scn.Sequential()
sum_filters = self.seg_net.num_filters + self.seed_net.num_filters + self.cluster_net.num_filters
self._resnet_block(self.output_gnn_features, sum_filters, self.num_gnn_features)
self._nin_block(self.output_gnn_features, self.num_gnn_features, self.num_gnn_features)
self.output_gnn_features.add(scn.OutputLayer(self.dimension))
'''
if self.ghost:
self.linear_ghost = scn.Sequential()
self._nin_block(self.linear_ghost, self.num_filters, 2)
# self.linear_ghost.add(scn.OutputLayer(self.dimension))
# PPN
# self.ppn = PPN(cfg)
if self.skip_mode == 'default':
'''
for p1, p2 in zip(encoder_planes, seg_planes):
self.seg_skip.add(scn.Identity())
'''
for p1, p2 in zip(encoder_planes, cluster_planes):
self.cluster_skip.add(scn.Identity())
for p1, p2 in zip(encoder_planes, seed_planes):
self.seed_skip.add(scn.Identity())
'''
self.ppn_transform = scn.Sequential()
ppn1_num_filters = seg_planes[self.ppn.ppn1_stride-self.ppn._num_strides]
self._nin_block(self.ppn_transform, encoder_planes[-1], ppn1_num_filters)
'''
elif self.skip_mode == '1x1':
'''
for p1, p2 in zip(encoder_planes, seg_planes):
self._nin_block(self.seg_skip, p1, p2)
'''
for p1, p2 in zip(encoder_planes, cluster_planes):
self._nin_block(self.cluster_skip, p1, p2)
for p1, p2 in zip(encoder_planes, seed_planes):
self._nin_block(self.seed_skip, p1, p2)
# self.ppn_transform = scn.Identity()
else:
raise ValueError('Invalid skip connection mode!')
# Freeze Layers
if self.encoder_freeze:
for p in self.encoder.parameters():
p.requires_grad = False
print('Encoder Freezed')
'''
if self.ppn_freeze:
for p in self.ppn.parameters():
p.requires_grad = False
print('PPN Freezed')
'''
'''
if self.segmentation_freeze:
for p in self.seg_net.parameters():
p.requires_grad = False
for p in self.output_segmentation.parameters():
p.requires_grad = False
print('Segmentation Branch Freezed')
'''
if self.embedding_freeze:
for p in self.cluster_net.parameters():
p.requires_grad = False
for p in self.output_cluster.parameters():
p.requires_grad = False
print('Clustering Branch Freezed')
if self.seediness_freeze:
for p in self.seed_net.parameters():
p.requires_grad = False
for p in self.output_seediness.parameters():
p.requires_grad = False
print('Seediness Branch Freezed')
# Pytorch Activations
self.tanh = nn.Tanh()
self.sigmoid = nn.Sigmoid()
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):
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 __init__(self,
inputshape,
reps,
nin_features,
nout_features,
nplanes,
flowdirs=['y2u', 'y2v'],
predict_classvec=False,
home_gpu=None,
features_per_layer=None,
show_sizes=False,
share_encoder_weights=False):
nn.Module.__init__(self)
"""
inputs
------
inputshape [list of int]: dimensions of the matrix or image
reps [int]: number of residual modules per layer (for both encoder and decoder)
nin_features [int]: number of features in the first convolutional layer
nout_features [int]: number of features that feed into the classification/regression layer
nplanes [int]: the depth of the U-Net
flowdirs [list of str]: which flow directions to implement, if two (y2u+y2v)
then we must process all three planes and produce two flow predictions.
if one, then only two planes are processed by encoder, and one flow predicted.
share_encoder_weights [bool]: if True, share the weights for the encoder
features_per_layer [list of int]: if provided, defines the feature size of each layer depth.
if None, calculated automatically.
show_sizes [bool]: if True, print sizes while running forward
predict_pixenum [bool]: if True, we predict a class vector representing what target column matches
"""
# set parameters
self.dimensions = 2 # not playing with 3D for now
self.home_gpu = home_gpu
self.predict_classvec = predict_classvec
# input shape: LongTensor, tuple, or list. Handled by InputLayer
# size of each spatial dimesion
self.inputshape = inputshape
if len(self.inputshape) != self.dimensions:
raise ValueError(
"expected inputshape to contain size of 2 dimensions only." +
"given %d values" % (len(self.inputshape)))
# mode variable: how to deal with repeated data
self.mode = 0
# for debug, show sizes of layers/planes
self._show_sizes = show_sizes
# nfeatures
self.nfeatures = nin_features
self.nout_features = nout_features
# plane structure
if features_per_layer is None:
self.nPlanes = [
self.nfeatures * 2**(n + 1) for n in xrange(nplanes)
]
else:
if (type(features_per_layer) is list
and len(features_per_layer) == nplanes
and type(features_per_layer[0]) is int):
self.nPlanes = features_per_layer
else:
raise ValueError(
"features_per_layer should be a list of int with number of elements equalling 'nplanes' argument"
)
if self._show_sizes:
print "Features per plane/layer: ", self.nPlanes
# repetitions (per plane)
self.reps = reps
# residual blocks
self.residual_blocks = True
# need encoder for both source and target
# then cat tensor
# and produce one decoder for flow, another decoder for visibility
# set which flows to run
for flowdir in flowdirs:
if flowdir not in SparseLArFlow.avail_flows:
raise ValueError(
"flowdir={} not available. Allowed flows: {}".format(
flowdir, SparseLArFlow.avail_flows))
self.flowdirs = flowdirs
self.nflows = len(self.flowdirs)
# do we share weights
self._share_encoder_weights = share_encoder_weights
# model:
# input
self.src_inputlayer = scn.InputLayer(self.dimensions,
self.inputshape,
mode=self.mode)
if 'y2u' in self.flowdirs:
self.tar1_inputlayer = scn.InputLayer(self.dimensions,
self.inputshape,
mode=self.mode)
else:
self.tar1_inputlayer = None
if 'y2v' in self.flowdirs:
self.tar2_inputlayer = scn.InputLayer(self.dimensions,
self.inputshape,
mode=self.mode)
else:
self.tar2_inputlayer = None
# stem
self.src_stem = scn.SubmanifoldConvolution(self.dimensions, 1,
self.nfeatures, 3, False)
if not self._share_encoder_weights:
# if not sharing weights, producer separate stems
if 'y2u' in self.flowdirs:
self.tar1_stem = scn.SubmanifoldConvolution(
self.dimensions, 1, self.nfeatures, 3, False)
else:
self.tar1_stem = None
if 'y2v' in self.flowdirs:
self.tar2_stem = scn.SubmanifoldConvolution(
self.dimensions, 1, self.nfeatures, 3, False)
else:
self.tar2_stem = None
# encoders
self.source_encoder = SparseEncoder("src", self.reps, self.nfeatures,
self.nPlanes)
if not self._share_encoder_weights:
# if not sharing weights, add additional encoders
if 'y2u' in self.flowdirs:
self.target1_encoder = SparseEncoder("tar1", self.reps,
self.nfeatures,
self.nPlanes)
else:
self.target1_encoder = None
if 'y2v' in self.flowdirs:
self.target2_encoder = SparseEncoder("tar2", self.reps,
self.nfeatures,
self.nPlanes)
else:
self.target2_encoder = None
if self._show_sizes:
self.source_encoder.set_verbose(True)
# concat
self.join_enclayers = []
for ilayer in xrange(len(self.nPlanes)):
self.join_enclayers.append(scn.JoinTable())
setattr(self, "join_enclayers%d" % (ilayer),
self.join_enclayers[ilayer])
# calculate decoder planes
self.decode_layers_inchs = []
self.decode_layers_outchs = []
for ilayer, enc_outchs in enumerate(reversed(self.nPlanes)):
# for input, we expect features from encoder + output features
enc_nfeatures = (self.nflows + 1) * enc_outchs
dec_nfeatures = 0 if ilayer == 0 else self.decode_layers_outchs[-1]
self.decode_layers_inchs.append(enc_nfeatures + dec_nfeatures)
self.decode_layers_outchs.append(enc_nfeatures)
if self._show_sizes:
print "decoder layers input chs: ", self.decode_layers_inchs
print "decoder layers output chs: ", self.decode_layers_outchs
# decoders
if 'y2u' in self.flowdirs:
self.flow1_decoder = SparseDecoder("flow1", self.reps,
self.decode_layers_inchs,
self.decode_layers_outchs)
else:
self.flow1_decoder = None
if 'y2v' in self.flowdirs:
self.flow2_decoder = SparseDecoder("flow2", self.reps,
self.decode_layers_inchs,
self.decode_layers_outchs)
else:
self.flow2_decoder = None
if self._show_sizes:
for fd in [self.flow1_decoder, self.flow2_decoder]:
if fd is not None:
fd.set_verbose(True)
# last deconv concat
if 'y2u' in self.flowdirs:
self.flow1_concat = scn.JoinTable()
else:
self.flow1_concat = None
if 'y2v' in self.flowdirs:
self.flow2_concat = scn.JoinTable()
else:
self.flow2_concat = None
# final feature set convolution
flow_resblock_inchs = (
self.nflows + 1) * self.nfeatures + self.decode_layers_outchs[-1]
if 'y2u' in self.flowdirs:
self.flow1_resblock = create_resnet_layer(self.reps,
flow_resblock_inchs,
self.nout_features)
else:
self.flow1_resblock = None
if 'y2v' in self.flowdirs:
self.flow2_resblock = create_resnet_layer(self.reps,
flow_resblock_inchs,
self.nout_features)
else:
self.flow2_resblock = None
# OUTPUT LAYER
if self.predict_classvec:
if 'y2u' in self.flowdirs:
self.flow1_out = scn.SubmanifoldConvolution(
self.dimensions, self.nout_features, self.inputshape[1], 1,
True)
else:
self.flow1_out = None
if 'y2v' in self.flowdirs:
self.flow2_out = scn.SubmanifoldConvolution(
self.dimensions, self.nout_features, self.inputshape[1], 1,
True)
else:
self.flow2_out = None
else:
# regression layer
if 'y2u' in self.flowdirs:
self.flow1_out = scn.SubmanifoldConvolution(
self.dimensions, self.nout_features, 1, 1, True)
else:
self.flow1_out = None
if 'y2v' in self.flowdirs:
self.flow2_out = scn.SubmanifoldConvolution(
self.dimensions, self.nout_features, 1, 1, True)
else:
self.flow2_out = None
