def discriminator_model(): """PatchGan discriminator model. Returns: Discriminator model """ initializer = tf.random_normal_initializer(0., 0.02) inp = tf.keras.layers.Input(shape=[None, None, 3], name='input_image') tar = tf.keras.layers.Input(shape=[None, None, 3], name='target_image') x = tf.keras.layers.concatenate([inp, tar]) # (bs, 256, 256, channels*2) down1 = downsample(64, 4, False)(x) # (bs, 128, 128, 64) down2 = downsample(128, 4)(down1) # (bs, 64, 64, 128) down3 = downsample(256, 4)(down2) # (bs, 32, 32, 256) zero_pad1 = tf.keras.layers.ZeroPadding2D()(down3) # (bs, 34, 34, 256) conv = tf.keras.layers.Conv2D(512, 4, strides=1, kernel_initializer=initializer, use_bias=False)(zero_pad1) # (bs, 31, 31, 512) batchnorm1 = tf.keras.layers.BatchNormalization()(conv) leaky_relu = tf.keras.layers.LeakyReLU()(batchnorm1) zero_pad2 = tf.keras.layers.ZeroPadding2D()(leaky_relu) # (bs, 33, 33, 512) last = tf.keras.layers.Conv2D( 1, 4, strides=1, kernel_initializer=initializer)(zero_pad2) # (bs, 30, 30, 1) return tf.keras.Model(inputs=[inp, tar], outputs=last)
def upsample(filters, size, apply_dropout=False): """Upsamples an input. Conv2DTranspose => Batchnorm => Dropout => Relu Args: filters: number of filters size: filter size apply_dropout: If True, adds the dropout layer Returns: Upsample Sequential Model """ initializer = tf.random_normal_initializer(0., 0.02) result = tf.keras.Sequential() result.add( tf.keras.layers.Conv2DTranspose(filters, size, strides=2, padding='same', kernel_initializer=initializer, use_bias=False)) result.add(tf.keras.layers.BatchNormalization()) if apply_dropout: result.add(tf.keras.layers.Dropout(0.5)) result.add(tf.keras.layers.ReLU()) return result
def downsample(filters, size, apply_batchnorm=True): """Downsamples an input. Conv2D => Batchnorm => LeakyRelu Args: filters: number of filters size: filter size apply_batchnorm: If True, adds the batchnorm layer Returns: Downsample Sequential Model """ initializer = tf.random_normal_initializer(0., 0.02) result = tf.keras.Sequential() result.add( tf.keras.layers.Conv2D(filters, size, strides=2, padding='same', kernel_initializer=initializer, use_bias=False)) if apply_batchnorm: result.add(tf.keras.layers.BatchNormalization()) result.add(tf.keras.layers.LeakyReLU()) return result
def build(self, input_shape): """Initialize impulse response.""" super(Reverb, self).build(input_shape) if self.trainable: initializer = tf.random_normal_initializer(mean=0, stddev=1e-6) self._ir = self.add_weight(name='ir', shape=[self._reverb_length], dtype=tf.float32, initializer=initializer)
def build(self, input_shape): """Initialize impulse response.""" super(FilteredNoiseReverb, self).build(input_shape) if self.trainable: initializer = tf.random_normal_initializer(mean=0, stddev=1e-2) self._magnitudes = self.add_weight( name='magnitudes', shape=[self._n_frames, self._n_filter_banks], dtype=tf.float32, initializer=initializer)
def generator_model(): """Modified u-net generator model. Returns: Generator model """ down_stack = [ downsample(64, 4, apply_batchnorm=False), # (bs, 128, 128, 64) downsample(128, 4), # (bs, 64, 64, 128) downsample(256, 4), # (bs, 32, 32, 256) downsample(512, 4), # (bs, 16, 16, 512) downsample(512, 4), # (bs, 8, 8, 512) downsample(512, 4), # (bs, 4, 4, 512) downsample(512, 4), # (bs, 2, 2, 512) downsample(512, 4), # (bs, 1, 1, 512) ] up_stack = [ upsample(512, 4, apply_dropout=True), # (bs, 2, 2, 1024) upsample(512, 4, apply_dropout=True), # (bs, 4, 4, 1024) upsample(512, 4, apply_dropout=True), # (bs, 8, 8, 1024) upsample(512, 4), # (bs, 16, 16, 1024) upsample(256, 4), # (bs, 32, 32, 512) upsample(128, 4), # (bs, 64, 64, 256) upsample(64, 4), # (bs, 128, 128, 128) ] initializer = tf.random_normal_initializer(0., 0.02) last = tf.keras.layers.Conv2DTranspose(OUTPUT_CHANNELS, 4, strides=2, padding='same', kernel_initializer=initializer, activation='tanh') # (bs, 256, 256, 3) concat = tf.keras.layers.Concatenate() inputs = tf.keras.layers.Input(shape=[None, None, 3]) x = inputs # Downsampling through the model skips = [] for down in down_stack: x = down(x) skips.append(x) skips = reversed(skips[:-1]) # Upsampling and establishing the skip connections for up, skip in zip(up_stack, skips): x = up(x) x = concat([x, skip]) x = last(x) return tf.keras.Model(inputs=inputs, outputs=x)
def build(self, unused_input_shape): """Initialize impulse response.""" if self.trainable: initializer = tf.random_normal_initializer(mean=0, stddev=1e-2) self._magnitudes = self.add_weight(name='magnitudes', shape=[1, self._n_filter_banks], dtype=tf.float32, initializer=initializer) self._decay = self.add_weight( name='decay', shape=[1], dtype=tf.float32, initializer=tf.constant_initializer(4.0)) self.built = True