def test_from_to_config(): base_layer = tf.keras.layers.Dense(1) sn = spectral_normalization.SpectralNormalization(base_layer) config = sn.get_config() new_sn = spectral_normalization.SpectralNormalization.from_config(config) assert sn.power_iterations == new_sn.power_iterations
def test_model_build(base_layer_fn, input_shape): inputs = tf.keras.layers.Input(shape=input_shape) base_layer = base_layer_fn() sn_layer = spectral_normalization.SpectralNormalization(base_layer) model = tf.keras.models.Sequential(layers=[inputs, sn_layer]) model.build() assert hasattr(model.layers[0], "u")
def test_normalization(): inputs = tf.keras.layers.Input(shape=[2, 2, 1]) base_layer = tf.keras.layers.Conv2D( 1, (2, 2), kernel_initializer=tf.constant_initializer(value=2), ) sn_layer = spectral_normalization.SpectralNormalization(base_layer) model = tf.keras.models.Sequential(layers=[inputs, sn_layer]) weights = np.squeeze(model.layers[0].w.numpy()) # This wrapper normalizes weights by the maximum eigen value eigen_val, _ = np.linalg.eig(weights) weights_normalized = weights / np.max(eigen_val) for training in [False, True]: _ = model( tf.constant(np.ones((1, 2, 2, 1), dtype=np.float32)), training=training, ) if training: w = weights_normalized else: w = weights np.testing.assert_allclose(w, np.squeeze(model.layers[0].w.numpy()))
def test_save_load_model(tmpdir): base_layer = tf.keras.layers.Dense(1) input_shape = [1] inputs = tf.keras.layers.Input(shape=input_shape) sn_layer = spectral_normalization.SpectralNormalization(base_layer) model = tf.keras.models.Sequential(layers=[inputs, sn_layer]) # initialize model model.predict(np.random.uniform(size=(2, 1))) model_path = str(tmpdir / "test.h5") model.save(model_path) new_model = tf.keras.models.load_model(model_path) assert model.layers[0].get_config() == new_model.layers[0].get_config()
def test_apply_layer(): images = tf.ones((1, 2, 2, 1)) sn_wrapper = spectral_normalization.SpectralNormalization( tf.keras.layers.Conv2D( 1, [2, 2], kernel_initializer=tf.constant_initializer(value=1)), input_shape=(2, 2, 1), ) result = sn_wrapper(images, training=False) result_train = sn_wrapper(images, training=True) expected_output = np.array([[[[4.0]]]], dtype=np.float32) np.testing.assert_allclose(result, expected_output) # max eigen value of 2x2 matrix of ones is 2 np.testing.assert_allclose(result_train, expected_output / 2) assert hasattr(sn_wrapper, "u")
def test_model_fit(base_layer_fn, input_shape, output_shape): inputs = tf.keras.layers.Input(shape=input_shape) base_layer = base_layer_fn() sn_layer = spectral_normalization.SpectralNormalization(base_layer) model = tf.keras.models.Sequential(layers=[inputs, sn_layer]) model.add(tf.keras.layers.Activation("relu")) model.compile(optimizer=tf.keras.optimizers.RMSprop(learning_rate=0.001), loss="mse") model.fit( np.random.random((2, *input_shape)), np.random.random((2, *output_shape)), epochs=3, batch_size=10, verbose=0, ) assert hasattr(model.layers[0], "u")
def test_no_kernel(): with pytest.raises(AttributeError): spectral_normalization.SpectralNormalization( tf.keras.layers.MaxPooling2D(2, 2)).build((2, 2))
def test_no_layer(): images = tf.random.uniform((2, 4, 43)) with pytest.raises(AssertionError): spectral_normalization.SpectralNormalization(images)