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)