def build(
cls,
image_size,
alpha=1.0,
input_tensor=None,
checkpoint_file=None):
"""Build a Transfer Network Model using keras' functional API.
Args:
image_size - the size of the input and output image (H, W)
alpha - a width parameter to scale the number of channels by
Returns:
model: a keras model object
"""
x = keras.layers.Input(
shape=(image_size[0], image_size[1], 3), tensor=input_tensor)
out = cls._convolution(x, int(alpha * 32), 9, strides=1)
out = cls._convolution(out, int(alpha * 64), 3, strides=2)
out = cls._convolution(out, int(alpha * 128), 3, strides=2)
out = cls._residual_block(out, int(alpha * 128))
out = cls._residual_block(out, int(alpha * 128))
out = cls._residual_block(out, int(alpha * 128))
out = cls._residual_block(out, int(alpha * 128))
out = cls._residual_block(out, int(alpha * 128))
out = cls._upsample(out, int(alpha * 64), 3)
out = cls._upsample(out, int(alpha * 32), 3)
# Add a layer of padding to keep sizes consistent.
# out = keras.layers.ZeroPadding2D(padding=(1, 1))(out)
out = cls._convolution(out, 3, 9, relu=False, padding='valid')
# Restrict outputs of pixel values to -1 and 1.
out = keras.layers.Activation('tanh')(out)
# Deprocess the image into valid image data. Note we'll need to define
# a custom layer for this in Core ML as well.
out = layers.DeprocessStylizedImage()(out)
model = keras.models.Model(inputs=x, outputs=out)
# Optionally load weights from a checkpoint
if checkpoint_file:
logger.info(
'Loading weights from checkpoint: %s' % checkpoint_file
)
if checkpoint_file.startswith('gs://'):
checkpoint_file = utils.copy_file_from_gcs(checkpoint_file)
# Change by_name to False because the names of layers may differ
model.load_weights(checkpoint_file, by_name=False)
return model
def build(cls,
image_size,
alpha=1.0,
input_tensor=None,
checkpoint_file=None):
"""Build a Smaller Transfer Network Model using keras' functional API.
This architecture removes some blocks of layers and reduces the size
of convolutions to save on computation.
Args:
image_size - the size of the input and output image (H, W)
alpha - a width parameter to scale the number of channels by
Returns:
model: a keras model object
"""
x = keras.layers.Input(shape=(image_size[0], image_size[1], 3),
tensor=input_tensor)
out = cls._convolution(x, int(alpha * 32), 9, strides=1)
out = cls._convolution(out, int(alpha * 32), 3, strides=2)
out = cls._convolution(out, int(alpha * 32), 3, strides=2)
out = cls._residual_block(out, int(alpha * 32))
out = cls._residual_block(out, int(alpha * 32))
out = cls._residual_block(out, int(alpha * 32))
out = cls._upsample(out, int(alpha * 32), 3)
out = cls._upsample(out, int(alpha * 32), 3)
out = cls._convolution(out, 3, 9, relu=False, padding='same')
# Restrict outputs of pixel values to -1 and 1.
out = keras.layers.Activation('tanh')(out)
# Deprocess the image into valid image data. Note we'll need to define
# a custom layer for this in Core ML as well.
out = layers.DeprocessStylizedImage()(out)
model = keras.models.Model(inputs=x, outputs=out)
# Optionally load weights from a checkpoint
if checkpoint_file:
logger.info('Loading weights from checkpoint: %s' %
checkpoint_file)
if checkpoint_file.startswith('gs://'):
checkpoint_file = utils.copy_file_from_gcs(checkpoint_file)
model.load_weights(checkpoint_file, by_name=True)
return model
