feat: save and load models with complex layers

README.md

@@ -53,6 +53,19 @@ model.compile(optimizer=optimizers.Adam(), loss='mse')
 
 An example working implementation of an autoencoder can be found [here](https://github.com/JesperDramsch/Complex-CNN-Seismic/).
 
+Saving and loading
+------------------
+
+Models with complexnn layers save to the `.keras` format and load again without `custom_objects`. The layers are registered with Keras when `complexnn` is imported, so import it before loading:
+
+```
+import complexnn
+import tensorflow as tf
+
+model.save("model.keras")
+model = tf.keras.models.load_model("model.keras")
+```
+
 Complex Format of Tensors
 -------------------------
 

complexnn/bn.py

@@ -10,6 +10,9 @@
 import tensorflow.keras.backend as K
 from tensorflow.keras import constraints, initializers, regularizers
 from tensorflow.keras.layers import InputSpec, Layer
+from tensorflow.keras.utils import register_keras_serializable
+
+from .utils import get_initializer
 
 from .utils import _ignore_input_shape_warning
 
@@ -23,7 +26,7 @@ def sanitizedInitGet(init):
     if init in ["sqrt_init"]:
         return sqrt_init
     else:
-        return initializers.get(init)
+        return get_initializer(init)
 
 
 def sanitizedInitSer(init):
@@ -256,6 +259,7 @@ def ComplexBN(
             return input_centred
 
 
+@register_keras_serializable(package="complexnn")
 class ComplexBatchNormalization(Layer):
     """Complex version of the real domain
     Batch normalization layer (Ioffe and Szegedy, 2014).
@@ -532,6 +536,11 @@ def update_moving_average(variable, value):
             axis=self.axis,
         )
 
+    def compute_output_shape(self, input_shape):
+        # Normalisation keeps the shape. Defining it lets Keras 3 skip
+        # tracing call() to infer the output shape.
+        return input_shape
+
     def get_config(self):
         config = {
             "axis": self.axis,

complexnn/conv.py

@@ -9,12 +9,13 @@
     InputSpec,
     Layer,
 )
+from tensorflow.keras.utils import register_keras_serializable
 
 from . import conv_utils
 from .bn import ComplexBN as complex_normalization
 from .bn import sqrt_init
 from .init import ComplexIndependentFilters, ComplexInit
-from .utils import _ignore_input_shape_warning
+from .utils import _ignore_input_shape_warning, get_initializer
 
 # WeightNorm_Conv still subclasses the Keras 2 convolution base from TensorFlow's
 # private copy of Keras. The class is broken on Keras 2 and Keras 3 and needs a
@@ -181,21 +182,23 @@ def sanitizedInitGet(init):
     elif init in ["complex", "complex_independent", "glorot_complex", "he_complex"]:
         return init
     else:
-        return initializers.get(init)
+        return get_initializer(init)
 
 
 def sanitizedInitSer(init):
     """sanitizedInitSer"""
     if init in [sqrt_init]:
         return "sqrt_init"
-    elif init == "complex" or isinstance(init, ComplexInit):
-        return "complex"
-    elif init == "complex_independent" or isinstance(init, ComplexIndependentFilters):
-        return "complex_independent"
+    elif init in ["complex", "complex_independent"]:
+        # Only these names stand for the initialisers that build() creates. An
+        # instance carries its own criterion and seed and is registered, so it
+        # goes through initializers.serialize like any other initialiser.
+        return init
     else:
         return initializers.serialize(init)
 
 
+@register_keras_serializable(package="complexnn")
 class ComplexConv(Layer):
     """Abstract nD complex convolution layer.
 
@@ -635,11 +638,13 @@ def get_config(self):
             "init_criterion": self.init_criterion,
             "spectral_parametrization": self.spectral_parametrization,
             "transposed": self.transposed,
+            "epsilon": self.epsilon,
         }
         base_config = super(ComplexConv, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
 
 
+@register_keras_serializable(package="complexnn")
 class ComplexConv1D(ComplexConv):
     """1D complex convolution layer.
     This layer creates a complex convolution kernel that is convolved
@@ -768,9 +773,12 @@ def __init__(
     def get_config(self):
         config = super(ComplexConv1D, self).get_config()
         config.pop("rank")
+        # __init__ always passes data_format="channels_last" to ComplexConv.
+        config.pop("data_format")
         return config
 
 
+@register_keras_serializable(package="complexnn")
 class ComplexConv2D(ComplexConv):
     """2D Complex convolution layer (e.g. spatial convolution over images).
     This layer creates a complex convolution kernel that is convolved
@@ -918,6 +926,7 @@ def get_config(self):
         return config
 
 
+@register_keras_serializable(package="complexnn")
 class ComplexConv3D(ComplexConv):
     """3D convolution layer (e.g. spatial convolution over volumes).
     This layer creates a complex convolution kernel that is convolved
@@ -1069,6 +1078,7 @@ def get_config(self):
         return config
 
 
+@register_keras_serializable(package="complexnn")
 class WeightNorm_Conv(Conv):
     """WeightNorm_Conv"""
 

complexnn/dense.py

@@ -6,10 +6,12 @@
 from tensorflow.keras import activations, constraints, initializers, regularizers
 from tensorflow.keras import backend as K
 from tensorflow.keras.layers import InputSpec, Layer
+from tensorflow.keras.utils import register_keras_serializable
 
-from .utils import _compute_fans, _ignore_input_shape_warning
+from .utils import _compute_fans, _ignore_input_shape_warning, get_initializer
 
 
+@register_keras_serializable(package="complexnn")
 class ComplexDense(Layer):
     """Regular complex densely-connected NN layer.
     `Dense` implements the operation:
@@ -79,11 +81,13 @@ def __init__(
         self.activation = activations.get(activation)
         self.use_bias = use_bias
         self.init_criterion = init_criterion
-        if kernel_initializer in {"complex"}:
+        # get_config stores any other initialiser as a dict, which a set
+        # membership test would try to hash, so compare with == instead.
+        if kernel_initializer == "complex":
             self.kernel_initializer = kernel_initializer
         else:
-            self.kernel_initializer = initializers.get(kernel_initializer)
-        self.bias_initializer = initializers.get(bias_initializer)
+            self.kernel_initializer = get_initializer(kernel_initializer)
+        self.bias_initializer = get_initializer(bias_initializer)
         self.kernel_regularizer = regularizers.get(kernel_regularizer)
         self.bias_regularizer = regularizers.get(bias_regularizer)
         self.activity_regularizer = regularizers.get(activity_regularizer)
@@ -128,7 +132,7 @@ def init_w_real(shape, dtype=None):
         def init_w_imag(shape, dtype=None):
             return rng.normal(size=kernel_shape, loc=0, scale=s)  # .astype(dtype)
 
-        if self.kernel_initializer in {"complex"}:
+        if self.kernel_initializer == "complex":
             real_init = init_w_real
             imag_init = init_w_imag
         else:
@@ -197,7 +201,7 @@ def compute_output_shape(self, input_shape):
         return tuple(output_shape)
 
     def get_config(self):
-        if self.kernel_initializer in {"complex"}:
+        if self.kernel_initializer == "complex":
             ki = self.kernel_initializer
         else:
             ki = initializers.serialize(self.kernel_initializer)

complexnn/init.py

@@ -5,10 +5,12 @@
 import tensorflow.keras.backend as K
 from numpy.random import RandomState
 from tensorflow.keras.initializers import Initializer
+from tensorflow.keras.utils import register_keras_serializable
 
 from .utils import _compute_fans
 
 
+@register_keras_serializable(package="complexnn")
 class IndependentFilters(Initializer):
     # This initialization constructs real-valued kernels
     # that are independent as much as possible from each other
@@ -103,6 +105,7 @@ def get_config(self):
         }
 
 
+@register_keras_serializable(package="complexnn")
 class ComplexIndependentFilters(Initializer):
     # This initialization constructs complex-valued kernels
     # that are independent as much as possible from each other
@@ -213,6 +216,7 @@ def get_config(self):
         }
 
 
+@register_keras_serializable(package="complexnn")
 class ComplexInit(Initializer):
     # The standard complex initialization using
     # either the He or the Glorot criterion.
@@ -277,7 +281,18 @@ def __call__(self, shape, dtype=None):
 
         return weight
 
+    def get_config(self):
+        return {
+            "nb_filters": self.nb_filters,
+            "kernel_size": self.kernel_size,
+            "input_dim": self.input_dim,
+            "weight_dim": self.weight_dim,
+            "criterion": self.criterion,
+            "seed": self.seed,
+        }
+
 
+@register_keras_serializable(package="complexnn")
 class SqrtInit(Initializer):
     def __call__(self, shape, dtype=None):
         return K.constant(1 / K.sqrt(2), shape=shape, dtype=dtype)

complexnn/norm.py

@@ -7,9 +7,10 @@
 import tensorflow.keras.backend as K
 from tensorflow.keras import constraints, initializers, regularizers
 from tensorflow.keras.layers import InputSpec, Layer
+from tensorflow.keras.utils import register_keras_serializable
 
 from .bn import ComplexBN as complex_normalization
-from .bn import sqrt_init
+from .bn import sanitizedInitGet, sanitizedInitSer, sqrt_init
 from .utils import _ignore_input_shape_warning
 
 
@@ -45,6 +46,7 @@ def layernorm(x, axis, epsilon, gamma, beta):
     return x_LN
 
 
+@register_keras_serializable(package="complexnn")
 class LayerNormalization(Layer):
     def __init__(
         self,
@@ -96,17 +98,16 @@ def get_config(self):
         config = {
             "epsilon": self.epsilon,
             "axis": self.axis,
-            "gamma_regularizer": self.gamma_regularizer.get_config()
-            if self.gamma_regularizer
-            else None,
-            "beta_regularizer": self.beta_regularizer.get_config()
-            if self.beta_regularizer
-            else None,
+            "beta_init": initializers.serialize(self.beta_init),
+            "gamma_init": initializers.serialize(self.gamma_init),
+            "gamma_regularizer": regularizers.serialize(self.gamma_regularizer),
+            "beta_regularizer": regularizers.serialize(self.beta_regularizer),
         }
         base_config = super(LayerNormalization, self).get_config()
         return dict(list(base_config.items()) + list(config.items()))
 
 
+@register_keras_serializable(package="complexnn")
 class ComplexLayerNorm(Layer):
     def __init__(
         self,
@@ -130,9 +131,11 @@ def __init__(
         self.axis = axis
         self.center = center
         self.scale = scale
-        self.beta_initializer = initializers.get(beta_initializer)
-        self.gamma_diag_initializer = initializers.get(gamma_diag_initializer)
-        self.gamma_off_initializer = initializers.get(gamma_off_initializer)
+        # sanitizedInitGet also accepts "sqrt_init", the name under which
+        # get_config stores the default gamma_diag_initializer.
+        self.beta_initializer = sanitizedInitGet(beta_initializer)
+        self.gamma_diag_initializer = sanitizedInitGet(gamma_diag_initializer)
+        self.gamma_off_initializer = sanitizedInitGet(gamma_off_initializer)
         self.beta_regularizer = regularizers.get(beta_regularizer)
         self.gamma_diag_regularizer = regularizers.get(gamma_diag_regularizer)
         self.gamma_off_regularizer = regularizers.get(gamma_off_regularizer)
@@ -268,17 +271,20 @@ def call(self, inputs):
             axis=self.axis,
         )
 
+    def compute_output_shape(self, input_shape):
+        # Normalisation keeps the shape. Defining it lets Keras 3 skip
+        # tracing call() to infer the output shape.
+        return input_shape
+
     def get_config(self):
         config = {
             "axis": self.axis,
             "epsilon": self.epsilon,
             "center": self.center,
             "scale": self.scale,
-            "beta_initializer": initializers.serialize(self.beta_initializer),
-            "gamma_diag_initializer": initializers.serialize(
-                self.gamma_diag_initializer
-            ),
-            "gamma_off_initializer": initializers.serialize(self.gamma_off_initializer),
+            "beta_initializer": sanitizedInitSer(self.beta_initializer),
+            "gamma_diag_initializer": sanitizedInitSer(self.gamma_diag_initializer),
+            "gamma_off_initializer": sanitizedInitSer(self.gamma_off_initializer),
             "beta_regularizer": regularizers.serialize(self.beta_regularizer),
             "gamma_diag_regularizer": regularizers.serialize(
                 self.gamma_diag_regularizer

complexnn/pool.py

@@ -5,12 +5,14 @@
 import tensorflow.keras as KE
 import tensorflow.keras.backend as KB
 import tensorflow.keras.layers as KL
+from tensorflow.keras.utils import register_keras_serializable
 
 #
 # Spectral Pooling Layer
 #
 
 
+@register_keras_serializable(package="complexnn")
 class SpectralPooling1D(KL.Layer):
     def __init__(self, topf=(0,), **kwargs):
         super(SpectralPooling1D, self).__init__()
@@ -55,6 +57,7 @@ def call(self, x, mask=None):
         return x
 
 
+@register_keras_serializable(package="complexnn")
 class SpectralPooling2D(KL.Layer):
     def __init__(self, **kwargs):
         super(SpectralPooling2D, self).__init__()