def test_episode_end_callback_is_called(self):
# Given
model = tf.keras.models.Sequential([
tf.keras.layers.Input((1,)),
tf.keras.layers.BatchNormalization(),
tf.keras.layers.Dense(1)
])
maml = MAML(model, loss=tf.keras.losses.MSE)
meta_train_x = np.array([
[1],
[2],
[3],
], dtype=np.float32)
meta_train_y = np.array([1, 2, 3])
def task_generator():
for i in range(3):
support_set = meta_train_x[i, :], meta_train_y[i]
query_set = meta_train_x[i, :], meta_train_y[i]
yield support_set, query_set
# When
mock_callback = MagicMock()
maml.meta_train(task_generator, n_episode=3, episode_end_callback=mock_callback)
# Then
for args in mock_callback.call_args_list:
self.assertTrue('episode_loss' in args[-1])
def test_meta_train_doesnt_crash_with_model_with_non_trainable_variables(self):
# Given
model = tf.keras.models.Sequential([
tf.keras.layers.Input((1,)),
tf.keras.layers.BatchNormalization(),
tf.keras.layers.Dense(1)
])
maml = MAML(model, loss=tf.keras.losses.MSE)
meta_train_x = np.array([
[1],
[2],
[3],
], dtype=np.float32)
meta_train_y = np.array([1, 2, 3])
def task_generator():
for i in range(3):
support_set = meta_train_x[i, :], meta_train_y[i]
query_set = meta_train_x[i, :], meta_train_y[i]
yield support_set, query_set
# When
maml.meta_train(task_generator, n_episode=3)
# Then
self.assertIsNotNone(model(np.array([[1]])))
def test_meta_learn_produces_right_model_after_1_step(self):
# See calculus derivations in theta.py
# Given
model = create_squared_perceptron()
maml = MAML(model, loss=tf.keras.losses.MSE)
meta_train_x = np.array([
[1],
[2],
[3],
], dtype=np.float32)
meta_train_y = np.array([1, 2, 3])
eval_x = np.array([
[1],
[2],
[3],
], dtype=np.float32) * 1e-10
eval_y = (1 - 2779171867128) ** 2 * eval_x ** 2
def task_generator():
for i in range(3):
support_set = meta_train_x[i, :], meta_train_y[i]
query_set = meta_train_x[i, :], meta_train_y[i]
yield support_set, query_set
# When
optimizer = tf.keras.optimizers.SGD(1.0)
maml.meta_train(task_generator, n_episode=1, alpha=1.0, optimizer=optimizer)
preds = maml.model(eval_x)
# Then
self.assertTrue(np.all(np.abs((preds - eval_y)) < 1e-1))
def test_instantiate_MAML(self):
# Given
model = tf.keras.models.Sequential()
loss = tf.keras.losses.MSE
# When
maml = MAML(model, loss)
# Then
self.assertIsNotNone(maml)
def test_fit_on_two_layer_perceptron_is_correct_when_used_with_a_loss(self):
# Given
model = create_2l_perceptron()
maml = MAML(model, loss=tf.keras.losses.MSE)
data_x = np.ones((1,))
data_y = np.zeros((1, 1))
expected_weights = [
-3 * np.ones((1, 2)),
-3 * np.ones((2, 1)),
]
# When
eval_model = maml.fit(data_x, data_y)
weight_set = [layer.kernel for layer in eval_model.layers]
# Then
for i_weight, weights in enumerate(weight_set):
self.assertTrue(np.all(weights == expected_weights[i_weight]))
def test_fit_on_two_layer_perceptron_is_correct_when_passed_a_value_for_alpha(self):
# Given
model = create_2l_perceptron()
maml = MAML(model, loss=lambda y, p: p)
data_x = np.ones((1,))
data_y = np.zeros((1, 1))
expected_weights = [
0.5 * np.ones((1, 2)),
0.5 * np.ones((2, 1)),
]
# When
eval_model = maml.fit(data_x, data_y, alpha=0.5)
weight_set = [layer.kernel for layer in eval_model.layers]
# Then
for i_weight, weights in enumerate(weight_set):
self.assertTrue(np.all(weights == expected_weights[i_weight]))
def test_task_generator_is_called_n_episode_times(self):
# Given
model = create_squared_perceptron()
maml = MAML(model, loss=tf.keras.losses.MSE)
meta_train_x = np.array([
[1],
[2],
[3],
], dtype=np.float32)
meta_train_y = np.array([1, -1, 7])
def task_generator():
for i in range(3):
support_set = meta_train_x[i, :], meta_train_y[i]
query_set = meta_train_x[i, :], meta_train_y[i]
yield support_set, query_set
# When
magic_task_generator = MagicMock(return_value=task_generator())
maml.meta_train(magic_task_generator, n_episode=7)
# Then
self.assertEqual(magic_task_generator.call_count, 7)
def test_multistep_model_doesnt_diverge(self):
# Given
def sine(n_task=5):
for _ in range(n_task):
x = np.random.uniform(-5, 5, size=(20, 1))
a = np.random.uniform(0.1, 5)
p = np.random.uniform(0, np.pi)
y = a * np.sin(x + p)
support = x[:10], y[:10]
query = x[10:], y[10:]
yield support, query
encoder = tf.keras.models.Sequential([
tf.keras.layers.Input((1,)),
tf.keras.layers.Dense(1),
tf.keras.layers.Dense(40, activation='relu'),
tf.keras.layers.Dense(40, activation='relu'),
tf.keras.layers.Dense(1),
])
model = MAML(encoder, tf.keras.losses.mse)
mock_callback = MagicMock()
def train():
model.meta_train(
sine,
n_episode=2,
alpha=0.01,
n_step=5,
optimizer=tf.keras.optimizers.Adam(1e-3, clipvalue=1.0),
episode_end_callback=mock_callback,
clipvalue=1
)
# When
train()
# Then
episode_loss = np.mean(mock_callback.call_args_list[-1][-1]['episode_loss'].numpy())
self.assertFalse(np.isinf(episode_loss))
self.assertFalse(np.isnan(episode_loss))
def test_meta_train_doesnt_grow_memory_linearly_with_n_episode(self):
# Given
# big beefy model
model = tf.keras.models.Sequential([
tf.keras.layers.Input((1,)),
tf.keras.layers.Dense(100),
tf.keras.layers.Dense(100),
tf.keras.layers.Dense(100),
tf.keras.layers.Dense(100),
tf.keras.layers.Dense(100),
])
maml = MAML(model, loss=tf.keras.losses.MSE)
meta_train_x = np.array([
[1],
[2],
[3],
], dtype=np.float32)
meta_train_y = np.array([1, 2, 3])
def task_generator():
for i in range(3):
support_set = meta_train_x[i, :], meta_train_y[i]
query_set = meta_train_x[i, :], meta_train_y[i]
yield support_set, query_set
# When
tracemalloc.start()
maml.meta_train(task_generator, n_episode=1)
current_mem_one_ep, _ = tracemalloc.get_traced_memory()
tracemalloc.stop()
tracemalloc.start()
maml.meta_train(task_generator, n_episode=10)
current_mem_ten_ep, _ = tracemalloc.get_traced_memory()
tracemalloc.stop()
# Then
self.assertLess(current_mem_ten_ep, 1.2 * current_mem_one_ep)