title endpoint layer pattern fchollethttps twitter comfchollet date created 20190510 last modified 20231122 description demonstration of the endpoint layer pattern layer that handles loss management accelerator gpu setup usage of endpoint layers in the functional api an endpoint layer has access to the model s targets and creates arbitrary losses in call using self addloss and metric updatestate this enables you to define losses and metrics that don t match the usual signature fnytrue ypred sampleweightnone note that you could have separate metrics for training and eval with this pattern compute the trainingtime loss value and add it to the layer using self addloss log the accuracy as a metric we could log arbitrary metrics including different metrics for training and inference return the inferencetime prediction tensor for predict exporting an inferenceonly model simply don t include targets in the model the weights stay the same usage of loss endpoint layers in subclassed models note that all inputs should be in the first argument since we want to be able to call model fitinputs setup usage of endpoint layers in the functional api an endpoint layer has access to the model s targets and creates arbitrary losses in call using self add_loss and metric update_state this enables you to define losses and metrics that don t match the usual signature fn y_true y_pred sample_weight none note that you could have separate metrics for training and eval with this pattern compute the training time loss value and add it to the layer using self add_loss log the accuracy as a metric we could log arbitrary metrics including different metrics for training and inference return the inference time prediction tensor for predict exporting an inference only model simply don t include targets in the model the weights stay the same usage of loss endpoint layers in subclassed models note that all inputs should be in the first argument since we want to be able to call model fit inputs
import os
os.environ["KERAS_BACKEND"] = "tensorflow"

import tensorflow as tf
import keras
import numpy as np




class LogisticEndpoint(keras.layers.Layer):
    def __init__(self, name=None):
        super().__init__(name=name)
        self.loss_fn = keras.losses.BinaryCrossentropy(from_logits=True)
        self.accuracy_metric = keras.metrics.BinaryAccuracy(name="accuracy")

    def call(self, logits, targets=None, sample_weight=None):
        if targets is not None:
            
            
            loss = self.loss_fn(targets, logits, sample_weight)
            self.add_loss(loss)

            
            
            self.accuracy_metric.update_state(targets, logits, sample_weight)

        
        return tf.nn.softmax(logits)


inputs = keras.Input((764,), name="inputs")
logits = keras.layers.Dense(1)(inputs)
targets = keras.Input((1,), name="targets")
sample_weight = keras.Input((1,), name="sample_weight")
preds = LogisticEndpoint()(logits, targets, sample_weight)
model = keras.Model([inputs, targets, sample_weight], preds)

data = {
    "inputs": np.random.random((1000, 764)),
    "targets": np.random.random((1000, 1)),
    "sample_weight": np.random.random((1000, 1)),
}

model.compile(keras.optimizers.Adam(1e-3))
model.fit(data, epochs=2)



inputs = keras.Input((764,), name="inputs")
logits = keras.layers.Dense(1)(inputs)
preds = LogisticEndpoint()(logits, targets=None, sample_weight=None)
inference_model = keras.Model(inputs, preds)

inference_model.set_weights(model.get_weights())

preds = inference_model.predict(np.random.random((1000, 764)))




class LogReg(keras.Model):
    def __init__(self):
        super().__init__()
        self.dense = keras.layers.Dense(1)
        self.logistic_endpoint = LogisticEndpoint()

    def call(self, inputs):
        
        
        logits = self.dense(inputs["inputs"])
        preds = self.logistic_endpoint(
            logits=logits,
            targets=inputs["targets"],
            sample_weight=inputs["sample_weight"],
        )
        return preds


model = LogReg()
data = {
    "inputs": np.random.random((1000, 764)),
    "targets": np.random.random((1000, 1)),
    "sample_weight": np.random.random((1000, 1)),
}

model.compile(keras.optimizers.Adam(1e-3))
model.fit(data, epochs=2)