title writing keras models with tensorflow numpy lukewoodhttps lukewood xyz date created 20210828 last modified 20210828 description overview of how to use the tensorflow numpy api to write keras models accelerator gpu introduction numpyhttps numpy org is a hugely successful python linear algebra library tensorflow recently launched tfnumpyhttps www tensorflow orgguidetfnumpy a tensorflow implementation of a large subset of the numpy api thanks to tfnumpy you can write keras layers or models in the numpy style the tensorflow numpy api has full integration with the tensorflow ecosystem features such as automatic differentiation tensorboard keras model callbacks tpu distribution and model exporting are all supported let s run through a few examples setup to test our models we will use the boston housing prices regression dataset subclassing keras model with tnp the most flexible way to make use of the keras api is to subclass the keras modelhttps keras ioapimodelsmodel class subclassing the model class gives you the ability to fully customize what occurs in the training loop this makes subclassing model a popular option for researchers in this example we will implement a model subclass that performs regression over the boston housing dataset using the tnp api note that differentiation and gradient descent is handled automatically when using the tnp api alongside keras first let s define a simple tnpforwardfeedregressionnetwork class relu activation function just like with any other keras model we can utilize any supported optimizer loss metrics or callbacks that we want let s see how the model performs great our model seems to be effectively learning to solve the problem at hand we can also write our own custom loss function using tnp implementing a keras layer based model with tnp if desired tnp can also be used in layer oriented keras code structure let s implement the same model but using a layered approach you can also seamlessly switch between tnp layers and native keras layers the model will have no issue using a normal dense layer or switching back to tnp layers the keras api offers a wide variety of layers the ability to use them alongside numpy code can be a huge time saver in projects distribution strategy tensorflow numpy and keras integrate with tensorflow distribution strategieshttps www tensorflow orgguidedistributedtraining this makes it simple to perform distributed training across multiple gpus or even an entire tpu pod we can fallback to a noop cpu strategy tensorboard integration one of the many benefits of using the keras api is the ability to monitor training through tensorboard using the tensorflow numpy api alongside keras allows you to easily leverage tensorboard to load the tensorboard from a jupyter notebook you can run the following magic loadext tensorboard to load the tensorboard from a jupyter notebook you can use the tensorboard magic tensorboard logdir logs the tensorboard monitor metrics and examine the training curve tensorboard training graphhttps i imgur comwsoufnz png the tensorboard also allows you to explore the computation graph used in your models tensorboard graph explorationhttps i imgur comtorezdl png the ability to introspect into your models can be valuable during debugging conclusion porting existing numpy code to keras models using the tensorflownumpy api is easy by integrating with keras you gain the ability to use existing keras callbacks metrics and optimizers easily distribute your training and use tensorboard migrating a more complex model such as a resnet to the tensorflow numpy api would be a great follow up learning exercise several open source numpy resnet implementations are available online introduction numpy https numpy org is a hugely successful python linear algebra library tensorflow recently launched tf_numpy https www tensorflow org guide tf_numpy a tensorflow implementation of a large subset of the numpy api thanks to tf_numpy you can write keras layers or models in the numpy style the tensorflow numpy api has full integration with the tensorflow ecosystem features such as automatic differentiation tensorboard keras model callbacks tpu distribution and model exporting are all supported let s run through a few examples setup to test our models we will use the boston housing prices regression dataset subclassing keras model with tnp the most flexible way to make use of the keras api is to subclass the keras model https keras io api models model class subclassing the model class gives you the ability to fully customize what occurs in the training loop this makes subclassing model a popular option for researchers in this example we will implement a model subclass that performs regression over the boston housing dataset using the tnp api note that differentiation and gradient descent is handled automatically when using the tnp api alongside keras first let s define a simple tnpforwardfeedregressionnetwork class relu activation function just like with any other keras model we can utilize any supported optimizer loss metrics or callbacks that we want let s see how the model performs great our model seems to be effectively learning to solve the problem at hand we can also write our own custom loss function using tnp implementing a keras layer based model with tnp if desired tnp can also be used in layer oriented keras code structure let s implement the same model but using a layered approach you can also seamlessly switch between tnp layers and native keras layers the model will have no issue using a normal dense layer or switching back to tnp layers the keras api offers a wide variety of layers the ability to use them alongside numpy code can be a huge time saver in projects distribution strategy tensorflow numpy and keras integrate with tensorflow distribution strategies https www tensorflow org guide distributed_training this makes it simple to perform distributed training across multiple gpus or even an entire tpu pod we can fallback to a no op cpu strategy tensorboard integration one of the many benefits of using the keras api is the ability to monitor training through tensorboard using the tensorflow numpy api alongside keras allows you to easily leverage tensorboard to load the tensorboard from a jupyter notebook you can run the following magic load_ext tensorboard to load the tensorboard from a jupyter notebook you can use the tensorboard magic tensorboard logdir logs the tensorboard monitor metrics and examine the training curve tensorboard training graph https i imgur com wsoufnz png the tensorboard also allows you to explore the computation graph used in your models tensorboard graph exploration https i imgur com torezdl png the ability to introspect into your models can be valuable during debugging conclusion porting existing numpy code to keras models using the tensorflow_numpy api is easy by integrating with keras you gain the ability to use existing keras callbacks metrics and optimizers easily distribute your training and use tensorboard migrating a more complex model such as a resnet to the tensorflow numpy api would be a great follow up learning exercise several open source numpy resnet implementations are available online
import os

os.environ["KERAS_BACKEND"] = "tensorflow"

import tensorflow as tf
import tensorflow.experimental.numpy as tnp
import keras
from keras import layers



(x_train, y_train), (x_test, y_test) = keras.datasets.boston_housing.load_data(
    path="boston_housing.npz", test_split=0.2, seed=113
)
input_dim = x_train.shape[1]


def evaluate_model(model: keras.Model):
    loss, percent_error = model.evaluate(x_test, y_test, verbose=0)
    print("Mean absolute percent error before training: ", percent_error)
    model.fit(x_train, y_train, epochs=200, verbose=0)
    loss, percent_error = model.evaluate(x_test, y_test, verbose=0)
    print("Mean absolute percent error after training:", percent_error)





class TNPForwardFeedRegressionNetwork(keras.Model):
    def __init__(self, blocks=None, **kwargs):
        super().__init__(**kwargs)
        if not isinstance(blocks, list):
            raise ValueError(f"blocks must be a list, got blocks={blocks}")
        self.blocks = blocks
        self.block_weights = None
        self.biases = None

    def build(self, input_shape):
        current_shape = input_shape[1]
        self.block_weights = []
        self.biases = []
        for i, block in enumerate(self.blocks):
            self.block_weights.append(
                self.add_weight(
                    shape=(current_shape, block),
                    trainable=True,
                    name=f"block-{i}",
                    initializer="glorot_normal",
                )
            )
            self.biases.append(
                self.add_weight(
                    shape=(block,),
                    trainable=True,
                    name=f"bias-{i}",
                    initializer="zeros",
                )
            )
            current_shape = block

        self.linear_layer = self.add_weight(
            shape=(current_shape, 1),
            name="linear_projector",
            trainable=True,
            initializer="glorot_normal",
        )

    def call(self, inputs):
        activations = inputs
        for w, b in zip(self.block_weights, self.biases):
            activations = tnp.matmul(activations, w) + b
            
            activations = tnp.maximum(activations, 0.0)

        return tnp.matmul(activations, self.linear_layer)




model = TNPForwardFeedRegressionNetwork(blocks=[3, 3])
model.compile(
    optimizer="adam",
    loss="mean_squared_error",
    metrics=[keras.metrics.MeanAbsolutePercentageError()],
)
evaluate_model(model)




def tnp_mse(y_true, y_pred):
    return tnp.mean(tnp.square(y_true - y_pred), axis=0)


keras.backend.clear_session()
model = TNPForwardFeedRegressionNetwork(blocks=[3, 3])
model.compile(
    optimizer="adam",
    loss=tnp_mse,
    metrics=[keras.metrics.MeanAbsolutePercentageError()],
)
evaluate_model(model)




def tnp_relu(x):
    return tnp.maximum(x, 0)


class TNPDense(keras.layers.Layer):
    def __init__(self, units, activation=None):
        super().__init__()
        self.units = units
        self.activation = activation

    def build(self, input_shape):
        self.w = self.add_weight(
            name="weights",
            shape=(input_shape[1], self.units),
            initializer="random_normal",
            trainable=True,
        )
        self.bias = self.add_weight(
            name="bias",
            shape=(self.units,),
            initializer="zeros",
            trainable=True,
        )

    def call(self, inputs):
        outputs = tnp.matmul(inputs, self.w) + self.bias
        if self.activation:
            return self.activation(outputs)
        return outputs


def create_layered_tnp_model():
    return keras.Sequential(
        [
            TNPDense(3, activation=tnp_relu),
            TNPDense(3, activation=tnp_relu),
            TNPDense(1),
        ]
    )


model = create_layered_tnp_model()
model.compile(
    optimizer="adam",
    loss="mean_squared_error",
    metrics=[keras.metrics.MeanAbsolutePercentageError()],
)
model.build((None, input_dim))
model.summary()

evaluate_model(model)




def create_mixed_model():
    return keras.Sequential(
        [
            TNPDense(3, activation=tnp_relu),
            
            layers.Dense(3, activation="relu"),
            
            TNPDense(1),
        ]
    )


model = create_mixed_model()
model.compile(
    optimizer="adam",
    loss="mean_squared_error",
    metrics=[keras.metrics.MeanAbsolutePercentageError()],
)
model.build((None, input_dim))
model.summary()

evaluate_model(model)





gpus = tf.config.list_logical_devices("GPU")
if gpus:
    strategy = tf.distribute.MirroredStrategy(gpus)
else:
    
    strategy = tf.distribute.get_strategy()
print("Running with strategy:", str(strategy.__class__.__name__))

with strategy.scope():
    model = create_layered_tnp_model()
    model.compile(
        optimizer="adam",
        loss="mean_squared_error",
        metrics=[keras.metrics.MeanAbsolutePercentageError()],
    )
    model.build((None, input_dim))
    model.summary()
    evaluate_model(model)



keras.backend.clear_session()



models = [
    (
        TNPForwardFeedRegressionNetwork(blocks=[3, 3]),
        "TNPForwardFeedRegressionNetwork",
    ),
    (create_layered_tnp_model(), "layered_tnp_model"),
    (create_mixed_model(), "mixed_model"),
]
for model, model_name in models:
    model.compile(
        optimizer="adam",
        loss="mean_squared_error",
        metrics=[keras.metrics.MeanAbsolutePercentageError()],
    )
    model.fit(
        x_train,
        y_train,
        epochs=200,
        verbose=0,
        callbacks=[keras.callbacks.TensorBoard(log_dir=f"logs/{model_name}")],
    )