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