File size: 3,226 Bytes
d1bfee5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 | import tensorflow as tf
from tensorflow.keras import layers
def downsample(filters, size, apply_instancenorm=True):
initializer = tf.random_normal_initializer(0., 0.02)
result = tf.keras.Sequential()
result.add(layers.Conv2D(filters, size, strides=2, padding='same',
kernel_initializer=initializer, use_bias=False))
if apply_instancenorm:
result.add(tf.keras.layers.GroupNormalization(groups=-1))
result.add(layers.LeakyReLU())
return result
def upsample(filters, size, apply_dropout=False):
initializer = tf.random_normal_initializer(0., 0.02)
result = tf.keras.Sequential()
result.add(layers.Conv2DTranspose(filters, size, strides=2, padding='same',
kernel_initializer=initializer, use_bias=False))
result.add(tf.keras.layers.GroupNormalization(groups=-1))
if apply_dropout:
result.add(layers.Dropout(0.5))
result.add(layers.ReLU())
return result
def resnet_block(filters, size=3):
initializer = tf.random_normal_initializer(0., 0.02)
result = tf.keras.Sequential()
result.add(layers.Conv2D(filters, size, padding='same', kernel_initializer=initializer, use_bias=False))
result.add(tf.keras.layers.GroupNormalization(groups=-1))
result.add(layers.ReLU())
result.add(layers.Conv2D(filters, size, padding='same', kernel_initializer=initializer, use_bias=False))
result.add(tf.keras.layers.GroupNormalization(groups=-1))
return result
def Generator(output_channels=3, num_resnet=9):
inputs = layers.Input(shape=[256, 256, 3])
# Downsampling
x = layers.Conv2D(64, 7, padding='same', kernel_initializer=tf.random_normal_initializer(0., 0.02), use_bias=False)(inputs)
x = tf.keras.layers.GroupNormalization(groups=-1)(x)
x = layers.ReLU()(x)
x = downsample(128, 3)(x)
x = downsample(256, 3)(x)
# Residual blocks
for _ in range(num_resnet):
res = resnet_block(256)(x)
x = layers.Add()([x, res])
# Upsampling
x = upsample(128, 3)(x)
x = upsample(64, 3)(x)
last = layers.Conv2D(output_channels, 7, padding='same', activation='tanh',
kernel_initializer=tf.random_normal_initializer(0., 0.02))(x)
return tf.keras.Model(inputs=inputs, outputs=last)
def Discriminator():
initializer = tf.random_normal_initializer(0., 0.02)
inputs = layers.Input(shape=[256, 256, 3])
down1 = downsample(64, 4, False)(inputs) # (bs, 128, 128, 64)
down2 = downsample(128, 4)(down1) # (bs, 64, 64, 128)
down3 = downsample(256, 4)(down2) # (bs, 32, 32, 256)
zero_pad1 = layers.ZeroPadding2D()(down3) # (bs, 34, 34, 256)
conv = layers.Conv2D(512, 4, strides=1, kernel_initializer=initializer, use_bias=False)(zero_pad1)
norm1 = tf.keras.layers.GroupNormalization(groups=-1)(conv)
leaky_relu = layers.LeakyReLU()(norm1)
zero_pad2 = layers.ZeroPadding2D()(leaky_relu)
last = layers.Conv2D(1, 4, strides=1, kernel_initializer=initializer)(zero_pad2)
return tf.keras.Model(inputs=inputs, outputs=last)
|