Save/load a Keras model with (constant) parameters

Question

My case similar to, but a little different from, "Save/load a keras model with constants"

I'm creating an object detection model (based on YOLO ) in tf.keras (TFv1.12, yes I know) whose raw outputs need to be post-processed into bounding boxes.

This involves a few parameters which are constant for the purposes of the model, but parameters to the script that builds the model : Eg the number of classes, and the matrix of "anchor" locations to generate boxes relative to.

My model will be loaded into a TFServing container, so I'm trying to make sure:

1. The transformation is encapsulated within the model, rather than making users do it or splitting out post-processing logic
2. The saved model artifacts (eg Keras h5, or TF pb+params) are sufficient to load and serve the model

What's the right way to do this?

As far as I've been able to tell, the following don't work:

import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import backend as K
from tensorflow.keras.layers import Lambda, Layer

# Lambda layer using closure scope triggers error when trying to load the model:
# seems like `param` is defined but some weird object
def make_output_lambda(param):
    def mylambda(raw_output):
        return raw_output + param
    return Lambda(mylambda)

# Even if the custom layer type is added to `custom_objects` on
# `tf.keras.models.load_model()` - it seems to get called without the positional
# arguments:
class MyCustomLayer(Layer):
    def __init__(self, param, **kwargs):
        super(YOLOHeadLayer, self).__init__(**kwargs)
        self.param = param  # or K.constant(param) - same overall problem

    def call(self, inputs):
        return inputs + self.param

# Keras throws an error when creating a `Model` that depends on a constant
# tensor which isn't an `Input` (and who wants a constant "Input"?)
def lambdatwo(inputs):
    return inputs[0] + inputs[1]
param_tensor = K.constant(param)
y = Lambda(lambdatwo)((raw_output, param_tensor))

Answer 1

You should use the add_weight method from the inherited Layer class, using the trainable=False flag to avoid updating your constants:

class MyCustomLayer(Layer):
    def __init__(self, param, **kwargs):
        super(MyCustomLayer, self).__init__(**kwargs)
        self.param = self.add_weight(
            shape=tf.shape(param),
            initializer=lambda: param,
            trainable=False
        )

    def call(self, inputs):
        return inputs + self.param