def test_get_latest_model(self):
"""Test latest model is returned as a dictionary."""
# Model_2 is newer than model_1, so we will get model_2.
model_paths = ml_train_task.get_last_saved_model(MODEL_DIR)
expected = {
'meta': self.model_2_meta_path,
'data': self.model_2_data_path,
'index': self.model_2_index_path
}
self.assertDictEqual(model_paths, expected)
def test_train_rnn(self):
"""Test train RNN model on a simple corpus."""
# No model exists in model directory.
self.assertFalse(
ml_train_task.get_last_saved_model(self.model_directory))
# The training should be fast (a few seconds) since sample corpus is
# extremely small.
result = ml_train_task.train_rnn(self.input_directory,
self.model_directory,
self.log_directory, self.batch_size,
self.hidden_state_size,
self.hidden_layer_size)
self.assertEqual(result.return_code, constants.ExitCode.SUCCESS)
self.assertFalse(result.timed_out)
# At least one model exists.
self.assertTrue(
ml_train_task.get_last_saved_model(self.model_directory))
def test_no_valid_model(self):
"""Test no model is returned if all models are invalid."""
# Remove one file from model_1 and one from model_2.
os.remove(self.model_1_data_path)
os.remove(self.model_2_index_path)
self.assertFalse(os.path.exists(self.model_1_data_path))
self.assertFalse(os.path.exists(self.model_2_index_path))
# Now we should get empty dictionary since both models are invalid.
model_paths = ml_train_task.get_last_saved_model(MODEL_DIR)
expected = {}
self.assertDictEqual(model_paths, expected)
def test_get_valid_model(self):
"""Test lastest model is not returned if it is invalid."""
# Remove one file from model_2, so model_2 is not valid.
os.remove(self.model_2_index_path)
self.assertFalse(os.path.exists(self.model_2_index_path))
# Now we should get model_1.
model_paths = ml_train_task.get_last_saved_model(MODEL_DIR)
expected = {
'meta': self.model_1_meta_path,
'data': self.model_1_data_path,
'index': self.model_1_index_path
}
self.assertDictEqual(model_paths, expected)
def test_small_corpus(self):
"""Test small corpus situation."""
# Increase batch size so the sample corpus appears small in this case.
self.batch_size = 100
result = ml_train_task.train_rnn(self.input_directory,
self.model_directory,
self.log_directory, self.batch_size,
self.hidden_state_size,
self.hidden_layer_size)
self.assertEqual(result.return_code,
constants.ExitCode.CORPUS_TOO_SMALL)
self.assertFalse(result.timed_out)
# No model exsits after execution.
self.assertFalse(
ml_train_task.get_last_saved_model(self.model_directory))