Custom Layers

You can write your own custom layers, and integrate them into your Tensorleap model.

To set up your custom layer, open the Integration Scripts editor and add a class that describes your custom layer. Then, register your custom layer with the @tensorleap_custom_layer decorator to support your layer. For example, here we add a custom dense layer:

from code_loader.inner_leap_binder.leapbinder_decorators import tensorleap_custom_layer

# This class must inherit from tf.keras.layers.Layer
@tensorleap_custom_layer(name='CustomDense')
class CustomDense(tf.keras.layers.Layer):
    def __init__(self, n, **args):
        super(CustomDense, self).__init__()
        self.n = n
        self.dense = tf.keras.layers.Dense(self.n)

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

    def get_config(self):
        config = super().get_config()
        config.update({
            "n": self.n,
        })
        return config

Now, you can upload a model that includes these layers by following the Import Model guide.

Last updated

Was this helpful?