Testing Keras serialization for custom layers

test_keras.py

"""MLP Layer based on tf.keras.layer."""
import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Dense
from tensorflow.keras.layers import Layer as KerasLayer
from tensorflow.keras.models import Model as KerasModel
from tensorflow.keras.models import Sequential
from tensorflow.python.framework import tensor_shape

# flake8: noqa


class CustomLayer(KerasLayer):
    def __init__(self, output_dim, **kwargs):
        self.output_dim = output_dim
        super().__init__(**kwargs)

    def build(self, input_shape):
        input_shape = tensor_shape.TensorShape(input_shape)
        self.kernel = self.add_weight(
            name='kernel',
            shape=[input_shape[-1].value, self.output_dim],
            initializer='zero',
            trainable=True,
            dtype=tf.float32)
        super().build(input_shape)

    def call(self, x):
        return tf.keras.backend.dot(x, self.kernel)

    def compute_output_shape(self, input_shape):
        input_shape = tensor_shape.TensorShape(input_shape)
        # input_shape = input_shape.with_rank_at_least(2)
        if input_shape[-1].value is None:
          raise ValueError(
              'The innermost dimension of input_shape must be defined, but saw: %s'
              % input_shape)
        return input_shape[:-1].concatenate(self.output_dim)

    def get_config(self):
        config = {
            'output_dim': self.output_dim,
        }
        base_config = super().get_config()
        return dict(list(base_config.items()) + list(config.items()))


class CustomModel(KerasModel):
    def __init__(self, num_classes=10):
        super().__init__(name='mlp')
        self.num_classes = num_classes

        # using Keras layers
        self.dense1 = Dense(32, activation='relu')
        self.dense2 = Dense(num_classes, activation='softmax')

    def call(self, inputs):
        x = self.dense1(inputs)
        return self.dense2(x)


if __name__ == "__main__":



    dense = Dense(10)
    print("\n### Dense layer : {}".format(dense.get_config()))  # wonderful

    custom_layer = CustomLayer(10)
    print("\n### Custom layer: {}".format(
        custom_layer.get_config()))  # basically empty

    model = Sequential()
    model.add(Dense(32, activation='relu', input_dim=10))
    model.add(Dense(2, activation='softmax'))
    print("\n### Keras model : {}".format(model.get_config()))

    model = Sequential()
    model.add(Dense(5, input_dim=2))
    model.add(CustomLayer(5))

    print("\n### Keras custom model : {}\n".format(model.get_config()))

    print("\nJSON : {}\n".format(model.to_json()))
    fresh_model = tf.keras.models.model_from_json(model.to_json(), custom_objects={'CustomLayer': CustomLayer})
    print("\nRestore from JSON : {}\n".format(fresh_model.get_config()))

    print("\nYAML : {}\n".format(model.to_yaml()))
    fresh_model = tf.keras.models.model_from_yaml(model.to_yaml(), custom_objects={'CustomLayer': CustomLayer})
    print("\nRestore from YAML: {}\n".format(fresh_model.get_config()))

    print("\n### Saving/loading weights...")
    weights = model.get_weights()
    weights[2] = np.ones_like(weights[2])
    model.set_weights(weights)
    print("\nweights being saved: {}".format(model.get_weights()))
    model.save_weights('./test_keras.weights')

    print("\nweights before load: {}".format(fresh_model.get_weights()))
    fresh_model.load_weights('./test_keras.weights')
    print("\nweights after load: {}".format(fresh_model.get_weights()))

    # model = CustomModel(num_classes=2)
    # print(model.get_config())
    """
    Console log:

(garage) bogon:code ryan$ python test_keras.py

### Dense layer : {'name': 'dense', 'trainable': True, 'dtype': None, 'units': 10, 'activation': 'linear', 'use_bias': True, 'kernel_initializer': {'class_name': 'VarianceScaling', 'config': {'scale': 1.0, 'mode': 'fan_avg', 'distribution': 'uniform', 'seed': None, 'dtype': 'float32'}}, 'bias_initializer': {'class_name': 'Zeros', 'config': {'dtype': 'float32'}}, 'kernel_regularizer': None, 'bias_regularizer': None, 'activity_regularizer': None, 'kernel_constraint': None, 'bias_constraint': None}

### Custom layer: {'name': 'custom_layer', 'trainable': True, 'dtype': None, 'output_dim': 10}

### Keras model : [{'class_name': 'Dense', 'config': {'name': 'dense_1', 'trainable': True, 'batch_input_shape': (None, 10), 'dtype': 'float32', 'units': 32, 'activation': 'relu', 'use_bias': True, 'kernel_initializer': {'class_name': 'VarianceScaling', 'config': {'scale': 1.0, 'mode': 'fan_avg', 'distribution': 'uniform', 'seed': None, 'dtype': 'float32'}}, 'bias_initializer': {'class_name': 'Zeros', 'config': {'dtype': 'float32'}}, 'kernel_regularizer': None, 'bias_regularizer': None, 'activity_regularizer': None, 'kernel_constraint': None, 'bias_constraint': None}}, {'class_name': 'Dense', 'config': {'name': 'dense_2', 'trainable': True, 'dtype': 'float32', 'units': 2, 'activation': 'softmax', 'use_bias': True, 'kernel_initializer': {'class_name': 'VarianceScaling', 'config': {'scale': 1.0, 'mode': 'fan_avg', 'distribution': 'uniform', 'seed': None, 'dtype': 'float32'}}, 'bias_initializer': {'class_name': 'Zeros', 'config': {'dtype': 'float32'}}, 'kernel_regularizer': None, 'bias_regularizer': None, 'activity_regularizer': None, 'kernel_constraint': None, 'bias_constraint': None}}]

### Keras custom model : [{'class_name': 'Dense', 'config': {'name': 'dense_3', 'trainable': True, 'batch_input_shape': (None, 2), 'dtype': 'float32', 'units': 5, 'activation': 'linear', 'use_bias': True, 'kernel_initializer': {'class_name': 'VarianceScaling', 'config': {'scale': 1.0, 'mode': 'fan_avg', 'distribution': 'uniform', 'seed': None, 'dtype': 'float32'}}, 'bias_initializer': {'class_name': 'Zeros', 'config': {'dtype': 'float32'}}, 'kernel_regularizer': None, 'bias_regularizer': None, 'activity_regularizer': None, 'kernel_constraint': None, 'bias_constraint': None}}, {'class_name': 'CustomLayer', 'config': {'name': 'custom_layer_1', 'trainable': True, 'dtype': 'float32', 'output_dim': 5}}]


JSON : {"class_name": "Sequential", "config": [{"class_name": "Dense", "config": {"name": "dense_3", "trainable": true, "batch_input_shape": [null, 2], "dtype": "float32", "units": 5, "activation": "linear", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null, "dtype": "float32"}}, "bias_initializer": {"class_name": "Zeros", "config": {"dtype": "float32"}}, "kernel_regularizer": null, "bias_regularizer": null, "activity_regularizer": null, "kernel_constraint": null, "bias_constraint": null}}, {"class_name": "CustomLayer", "config": {"name": "custom_layer_1", "trainable": true, "dtype": "float32", "output_dim": 5}}], "keras_version": "2.1.6-tf", "backend": "tensorflow"}


Restore from JSON : [{'class_name': 'Dense', 'config': {'name': 'dense_3', 'trainable': True, 'batch_input_shape': (None, 2), 'dtype': 'float32', 'units': 5, 'activation': 'linear', 'use_bias': True, 'kernel_initializer': {'class_name': 'VarianceScaling', 'config': {'scale': 1.0, 'mode': 'fan_avg', 'distribution': 'uniform', 'seed': None, 'dtype': 'float32'}}, 'bias_initializer': {'class_name': 'Zeros', 'config': {'dtype': 'float32'}}, 'kernel_regularizer': None, 'bias_regularizer': None, 'activity_regularizer': None, 'kernel_constraint': None, 'bias_constraint': None}}, {'class_name': 'CustomLayer', 'config': {'name': 'custom_layer_1', 'trainable': True, 'dtype': 'float32', 'output_dim': 5}}]


YAML : backend: tensorflow
class_name: Sequential
config:
- class_name: Dense
  config:
    activation: linear
    activity_regularizer: null
    batch_input_shape: !!python/tuple [null, 2]
    bias_constraint: null
    bias_initializer:
      class_name: Zeros
      config: {dtype: float32}
    bias_regularizer: null
    dtype: float32
    kernel_constraint: null
    kernel_initializer:
      class_name: VarianceScaling
      config: {distribution: uniform, dtype: float32, mode: fan_avg, scale: 1.0, seed: null}
    kernel_regularizer: null
    name: dense_3
    trainable: true
    units: 5
    use_bias: true
- class_name: CustomLayer
  config: {dtype: float32, name: custom_layer_1, output_dim: 5, trainable: true}
keras_version: 2.1.6-tf



Restore from YAML: [{'class_name': 'Dense', 'config': {'name': 'dense_3', 'trainable': True, 'batch_input_shape': (None, 2), 'dtype': 'float32', 'units': 5, 'activation': 'linear', 'use_bias': True, 'kernel_initializer': {'class_name': 'VarianceScaling', 'config': {'scale': 1.0, 'mode': 'fan_avg', 'distribution': 'uniform', 'seed': None, 'dtype': 'float32'}}, 'bias_initializer': {'class_name': 'Zeros', 'config': {'dtype': 'float32'}}, 'kernel_regularizer': None, 'bias_regularizer': None, 'activity_regularizer': None, 'kernel_constraint': None, 'bias_constraint': None}}, {'class_name': 'CustomLayer', 'config': {'name': 'custom_layer_1', 'trainable': True, 'dtype': 'float32', 'output_dim': 5}}]


### Saving/loading weights...
2019-01-21 15:36:22.616775: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: AVX2 FMA

weights being saved: [array([[-0.2773376 , -0.31055903, -0.08883762, -0.23736203, -0.72527164],
       [-0.09298939, -0.4919746 ,  0.09247792,  0.54293466,  0.69436574]],
      dtype=float32), array([0., 0., 0., 0., 0.], dtype=float32), array([[1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1.]], dtype=float32)]

weights before load: [array([[ 0.29720736, -0.57115805,  0.80974674,  0.34152603,  0.805583  ],
       [-0.73444355, -0.74507445, -0.14655483, -0.84805375,  0.40057862]],
      dtype=float32), array([0., 0., 0., 0., 0.], dtype=float32), array([[0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]], dtype=float32)]

weights after load: [array([[-0.2773376 , -0.31055903, -0.08883762, -0.23736203, -0.72527164],
       [-0.09298939, -0.4919746 ,  0.09247792,  0.54293466,  0.69436574]],
      dtype=float32), array([0., 0., 0., 0., 0.], dtype=float32), array([[1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1.],
       [1., 1., 1., 1., 1.]], dtype=float32)]
    """