diff --git a/README.md b/README.md index eb32873..40f176b 100644 --- a/README.md +++ b/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 ------------------------- diff --git a/complexnn/bn.py b/complexnn/bn.py index 23d9ac0..245080d 100644 --- a/complexnn/bn.py +++ b/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, diff --git a/complexnn/conv.py b/complexnn/conv.py index 32fd3bb..d29d70d 100644 --- a/complexnn/conv.py +++ b/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""" diff --git a/complexnn/dense.py b/complexnn/dense.py index 5d1f9de..f84bde4 100644 --- a/complexnn/dense.py +++ b/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) diff --git a/complexnn/init.py b/complexnn/init.py index 4c652d0..fb25351 100644 --- a/complexnn/init.py +++ b/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) diff --git a/complexnn/norm.py b/complexnn/norm.py index 4f84af6..950257d 100644 --- a/complexnn/norm.py +++ b/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 diff --git a/complexnn/pool.py b/complexnn/pool.py index c20c77c..0e638d6 100644 --- a/complexnn/pool.py +++ b/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__() diff --git a/complexnn/utils.py b/complexnn/utils.py index 5043aa2..bdabfd9 100644 --- a/complexnn/utils.py +++ b/complexnn/utils.py @@ -8,7 +8,9 @@ import numpy as np import tensorflow as tf import tensorflow.keras.backend as K +from tensorflow.keras import initializers from tensorflow.keras.layers import Layer +from tensorflow.keras.utils import get_registered_object, register_keras_serializable # Keras 3 warns whenever a layer receives `input_shape` or `input_dim` and # recommends `keras.Input(shape)` instead. The complex layers keep accepting @@ -49,6 +51,25 @@ def _ignore_input_shape_warning(kwargs): yield +def get_initializer(identifier): + """Resolve an initialiser like ``initializers.get``, complexnn ones included. + + The legacy .h5 loader looks a serialised initialiser up by its bare class + name, while complexnn registers its classes as ``complexnn>ClassName``, so + a ``ComplexInit`` saved with its config would not load. When Keras cannot + resolve such a config, look it up by its registered name instead. + """ + try: + return initializers.get(identifier) + except (ValueError, TypeError): + if isinstance(identifier, dict): + registered = str(identifier.get("registered_name") or "") + cls = get_registered_object(registered) if registered else None + if cls is not None and registered.startswith("complexnn>"): + return cls.from_config(identifier.get("config") or {}) + raise + + # # GetReal/GetImag Lambda layer Implementation # @@ -152,6 +173,7 @@ def _compute_fans(shape, data_format="channels_last"): return fan_in, fan_out +@register_keras_serializable(package="complexnn") class GetReal(Layer): def call(self, inputs): return get_realpart(inputs) @@ -160,6 +182,7 @@ def compute_output_shape(self, input_shape): return getpart_output_shape(input_shape) +@register_keras_serializable(package="complexnn") class GetImag(Layer): def call(self, inputs): return get_imagpart(inputs) @@ -168,6 +191,7 @@ def compute_output_shape(self, input_shape): return getpart_output_shape(input_shape) +@register_keras_serializable(package="complexnn") class GetAbs(Layer): def call(self, inputs): return get_abs(inputs) diff --git a/tests/test_bn.py b/tests/test_bn.py index 834f529..67acb9f 100644 --- a/tests/test_bn.py +++ b/tests/test_bn.py @@ -48,3 +48,13 @@ def test_fit_updates_moving_mean(self): self.x, np.zeros_like(self.x), batch_size=32, epochs=1, verbose=0 ) self.assertFalse(np.allclose(self.bn.moving_mean.numpy(), before)) + + +class TestComplexBatchNormalizationShape(unittest.TestCase): + """ComplexBatchNormalization keeps the input shape""" + + def test_compute_output_shape(self): + layer = conn.ComplexBN() + for shape in [(None, 8), (None, 4, 4, 8)]: + with self.subTest(shape=shape): + self.assertEqual(tuple(layer.compute_output_shape(shape)), shape) diff --git a/tests/test_conv.py b/tests/test_conv.py index 4d82ecb..108998d 100644 --- a/tests/test_conv.py +++ b/tests/test_conv.py @@ -56,3 +56,62 @@ def test_conv2Dtranspose(self): true = (None, 128, 128, 4) calc = model.output_shape self.assertEqual(true, calc) + + +class TestConvConfig(unittest.TestCase): + """Conv layer configs rebuild the same layer""" + + def test_conv1d_from_config(self): + """ComplexConv1D sets data_format itself, so its config must not repeat it""" + layer = conn.ComplexConv1D(filters=2, kernel_size=3, padding="same") + config = layer.get_config() + clone = conn.ComplexConv1D.from_config(config) + self.assertEqual(clone.get_config(), config) + self.assertEqual(clone.data_format, "channels_last") + + def test_epsilon_in_config(self): + """epsilon changes the normalised kernel, so from_config must keep it""" + layer = conn.ComplexConv2D( + filters=2, kernel_size=3, normalize_weight=True, epsilon=1e-3 + ) + clone = conn.ComplexConv2D.from_config(layer.get_config()) + self.assertEqual(clone.epsilon, 1e-3) + + def test_complex_initializer_instance_in_config(self): + """A complexnn initialiser instance keeps its own config, seed included""" + initializers = [ + conn.ComplexInit( + kernel_size=(3, 3), + input_dim=2, + weight_dim=2, + nb_filters=2, + criterion="glorot", + seed=5, + ), + conn.ComplexIndependentFilters( + kernel_size=(3, 3), + input_dim=2, + weight_dim=2, + nb_filters=2, + criterion="he", + seed=5, + ), + ] + for init in initializers: + with self.subTest(type(init).__name__): + layer = conn.ComplexConv2D( + filters=2, kernel_size=3, kernel_initializer=init + ) + config = layer.get_config() + clone = conn.ComplexConv2D.from_config(config) + self.assertIs(type(clone.kernel_initializer), type(init)) + self.assertEqual( + clone.kernel_initializer.get_config(), init.get_config() + ) + self.assertEqual(clone.get_config(), config) + + layer.build((None, 6, 6, 4)) + clone.build((None, 6, 6, 4)) + np.testing.assert_array_equal( + np.asarray(clone.kernel), np.asarray(layer.kernel) + ) diff --git a/tests/test_dense.py b/tests/test_dense.py index b65ccc1..98721a9 100644 --- a/tests/test_dense.py +++ b/tests/test_dense.py @@ -18,3 +18,25 @@ def test_outputs_dense(self): true = (None, 16 * 2) calc = layer.compute_output_shape(input_shape) self.assertEqual(true, calc) + + +class TestDenseConfig(unittest.TestCase): + """ComplexDense configs rebuild the same kernel initialiser""" + + def test_kernel_initializer_round_trip(self): + """A kernel initialiser stored as a dict is read back, not hashed""" + cases = [ + ("complex", str), + ("glorot_uniform", tf.keras.initializers.GlorotUniform), + ( + conn.ComplexInit(kernel_size=(4, 4), input_dim=4, weight_dim=2), + conn.ComplexInit, + ), + ] + for kernel_initializer, expected_type in cases: + with self.subTest(expected_type.__name__): + layer = conn.ComplexDense(4, kernel_initializer=kernel_initializer) + config = layer.get_config() + clone = conn.ComplexDense.from_config(config) + self.assertIsInstance(clone.kernel_initializer, expected_type) + self.assertEqual(clone.get_config(), config) diff --git a/tests/test_init.py b/tests/test_init.py new file mode 100644 index 0000000..0695ab5 --- /dev/null +++ b/tests/test_init.py @@ -0,0 +1,30 @@ +import unittest + +import complexnn as conn +import numpy as np + + +def make_initializers(): + """One of each complex initialiser, each filling a (4, 6) dense kernel""" + return [ + conn.ComplexInit( + kernel_size=(4, 6), input_dim=4, weight_dim=2, criterion="glorot", seed=3 + ), + conn.IndependentFilters(kernel_size=(4, 6), input_dim=4, weight_dim=2, seed=3), + conn.ComplexIndependentFilters( + kernel_size=(4, 3), input_dim=4, weight_dim=2, seed=3 + ), + ] + + +class TestInitializerConfig(unittest.TestCase): + """Initialiser configs rebuild an initialiser that gives the same values""" + + shape = (4, 6) + + def test_config_round_trip(self): + for init in make_initializers(): + with self.subTest(type(init).__name__): + clone = type(init).from_config(init.get_config()) + self.assertEqual(clone.get_config(), init.get_config()) + np.testing.assert_array_equal(clone(self.shape), init(self.shape)) diff --git a/tests/test_norm.py b/tests/test_norm.py new file mode 100644 index 0000000..41544fa --- /dev/null +++ b/tests/test_norm.py @@ -0,0 +1,55 @@ +import unittest + +import complexnn as conn +import tensorflow as tf +from complexnn.bn import sqrt_init + + +class TestComplexLayerNormConfig(unittest.TestCase): + """ComplexLayerNorm configs rebuild the same layer""" + + def test_from_config(self): + layer = conn.ComplexLayerNorm() + config = layer.get_config() + clone = conn.ComplexLayerNorm.from_config(config) + self.assertEqual(clone.get_config(), config) + + def test_sqrt_init_serialises_by_name(self): + """The default gamma initialiser is stored by name and restored as sqrt_init""" + config = conn.ComplexLayerNorm().get_config() + self.assertEqual(config["gamma_diag_initializer"], "sqrt_init") + clone = conn.ComplexLayerNorm.from_config(config) + self.assertIs(clone.gamma_diag_initializer, sqrt_init) + + def test_keras_initializer_round_trip(self): + """Keras initialisers still go through the Keras registry""" + layer = conn.ComplexLayerNorm(gamma_diag_initializer="ones") + clone = conn.ComplexLayerNorm.from_config(layer.get_config()) + self.assertIsInstance(clone.gamma_diag_initializer, tf.keras.initializers.Ones) + + +class TestLayerNormalizationConfig(unittest.TestCase): + """LayerNormalization configs keep their initialisers and regularisers""" + + def test_from_config(self): + layer = conn.LayerNormalization( + beta_init="ones", + gamma_init="zeros", + gamma_regularizer=tf.keras.regularizers.L2(0.01), + ) + config = layer.get_config() + clone = conn.LayerNormalization.from_config(config) + self.assertEqual(clone.get_config(), config) + self.assertIsInstance(clone.beta_init, tf.keras.initializers.Ones) + self.assertIsInstance(clone.gamma_init, tf.keras.initializers.Zeros) + self.assertAlmostEqual(clone.gamma_regularizer.get_config()["l2"], 0.01) + + +class TestComplexLayerNormShape(unittest.TestCase): + """ComplexLayerNorm keeps the input shape""" + + def test_compute_output_shape(self): + layer = conn.ComplexLayerNorm() + for shape in [(None, 8), (None, 4, 4, 8)]: + with self.subTest(shape=shape): + self.assertEqual(tuple(layer.compute_output_shape(shape)), shape) diff --git a/tests/test_saving.py b/tests/test_saving.py new file mode 100644 index 0000000..627c6f3 --- /dev/null +++ b/tests/test_saving.py @@ -0,0 +1,277 @@ +import inspect +import json +import os +import tempfile +import unittest +import warnings + +import complexnn as conn +import numpy as np +import tensorflow as tf +from tensorflow.keras.utils import get_registered_name, get_registered_object + +# Every working layer, with an input shape it accepts. Each entry builds a new +# layer, so tests never share an instance. Layers that take arguments set some +# of them away from their defaults: a value that from_config does not restore +# then changes the config, and a lost argument that affects the output +# (activation, strides, dilation_rate, epsilon) changes the output. Arguments +# missing from the config altogether are caught by TestConfigCompleteness. +LAYERS = { + "ComplexDense": ( + (8,), + lambda: conn.ComplexDense(4, activation="tanh", init_criterion="glorot"), + ), + # A seeded Keras initialiser: Keras 2 warns when an unseeded one is called + # twice, and ComplexDense calls it for the real and the imaginary kernel. + "ComplexDense GlorotUniform": ( + (8,), + lambda: conn.ComplexDense( + 4, kernel_initializer=tf.keras.initializers.GlorotUniform(seed=1) + ), + ), + "ComplexDense ComplexInit": ( + (8,), + lambda: conn.ComplexDense( + 4, + kernel_initializer=conn.ComplexInit( + kernel_size=(4, 4), input_dim=4, weight_dim=2 + ), + ), + ), + "ComplexConv1D": ( + (10, 4), + lambda: conn.ComplexConv1D( + 2, 3, padding="same", dilation_rate=2, activation="tanh" + ), + ), + "ComplexConv2D": ( + (6, 6, 4), + lambda: conn.ComplexConv2D( + 2, 3, strides=2, padding="same", init_criterion="glorot" + ), + ), + "ComplexConv2D normalize_weight": ( + (6, 6, 4), + lambda: conn.ComplexConv2D( + 2, 3, padding="same", normalize_weight=True, epsilon=1e-2 + ), + ), + "ComplexConv2D ComplexInit": ( + (6, 6, 4), + lambda: conn.ComplexConv2D( + 2, + 3, + padding="same", + kernel_initializer=conn.ComplexInit( + kernel_size=(3, 3), input_dim=2, weight_dim=2, nb_filters=2, seed=5 + ), + ), + ), + "ComplexConv2D transposed": ( + (6, 6, 4), + lambda: conn.ComplexConv2D(1, 3, strides=2, padding="same", transposed=True), + ), + "ComplexConv3D": ( + (4, 4, 4, 4), + lambda: conn.ComplexConv3D( + 2, 3, padding="same", dilation_rate=2, use_bias=False + ), + ), + "ComplexBN": ((4, 4, 8), lambda: conn.ComplexBN(epsilon=1e-2, momentum=0.5)), + # On a 2D input both layer norms reduce over no axes and return beta, so + # they get inputs with a feature axis on which the output depends on gamma. + "ComplexLayerNorm": ((4, 4, 8), lambda: conn.ComplexLayerNorm(epsilon=1e-2)), + "LayerNormalization": ((4, 8), lambda: conn.LayerNormalization(epsilon=1e-2)), + "GetReal": ((8,), lambda: conn.GetReal()), + "GetImag": ((4, 4, 8), lambda: conn.GetImag()), + "GetAbs": ((8,), lambda: conn.GetAbs()), +} + + +def build_model(shape, make_layer, rng): + """A Sequential model around one layer, with every weight moved off its initial value""" + model = tf.keras.Sequential([tf.keras.Input(shape), make_layer()]) + # With moved weights, matching outputs show that the weights were restored + # and not just initialised again. The noise is small enough to keep the + # normalisation variances positive. + model.set_weights( + [w + rng.uniform(-0.1, 0.1, w.shape) for w in model.get_weights()] + ) + return model + + +class TestLayerConfig(unittest.TestCase): + """Each working layer rebuilds from its config and computes the same output""" + + def test_from_config(self): + rng = np.random.default_rng(0) + for name, (shape, make_layer) in LAYERS.items(): + with self.subTest(name): + layer = build_model(shape, make_layer, rng).layers[-1] + config = layer.get_config() + clone = type(layer).from_config(config) + self.assertEqual(clone.get_config(), config) + + x = rng.standard_normal((3, *shape)).astype("float32") + clone(x) # builds the clone, so it can take the weights + clone.set_weights(layer.get_weights()) + np.testing.assert_allclose( + np.asarray(clone(x)), np.asarray(layer(x)), rtol=1e-6, atol=1e-6 + ) + + +def init_parameters(cls): + """Named __init__ arguments of cls and of the complexnn classes it passes them to""" + names = set() + for klass in cls.__mro__: + if "__init__" in vars(klass) and klass.__module__.startswith("complexnn."): + for param in inspect.signature(klass.__init__).parameters.values(): + if param.name != "self" and param.kind not in ( + param.VAR_POSITIONAL, + param.VAR_KEYWORD, + ): + names.add(param.name) + return names + + +# Arguments that get_config leaves out. ComplexConv1D/2D/3D pass rank to +# ComplexConv themselves, and ComplexConv1D also data_format, so the config +# must not repeat them. They accept seed but never use it, so there is +# nothing to store yet. +NOT_IN_CONFIG = { + conn.ComplexConv1D: {"rank", "data_format", "seed"}, + conn.ComplexConv2D: {"rank", "seed"}, + conn.ComplexConv3D: {"rank", "seed"}, +} + + +class TestConfigCompleteness(unittest.TestCase): + """get_config stores every argument that __init__ takes""" + + def test_every_argument_in_config(self): + checked = set() + for _, make_layer in LAYERS.values(): + layer = make_layer() + cls = type(layer) + if cls in checked: + continue + checked.add(cls) + with self.subTest(cls.__name__): + missing = init_parameters(cls) - set(layer.get_config()) + # assertEqual rather than a subset check, so the allow-list has + # to shrink once a gap is closed. + self.assertEqual(missing, NOT_IN_CONFIG.get(cls, set())) + + +class TestRegistration(unittest.TestCase): + """Every class that complexnn exports is registered for serialisation""" + + def test_public_classes_are_registered(self): + classes = [ + obj + for obj in vars(conn).values() + if inspect.isclass(obj) and obj.__module__.startswith("complexnn.") + ] + self.assertGreater(len(classes), 0) + for cls in classes: + with self.subTest(cls.__name__): + registered_name = "complexnn>" + cls.__name__ + self.assertEqual(get_registered_name(cls), registered_name) + self.assertIs(get_registered_object(registered_name), cls) + + +class TestKerasFileRoundTrip(unittest.TestCase): + """Models with complex layers save to .keras and load without custom_objects""" + + def setUp(self): + tmp = tempfile.TemporaryDirectory() + self.addCleanup(tmp.cleanup) + self.tmp = tmp.name + + def save_and_load(self, model, name): + path = os.path.join(self.tmp, name + ".keras") + with warnings.catch_warnings(): + # TensorFlow variables implement __array__ without the copy keyword + # that numpy 2 passes, so Keras 3 warns while saving any model, + # built-in layers included. The warning is not about complexnn. + warnings.filterwarnings( + "ignore", + message="__array__ implementation doesn't accept a copy keyword", + category=DeprecationWarning, + ) + model.save(path) + return tf.keras.models.load_model(path) + + def test_layers(self): + rng = np.random.default_rng(1) + for i, (name, (shape, make_layer)) in enumerate(LAYERS.items()): + with self.subTest(name): + model = build_model(shape, make_layer, rng) + loaded = self.save_and_load(model, f"layer_{i}") + self.assertIs(type(loaded.layers[-1]), type(model.layers[-1])) + restored = loaded.get_weights() + self.assertEqual(len(restored), len(model.get_weights())) + for w_loaded, w_saved in zip( + restored, model.get_weights(), strict=True + ): + np.testing.assert_array_equal(w_loaded, w_saved) + + x = rng.standard_normal((3, *shape)).astype("float32") + np.testing.assert_allclose( + loaded.predict(x, verbose=0), + model.predict(x, verbose=0), + rtol=1e-6, + atol=1e-6, + ) + + def test_h5_with_complex_initializer(self): + """Legacy .h5 files keep a complexnn initialiser instance + + Keras 2's .h5 loader looks an initialiser up by its bare class name, + which is not where complexnn registers it. + """ + rng = np.random.default_rng(2) + for name in ("ComplexDense ComplexInit", "ComplexConv2D ComplexInit"): + shape, make_layer = LAYERS[name] + with self.subTest(name): + model = build_model(shape, make_layer, rng) + path = os.path.join(self.tmp, name.replace(" ", "_") + ".h5") + with warnings.catch_warnings(): + # Keras 3 warns that .h5 is a legacy format. + warnings.simplefilter("ignore") + model.save(path) + loaded = tf.keras.models.load_model(path) + restored = loaded.layers[-1].kernel_initializer + self.assertIs(type(restored), conn.ComplexInit) + # The file stores JSON, which turns tuples into lists. + saved_config = model.layers[-1].kernel_initializer.get_config() + self.assertEqual( + restored.get_config(), json.loads(json.dumps(saved_config)) + ) + x = rng.standard_normal((3, *shape)).astype("float32") + np.testing.assert_allclose( + loaded.predict(x, verbose=0), + model.predict(x, verbose=0), + rtol=1e-6, + atol=1e-6, + ) + + def test_initializers(self): + """A built-in layer keeps a complex initialiser through saving and loading""" + initializers = [ + conn.ComplexInit(kernel_size=(4, 6), input_dim=4, weight_dim=2, seed=3), + conn.IndependentFilters( + kernel_size=(4, 6), input_dim=4, weight_dim=2, seed=3 + ), + conn.ComplexIndependentFilters( + kernel_size=(4, 3), input_dim=4, weight_dim=2, seed=3 + ), + ] + for i, init in enumerate(initializers): + with self.subTest(type(init).__name__): + dense = tf.keras.layers.Dense(6, kernel_initializer=init) + model = tf.keras.Sequential([tf.keras.Input((4,)), dense]) + loaded = self.save_and_load(model, f"initializer_{i}") + restored = loaded.layers[-1].kernel_initializer + self.assertIs(type(restored), type(init)) + np.testing.assert_array_equal(restored((4, 6)), init((4, 6)))