Support dynamically added variables

[skydoorkai]

Question:
Do we need to support the two types of keras model (Graph network from Keras Functional API, or subclass model), or just one?

There are two types of keras model:

1. Graph network from Keras Functional and Sequential APIs.
2. Subclassed networks are used when a user subclasses the Model class.

More Keras features are supported with graph networks, such as serialization, whole model saving, etc.

Example of functional API:

inputs = tf.keras.Input(shape=(10,))
x = keras.layers.Dense(1)(inputs)
outputs = tf.nn.relu(x)
network = tf.keras.Model(inputs, outputs)

inputs and outputs can be a tensor, or a list of tensors.

Example of subclass model:

class MyModel(keras.Model):
  def __init__(self):
    super(MyModel, self).__init__(name='my_model', dynamic=False)
    self.layer1 = keras.layers.Dense(10, activation='relu')
  def call(self, inputs):
    return self.layer1(inputs)

Functional API knows the input shape, while subclass model does not. In order to use model.call() for subclass model, model.build(input_shapes) need to be called in advance to let keras model knows the input shape. `input_shapes' is a shape or a list of shapes.

ElasticDL user-defined class for keras model definition

Functional API

We used to use functional API for keras model, such as in worker_test.py

class TestModel(object):
    def __init__(self):
        input1 = tf.keras.layers.Input(shape=(1,))
        x1 = tf.keras.layers.Dense(1)(input1)
        self._model = tf.keras.models.Model(input1, x1)

    def get_keras_model(self):
        return self._model

    def output(self, data):
        return self._model.call(data['x'])

    @staticmethod
    def loss(output, data):
        return tf.reduce_mean(tf.square(output - data['y']))

    @staticmethod
    def input_fn(records):
        x_list = []
        y_list = []
        # deserialize
        for r in records:
            parsed = np.frombuffer(r, dtype='float32')
            x_list.append([parsed[0]])
            y_list.append([parsed[1]])
        # batching
        batch_size = len(x_list)
        xs = np.concatenate(x_list, axis=0)
        xs = np.reshape(xs, (batch_size, 1))
        ys = np.concatenate(y_list, axis=0)
        ys = np.reshape(xs, (batch_size, 1))
        return {'x': xs, 'y': ys}

    @staticmethod
    def optimizer(lr=0.1):
        return tf.train.GradientDescentOptimizer(lr)

Subclass model

Later, we changed to subclass model worker_test.py:

class TestModel(tf.keras.Model):
    def __init__(self):
        super(TestModel, self).__init__(name='test_model')
        self.dense = tf.keras.layers.Dense(1)

    def call(self, inputs):
        return self.dense(inputs)

    @staticmethod
    def input_shapes():
        return (1, 1)

    @staticmethod
    def input_names():
        return ['x']

    @staticmethod
    def loss(outputs, labels):
        return tf.reduce_mean(tf.square(outputs - labels['y'])) 

    @staticmethod
    def input_fn(records):
        x_list = []
        y_list = []
        # deserialize
        for r in records:
            parsed = np.frombuffer(r, dtype='float32')
            x_list.append([parsed[0]])
            y_list.append([parsed[1]])
        # batching
        batch_size = len(x_list)
        xs = np.concatenate(x_list, axis=0)
        xs = np.reshape(xs, (batch_size, 1))
        ys = np.concatenate(y_list, axis=0)
        ys = np.reshape(xs, (batch_size, 1))
        return {'x': xs, 'y': ys}

    @staticmethod
    def optimizer(lr=0.1):
        return tf.train.GradientDescentOptimizer(lr)

[yuyicg]

From the keras website, it seems that only sequential model and functional API are supported. subclass model inherited from tf.keras.Model only exists in tensorflow keras which provides users with a more flexible way of describing forward path with eager execution mode, and eager mode is not supported by all the backends of original keras.

[skydoorkai]

If dynamic graph (eager execution with conditional branch) is used in subclass keras model, the number of the model variables will change in the training. Thus, we cannot get all the model variables in advance. PS have to support dynamic variable adding, as worker may report gradients for the variables that PS does not know.

Example of dynamic graph, in which 2 variables are created after the first data call, and 4 variables are created after the second data call.

import tensorflow as tf
import numpy as np

if int(tf.__version__.split(".")[0]) < 2:
    tf.enable_eager_execution()

class TestModel(tf.keras.Model):
    def __init__(self):
        super(TestModel, self).__init__()
        self.dense_layer = tf.keras.layers.Dense(1)
        self.layer1 = tf.keras.layers.Dense(1)
        self.layer2 = tf.keras.layers.Dense(1)

    def call(self, inputs):
        x = self.dense_layer(inputs)
        if inputs[0] > 0.1:
            return x
        else:
            return self.layer1(x)


model = TestModel()

inputs = np.full((1, 1), 0.2)
model.call(inputs)
print("variable number is ", len(model.trainable_variables), " after data 1")

inputs = np.full((1, 1), 0.02)
model.call(inputs)
print("variable number is ", len(model.trainable_variables), " after data 2")

Output:

variable number is  2  after data 1
variable number is  4  after data 2

[wangkuiyi]

OMG. This is a great point. Thanks for surfacing it.

[terrytangyuan]

@skydoorkai Is this fixed since now we support eager mode?

[skydoorkai]

ElasticDL assumes that all variables will be created after one minibatch model.call.

In the eager mode, ifmodel.call has conditional branches that use different variables in different branches, there may be extra variables created in the following model.call with new minibatch data.

In order to solve this issue, ElasticDL has to support dynamically added variables.
This is not fixed and we need to evaluate if this is an important issue that we need to fix.
Almost all the frequently used models do not have such an issue.

[skydoorkai]

We can close this issue for now and revisit it if needed.