def __init__(self, mixed_mag, is_training, reuse, name):
"""
input_tensor: Tensor with shape [batch_size, height, width, channels]
is_training: Boolean - should the model be trained on the current input or not
name: Model instance name
"""
with tf.variable_scope(name):
self.mixed_mag = mixed_mag
with tf.variable_scope('Convolution'):
net = mf.relu(mixed_mag)
net = mf.conv(net, filters=128, kernel_size=5, stride=(1, 1))
net = mf.batch_norm(net, is_training=is_training, reuse=reuse)
self.conv1 = net
with tf.variable_scope('Primary_Caps'):
net = mf.relu(net)
net = mf.conv(net, filters=128, kernel_size=5, stride=(1, 1))
net = mf.batch_norm(net, is_training=is_training, reuse=reuse)
self.primary_caps = net
with tf.variable_scope('Seg_Caps'):
net = mf.relu(net)
net = mf.conv(net, filters=16, kernel_size=5, stride=(1, 1))
net = mf.batch_norm(net, is_training=is_training, reuse=reuse)
self.seg_caps = net
with tf.variable_scope('Mask'):
net = mf.relu(net)
net = mf.conv(mixed_mag, filters=1, kernel_size=5, stride=(1, 1))
self.voice_mask = net
self.output = net
def __init__(self, mixed_mag, name):
"""
A basic capsule network operating on magnitude spectrograms.
"""
with tf.variable_scope(name):
self.mixed_mag = mixed_mag
with tf.variable_scope('Convolution'):
net = mf.conv(mixed_mag, filters=128, kernel_size=5, stride=(1, 1))
# Reshape layer to be 1 capsule x [filters] atoms
_, H, W, C = net.get_shape()
net = layers.Reshape((H.value, W.value, 1, C.value))(net)
self.conv1 = net
with tf.variable_scope('Primary_Caps'):
net = capsule_layers.ConvCapsuleLayer(kernel_size=5, num_capsule=8, num_atoms=8, strides=1,
padding='same',
routings=1, name='primarycaps')(net)
self.primary_caps = net
with tf.variable_scope('Seg_Caps'):
net = capsule_layers.ConvCapsuleLayer(kernel_size=1, num_capsule=1, num_atoms=8, strides=1,
padding='same',
routings=3, name='seg_caps')(net)
self.seg_caps = net
with tf.variable_scope('Reconstruction'):
net = capsule_layers.ConvCapsuleLayer(kernel_size=1, num_capsule=1, num_atoms=1, strides=1,
padding='same',
routings=3, name='reconstruction')(net)
net = tf.squeeze(net, -1)
self.output = net
def __init__(self, input_tensor, name):
"""
A basic capsule network operating on magnitude spectrograms.
"""
with tf.variable_scope(name):
self.input_tensor = input_tensor
if tf.rank(self.input_tensor) == 3:
self.out_depth = 1
else:
self.out_depth = input_tensor.shape[3].value
with tf.variable_scope('layer_1'):
net = mf.conv(input_tensor, filters=128, kernel_size=5, stride=(1, 1))
# Reshape layer to be 1 capsule x [filters] atoms
_, H, W, C = net.get_shape()
net = layers.Reshape((H.value, W.value, 1, C.value))(net)
self.conv1 = net
net = capsule_layers.ConvCapsuleLayer(kernel_size=5, num_capsule=8, num_atoms=16, strides=1,
padding='same',
routings=1, name='layer_2')(net)
self.primary_caps = net
net = capsule_layers.ConvCapsuleLayer(kernel_size=1, num_capsule=1, num_atoms=16, strides=1,
padding='same',
routings=3, name='layer_3')(net)
self.seg_caps = net
net = capsule_layers.ConvCapsuleLayer(kernel_size=1, num_capsule=self.out_depth, num_atoms=1, strides=1,
padding='same',
routings=3, name='mask')(net)
net = tf.squeeze(net, -1)
self.output = net
def __init__(self, mixed_mag, name='SegCaps_CapsNetBasic'):
"""
input_tensor: Tensor with shape [batch_size, height, width, channels]
is_training: Boolean - should the model be trained on the current input or not
name: Model instance name
"""
with tf.variable_scope(name):
self.mixed_mag = mixed_mag
with tf.variable_scope('Convolution'):
conv1 = mf.conv(mixed_mag,
filters=128,
kernel_size=5,
stride=(1, 1))
# Reshape layer to be 1 capsule x [filters] atoms
_, H, W, C = conv1.get_shape()
conv1 = layers.Reshape((H.value, W.value, 1, C.value))(conv1)
self.conv1 = conv1
with tf.variable_scope('Primary_Caps'):
primary_caps = capsule_layers.ConvCapsuleLayer(
kernel_size=5,
num_capsule=8,
num_atoms=32,
strides=1,
padding='same',
routings=1,
name='primarycaps')(conv1)
self.primary_caps = primary_caps
with tf.variable_scope('Seg_Caps'):
seg_caps = capsule_layers.ConvCapsuleLayer(
kernel_size=1,
num_capsule=1,
num_atoms=16,
strides=1,
padding='same',
routings=3,
name='seg_caps')(primary_caps)
self.seg_caps = seg_caps
with tf.variable_scope('Reconstruction'):
reconstruction = capsule_layers.ConvCapsuleLayer(
kernel_size=1,
num_capsule=1,
num_atoms=1,
strides=1,
padding='same',
routings=3,
name='seg_caps')(primary_caps)
reconstruction = tf.squeeze(reconstruction, -1)
self.output = reconstruction
def __init__(self, input_tensor, is_training, reuse, name):
"""
input_tensor: Tensor with shape [batch_size, height, width, channels]
is_training: Boolean - should the model be trained on the current input or not
name: Model instance name
"""
with tf.variable_scope(name):
self.input_tensor = input_tensor
if tf.rank(self.input_tensor) == 3:
self.out_depth = 1
else:
self.out_depth = input_tensor.shape[3].value
with tf.variable_scope('layer_1'):
net = mf.relu(input_tensor)
net = mf.conv(net, filters=128, kernel_size=5, stride=(1, 1))
net = mf.batch_norm(net, is_training=is_training, reuse=reuse)
self.l1 = net
with tf.variable_scope('layer_2'):
net = mf.relu(net)
net = mf.conv(net, filters=128, kernel_size=5, stride=(1, 1))
net = mf.batch_norm(net, is_training=is_training, reuse=reuse)
self.l2 = net
with tf.variable_scope('layer_3'):
net = mf.relu(net)
net = mf.conv(net, filters=16, kernel_size=5, stride=(1, 1))
net = mf.batch_norm(net, is_training=is_training, reuse=reuse)
self.l3 = net
with tf.variable_scope('mask'):
net = mf.relu(net)
net = mf.conv(net,
filters=self.out_depth,
kernel_size=5,
stride=(1, 1))
self.voice_mask = net
self.output = net
def __init__(self, input_tensor, is_training, reuse):
self.input_tensor = input_tensor
with tf.variable_scope('encoder'):
with tf.variable_scope('layer-1'):
net = mf.conv(self.input_tensor,
filters=16,
kernel_size=5,
stride=(2, 2))
self.l1 = net
with tf.variable_scope('layer-2'):
net = mf.lrelu(net)
net = mf.conv(net, filters=32, kernel_size=5, stride=(2, 2))
net = mf.batch_norm(net, is_training=is_training, reuse=reuse)
self.l2 = net
with tf.variable_scope('layer-3'):
net = mf.lrelu(net)
net = mf.conv(net, filters=64, kernel_size=5, stride=(2, 2))
net = mf.batch_norm(net, is_training=is_training, reuse=reuse)
self.l3 = net
with tf.variable_scope('layer-4'):
net = mf.lrelu(net)
net = mf.conv(net, filters=128, kernel_size=5, stride=(2, 2))
net = mf.batch_norm(net, is_training=is_training, reuse=reuse)
self.l4 = net
with tf.variable_scope('layer-5'):
net = mf.lrelu(net)
net = mf.conv(net, filters=256, kernel_size=5, stride=(2, 2))
net = mf.batch_norm(net, is_training=is_training, reuse=reuse)
self.l5 = net
with tf.variable_scope('layer-6'):
net = mf.lrelu(net)
net = mf.conv(net, filters=512, kernel_size=5, stride=(2, 2))
self.output = net