<jupyter_start><jupyter_text>Custom Tensorflow Training> Extending tf.keras for custom training functionality<jupyter_code># export
import os
from nbdev.showdoc import *
from fastcore.test import *


import tensorflow as tf
import sklearn
import numpy as np
from sklearn.datasets import fetch_california_housing
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
import matplotlib as mpl
import matplotlib.pyplot as plt
mpl.rc('axes', labelsize=14)
mpl.rc('xtick', labelsize=12)
mpl.rc('ytick', labelsize=12)


test_eq(sklearn.__version__ > "0.20", True)
test_eq(tf.__version__ > "2.0.0", True)
#hide
from nbdev.showdoc import *
TEMP_DIR = "tmp"
if not os.path.exists('tmp'):
    os.makedirs('tmp')<jupyter_output><empty_output><jupyter_text>Custom Loss Functions<jupyter_code>housing = sklearn.datasets.fetch_california_housing()
X_train_full, X_test, y_train_full, y_test = train_test_split(
    housing.data, housing.target.reshape(-1, 1), random_state=42)
X_train, X_valid, y_train, y_valid = train_test_split(
    X_train_full, y_train_full, random_state=42)

scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_valid_scaled = scaler.transform(X_valid)
X_test_scaled = scaler.transform(X_test)
def huber_fn(y_true, y_pred):
    error = y_true - y_pred
    is_small_error = tf.abs(error) < 1
    squared_loss = tf.square(error) / 2
    linear_loss  = tf.abs(error) - 0.5
    return tf.where(is_small_error, squared_loss, linear_loss)
plt.figure(figsize=(8, 3.5))
z = np.linspace(-4, 4, 200)
plt.plot(z, huber_fn(0, z), "b-", linewidth=2, label="huber($z$)")
plt.plot(z, z**2 / 2, "b:", linewidth=1, label=r"$\frac{1}{2}z^2$")
plt.plot([-1, -1], [0, huber_fn(0., -1.)], "r--")
plt.plot([1, 1], [0, huber_fn(0., 1.)], "r--")
plt.gca().axhline(y=0, color='k')
plt.gca().axvline(x=0, color='k')
plt.axis([-4, 4, 0, 4])
plt.grid(True)
plt.xlabel("$z$")
plt.legend(fontsize=14)
plt.title("Huber loss", fontsize=14)
plt.show()
input_shape = X_train.shape[1:]

model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(30, activation="selu", kernel_initializer="lecun_normal",
                       input_shape=input_shape),
    tf.keras.layers.Dense(1),
])

model.compile(loss=huber_fn, optimizer="nadam", metrics=["mae"])
model.fit(X_train_scaled, y_train, epochs=2,
          validation_data=(X_valid_scaled, y_valid))
CUSTOM_MODEL = os.path.join(TEMP_DIR, "my_model_with_a_custom_loss.h5")
CUSTOM_MODEL_THRESHOLD = os.path.join(TEMP_DIR, "my_model_with_a_custom_loss_threshold.h5")
CUSTOM_MODEL_LOSS_CLASS = os.path.join(TEMP_DIR, "my_model_with_a_custom_loss_class.h5")
CUSTOM_MODEL_CUSTOM_PARTS = os.path.join(TEMP_DIR, "my_model_with_many_custom_parts.h5")
CUSTOM_MODEL_CUSTOM_PARTS_2 = os.path.join(TEMP_DIR, "my_model_with_many_custom_parts_2.h5")
model.save(CUSTOM_MODEL)
del model
model = tf.keras.models.load_model(CUSTOM_MODEL,
                    custom_objects={"huber_fn": huber_fn})
model.fit(X_train_scaled, y_train, epochs=2,
          validation_data=(X_valid_scaled, y_valid))
def create_huber(threshold=1.0):
    def huber_fn(y_true, y_pred):
        error = y_true - y_pred
        is_small_error = tf.abs(error) < threshold
        squared_loss = tf.square(error) / 2
        linear_loss  = threshold * tf.abs(error) - threshold**2 / 2
        return tf.where(is_small_error, squared_loss, linear_loss)
    return huber_fn
model.compile(loss=create_huber(2.0), optimizer="nadam", metrics=["mae"])
model.fit(X_train_scaled, y_train, epochs=2,
          validation_data=(X_valid_scaled, y_valid))
model.save(CUSTOM_MODEL_THRESHOLD)
del model
model = tf.keras.models.load_model(CUSTOM_MODEL_THRESHOLD, 
        custom_objects={"huber_fn": create_huber(2.0)})
model.fit(X_train_scaled, y_train, epochs=2,
          validation_data=(X_valid_scaled, y_valid))
class HuberLoss(tf.keras.losses.Loss):
    def __init__(self, threshold=1.0, **kwargs):
        self.threshold = threshold
        super().__init__(**kwargs)
    def call(self, y_true, y_pred):
        error = y_true - y_pred
        is_small_error = tf.abs(error) < self.threshold
        squared_loss = tf.square(error) / 2
        linear_loss  = self.threshold * tf.abs(error) - self.threshold**2 / 2
        return tf.where(is_small_error, squared_loss, linear_loss)
    def get_config(self):
        base_config = super().get_config()
        return {**base_config, "threshold": self.threshold}
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(30, activation="selu", kernel_initializer="lecun_normal",
                       input_shape=input_shape),
    tf.keras.layers.Dense(1),
])
model.compile(loss=HuberLoss(2.), optimizer="nadam", metrics=["mae"])
model.fit(X_train_scaled, y_train, epochs=2,
          validation_data=(X_valid_scaled, y_valid))
model.save(CUSTOM_MODEL_LOSS_CLASS, save_format='h5')
# tf.saved_model.save(model, CUSTOM_MODEL_LOSS_CLASS)
# del model
# open issue: https://github.com/tensorflow/tensorflow/issues/25938
# model = tf.keras.models.load_model(CUSTOM_MODEL_LOSS_CLASS, 
#                         {"HuberLoss": HuberLoss})

# model = tf.saved_model.load(CUSTOM_MODEL_LOSS_CLASS)                       
# related to : https://github.com/tensorflow/tensorflow/issues/25938
model.fit(X_train_scaled, y_train, epochs=2,
          validation_data=(X_valid_scaled, y_valid))<jupyter_output>Epoch 1/2
363/363 [==============================] - 1s 2ms/step - loss: 0.2432 - mae: 0.5113 - val_loss: 0.3272 - val_mae: 0.5405
Epoch 2/2
363/363 [==============================] - 1s 2ms/step - loss: 0.2323 - mae: 0.5007 - val_loss: 0.2340 - val_mae: 0.4888<jupyter_text>NotesSaving tensorflow models with custom components is still very clunky.* Saving loss classes of `tf.keras.losses.Loss` throws errors (related to: [this PR](https://github.com/tensorflow/tensorflow/issues/25938)* `h5py` works most cases except for saving `tf.keras.losses.Loss` classes.* Ideally would like to save all of this as a `SavedModel`, due to integreation with tensroflow serving.  Custom Activation Functions, Initializers, Regularizers and Constants<jupyter_code>def clear_keras_session():
    tf.keras.backend.clear_session()
    np.random.seed(42)
    tf.random.set_seed(42)
clear_keras_session()
def my_softplus(z): # return value is just tf.nn.softplus(z)
    return tf.math.log(tf.exp(z) + 1.0)

def my_glorot_initializer(shape, dtype=tf.float32):
    stddev = tf.sqrt(2. / (shape[0] + shape[1]))
    return tf.random.normal(shape, stddev=stddev, dtype=dtype)

def my_l1_regularizer(weights):
    return tf.reduce_sum(tf.abs(0.01 * weights))

def my_positive_weights(weights): # return value is just tf.nn.relu(weights)
    return tf.where(weights < 0., tf.zeros_like(weights), weights)
layer = tf.keras.layers.Dense(1, activation=my_softplus,
                           kernel_initializer=my_glorot_initializer,
                           kernel_regularizer=my_l1_regularizer,
                           kernel_constraint=my_positive_weights)
clear_keras_session()
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(30, activation="selu", kernel_initializer="lecun_normal",
                       input_shape=input_shape),
    tf.keras.layers.Dense(1, activation=my_softplus,
                       kernel_regularizer=my_l1_regularizer,
                       kernel_constraint=my_positive_weights,
                       kernel_initializer=my_glorot_initializer),
])
model.compile(loss="mse", optimizer="nadam", metrics=["mae"])
model.fit(X_train_scaled, y_train, epochs=2,
          validation_data=(X_valid_scaled, y_valid))
model.save(CUSTOM_MODEL_CUSTOM_PARTS)
del model
model = tf.keras.models.load_model(
    CUSTOM_MODEL_CUSTOM_PARTS,
    custom_objects={
       "my_l1_regularizer": my_l1_regularizer,
       "my_positive_weights": my_positive_weights,
       "my_glorot_initializer": my_glorot_initializer,
       "my_softplus": my_softplus,
    })
class MyL1Regularizer(tf.keras.regularizers.Regularizer):
    def __init__(self, factor):
        self.factor = factor
    def __call__(self, weights):
        return tf.reduce_sum(tf.abs(self.factor * weights))
    def get_config(self):
        return {"factor": self.factor}
clear_keras_session()
model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(30, activation="selu", kernel_initializer="lecun_normal",
                       input_shape=input_shape),
    tf.keras.layers.Dense(1, activation=my_softplus,
                       kernel_regularizer=MyL1Regularizer(0.01),
                       kernel_constraint=my_positive_weights,
                       kernel_initializer=my_glorot_initializer),
])
model.compile(loss="mse", optimizer="nadam", metrics=["mae"])
model.fit(X_train_scaled, y_train, epochs=2,
          validation_data=(X_valid_scaled, y_valid))
model.save(CUSTOM_MODEL_CUSTOM_PARTS_2)
model = tf.keras.models.load_model(
    CUSTOM_MODEL_CUSTOM_PARTS_2,
    custom_objects={
       "MyL1Regularizer": MyL1Regularizer,
       "my_positive_weights": my_positive_weights,
       "my_glorot_initializer": my_glorot_initializer,
       "my_softplus": my_softplus,
    })<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1319830120875531, 'fraction_numerical': 0.022650549929217032, 'mean_word_length': 4.254004576659039, 'pattern_counts': {'":': 13, '<': 16, '<?xml version=': 0, '>': 15, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '31053', 'license': 'Apache-2.0', 'n_tokens_mistral': 3494, 'n_tokens_neox': 3334, 'n_words': 542, 'path': 'notebooks/custom_tf_training.ipynb', 'repo_name': 'jovsa/jovsatools', 'size': '9183'}