def test_init_from_checkpoint_global_step(self): """Tests that a simple model trains and exported models are valid.""" gin.bind_parameter('tf.estimator.RunConfig.save_checkpoints_steps', 100) gin.bind_parameter('tf.estimator.RunConfig.keep_checkpoint_max', 10) model_dir = self.create_tempdir().full_path mock_t2r_model = mocks.MockT2RModel( preprocessor_cls=noop_preprocessor.NoOpPreprocessor) mock_input_generator_train = mocks.MockInputGenerator( batch_size=_BATCH_SIZE) train_eval.train_eval_model( t2r_model=mock_t2r_model, input_generator_train=mock_input_generator_train, max_train_steps=_MAX_TRAIN_STEPS, model_dir=model_dir, eval_steps=_EVAL_STEPS, eval_throttle_secs=_EVAL_THROTTLE_SECS, create_exporters_fn=train_eval.create_default_exporters) # The model trains for 1000 steps and saves a checkpoint each 100 steps and # keeps 10 -> len == 10. self.assertLen( tf.io.gfile.glob(os.path.join(model_dir, 'model*.meta')), 10) # The continuous training has its own directory. continue_model_dir = self.create_tempdir().full_path init_from_checkpoint_fn = functools.partial( abstract_model.default_init_from_checkpoint_fn, checkpoint=model_dir) continue_mock_t2r_model = mocks.MockT2RModel( preprocessor_cls=noop_preprocessor.NoOpPreprocessor, init_from_checkpoint_fn=init_from_checkpoint_fn) continue_mock_input_generator_train = mocks.MockInputGenerator( batch_size=_BATCH_SIZE) train_eval.train_eval_model( t2r_model=continue_mock_t2r_model, input_generator_train=continue_mock_input_generator_train, model_dir=continue_model_dir, max_train_steps=_MAX_TRAIN_STEPS + 100, eval_steps=_EVAL_STEPS, eval_throttle_secs=_EVAL_THROTTLE_SECS, create_exporters_fn=train_eval.create_default_exporters) # If the model was successful restored including the global step, only 1 # additional checkpoint to the init one should be created -> len == 2. self.assertLen( tf.io.gfile.glob(os.path.join(continue_model_dir, 'model*.meta')), 2)
def test_create_serving_input_receiver_numpy(self): (model_dir, mock_t2r_model, prediction_ref) = self._train_and_eval_reference_model('numpy') exporter = default_export_generator.DefaultExportGenerator() exporter.set_specification_from_model(mock_t2r_model) # Export trained serving estimator. estimator_exporter = tf.estimator.Estimator( model_fn=mock_t2r_model.model_fn, config=tf.estimator.RunConfig(model_dir=model_dir)) serving_input_receiver_fn = ( exporter.create_serving_input_receiver_numpy_fn()) exported_savedmodel_path = estimator_exporter.export_saved_model( export_dir_base=model_dir, serving_input_receiver_fn=serving_input_receiver_fn, checkpoint_path=tf.train.latest_checkpoint(model_dir)) # Load trained and exported serving estimator, run prediction and assert # it is the same as before exporting. feed_predictor_fn = tf.contrib.predictor.from_saved_model( exported_savedmodel_path) mock_input_generator = mocks.MockInputGenerator(batch_size=BATCH_SIZE) features, labels = mock_input_generator.create_numpy_data() for pos, value in enumerate(prediction_ref): actual = feed_predictor_fn({'x': features[pos, :].reshape(1, -1) })['logit'].flatten() predicted = value['logit'].flatten() np.testing.assert_almost_equal(actual=actual, desired=predicted, decimal=4) if labels[pos] > 0: self.assertGreater(predicted[0], 0) else: self.assertLess(predicted[0], 0)
def _train_and_eval_reference_model(self, path): model_dir = self.create_tempdir().full_path mock_t2r_model = mocks.MockT2RModel( preprocessor_cls=noop_preprocessor.NoOpPreprocessor) # We create a tpu estimator for potential training. estimator = tf.contrib.tpu.TPUEstimator( model_fn=mock_t2r_model.model_fn, use_tpu=mock_t2r_model.is_device_tpu, config=tf.contrib.tpu.RunConfig(model_dir=model_dir), train_batch_size=BATCH_SIZE, eval_batch_size=BATCH_SIZE) mock_input_generator = mocks.MockInputGenerator(batch_size=BATCH_SIZE) mock_input_generator.set_specification_from_model( mock_t2r_model, tf.estimator.ModeKeys.TRAIN) # We optimize our network. estimator.train(input_fn=mock_input_generator.create_dataset_input_fn( mode=tf.estimator.ModeKeys.TRAIN), max_steps=MAX_STEPS) # Verify that the serving estimator does exactly the same as the normal # estimator with all the parameters. estimator_predict = tf.estimator.Estimator( model_fn=mock_t2r_model.model_fn, config=tf.estimator.RunConfig(model_dir=model_dir)) prediction_ref = estimator_predict.predict( input_fn=mock_input_generator.create_dataset_input_fn( mode=tf.estimator.ModeKeys.EVAL)) return model_dir, mock_t2r_model, prediction_ref
def test_with_mock_training(self): model_dir = self.create_tempdir().full_path mock_t2r_model = mocks.MockT2RModel( preprocessor_cls=noop_preprocessor.NoOpPreprocessor, device_type='tpu', use_avg_model_params=True) mock_input_generator = mocks.MockInputGenerator(batch_size=_BATCH_SIZE) export_dir = os.path.join(model_dir, _EXPORT_DIR) hook_builder = async_export_hook_builder.AsyncExportHookBuilder( export_dir=export_dir, create_export_fn=async_export_hook_builder.default_create_export_fn ) gin.parse_config('tf.contrib.tpu.TPUConfig.iterations_per_loop=1') gin.parse_config('tf.estimator.RunConfig.save_checkpoints_steps=1') # We optimize our network. train_eval.train_eval_model(t2r_model=mock_t2r_model, input_generator_train=mock_input_generator, train_hook_builders=[hook_builder], model_dir=model_dir, max_train_steps=_MAX_STEPS) self.assertNotEmpty(tf.io.gfile.listdir(model_dir)) self.assertNotEmpty(tf.io.gfile.listdir(export_dir)) for exported_model_dir in tf.io.gfile.listdir(export_dir): self.assertNotEmpty( tf.io.gfile.listdir( os.path.join(export_dir, exported_model_dir))) predictor = exported_savedmodel_predictor.ExportedSavedModelPredictor( export_dir=export_dir) self.assertTrue(predictor.restore())
def test_with_mock_training(self): model_dir = self.create_tempdir().full_path mock_t2r_model = mocks.MockT2RModel( preprocessor_cls=noop_preprocessor.NoOpPreprocessor, device_type='cpu') mock_input_generator = mocks.MockInputGenerator(batch_size=_BATCH_SIZE) default_create_export_fn = functools.partial( async_export_hook_builder.default_create_export_fn, batch_sizes_for_export=_BATCH_SIZES_FOR_EXPORT) export_dir = os.path.join(model_dir, _EXPORT_DIR) default_create_export_fn = functools.partial( async_export_hook_builder.default_create_export_fn, batch_sizes_for_export=_BATCH_SIZES_FOR_EXPORT) hook_builder = async_export_hook_builder.AsyncExportHookBuilder( export_dir=export_dir, create_export_fn=default_create_export_fn) default_create_export_fn = functools.partial( async_export_hook_builder.default_create_export_fn, batch_sizes_for_export=_BATCH_SIZES_FOR_EXPORT) # We optimize our network. train_eval.train_eval_model( t2r_model=mock_t2r_model, input_generator_train=mock_input_generator, train_hook_builders=[hook_builder], model_dir=model_dir, max_train_steps=_MAX_STEPS) self.assertNotEmpty(tf.io.gfile.listdir(model_dir)) self.assertNotEmpty(tf.io.gfile.listdir(export_dir)) for exported_model_dir in tf.io.gfile.listdir(export_dir): self.assertNotEmpty( tf.io.gfile.listdir(os.path.join(export_dir, exported_model_dir))) predictor = exported_savedmodel_predictor.ExportedSavedModelPredictor( export_dir=export_dir) self.assertTrue(predictor.restore())
def test_set_preprocess_fn(self): mock_input_generator = mocks.MockInputGenerator(batch_size=BATCH_SIZE) preprocessor = noop_preprocessor.NoOpPreprocessor() with self.assertRaises(ValueError): # This should raise since we pass a function with `mode` not already # filled in either by a closure or functools.partial. mock_input_generator.set_preprocess_fn(preprocessor.preprocess) preprocess_fn = functools.partial(preprocessor.preprocess, labels=None) with self.assertRaises(ValueError): # This should raise since we pass a partial function but `mode` # is not abstracted away. mock_input_generator.set_preprocess_fn(preprocess_fn)
def test_create_serving_input_receiver_tf_example(self, multi_dataset): (model_dir, mock_t2r_model, prediction_ref) = self._train_and_eval_reference_model( 'tf_example', multi_dataset=multi_dataset) # Now we can actually export our serving estimator. estimator_exporter = tf.estimator.Estimator( model_fn=mock_t2r_model.model_fn, config=tf.estimator.RunConfig(model_dir=model_dir)) exporter = default_export_generator.DefaultExportGenerator() exporter.set_specification_from_model(mock_t2r_model) serving_input_receiver_fn = ( exporter.create_serving_input_receiver_tf_example_fn()) exported_savedmodel_path = estimator_exporter.export_saved_model( export_dir_base=model_dir, serving_input_receiver_fn=serving_input_receiver_fn, checkpoint_path=tf.train.latest_checkpoint(model_dir)) # Now we can load our exported estimator graph, there are no dependencies # on the model_fn or preprocessor anymore. feed_predictor_fn = tf.contrib.predictor.from_saved_model( exported_savedmodel_path) mock_input_generator = mocks.MockInputGenerator(batch_size=BATCH_SIZE) features, labels = mock_input_generator.create_numpy_data() for pos, value in enumerate(prediction_ref): # We have to create our serialized tf.Example proto. example = tf.train.Example() example.features.feature[ 'measured_position'].float_list.value.extend(features[pos]) serialized_example = np.array(example.SerializeToString()).reshape( 1, ) if multi_dataset: feed_dict = { 'input_example_dataset1': serialized_example, 'input_example_dataset2': serialized_example } else: feed_dict = {'input_example_tensor': serialized_example} actual = feed_predictor_fn(feed_dict)['logit'].flatten() predicted = value['logit'].flatten() np.testing.assert_almost_equal(actual=actual, desired=predicted, decimal=4) if labels[pos] > 0: self.assertGreater(predicted[0], 0) else: self.assertLess(predicted[0], 0)
def test_freezing_some_variables(self): """Tests we can freeze training for parts of the network.""" def freeze_biases(var): # Update all variables except bias variables. return 'bias' not in var.name gin.bind_parameter('tf.estimator.RunConfig.save_checkpoints_steps', 100) gin.bind_parameter('create_train_op.filter_trainables_fn', freeze_biases) model_dir = self.create_tempdir().full_path mock_t2r_model = mocks.MockT2RModel( preprocessor_cls=noop_preprocessor.NoOpPreprocessor) mock_input_generator_train = mocks.MockInputGenerator( batch_size=_BATCH_SIZE) train_eval.train_eval_model( t2r_model=mock_t2r_model, input_generator_train=mock_input_generator_train, max_train_steps=_MAX_TRAIN_STEPS, model_dir=model_dir) start_checkpoint = tf.train.NewCheckpointReader( os.path.join(model_dir, 'model.ckpt-0')) last_checkpoint = tf.train.NewCheckpointReader( tf.train.latest_checkpoint(model_dir)) for var_name, _ in tf.train.list_variables(model_dir): # Some of the batch norm moving averages are constant over training on the # mock data used. if 'batch_norm' in var_name: continue if 'bias' not in var_name: # Should update. self.assertNotAllClose(start_checkpoint.get_tensor(var_name), last_checkpoint.get_tensor(var_name), atol=1e-3) else: # Should not update. self.assertAllClose(start_checkpoint.get_tensor(var_name), last_checkpoint.get_tensor(var_name), atol=1e-3)
def test_train_eval_model(self): """Tests that a simple model trains and exported models are valid.""" gin.bind_parameter('tf.estimator.RunConfig.save_checkpoints_steps', 100) model_dir = self.create_tempdir().full_path mock_t2r_model = mocks.MockT2RModel( preprocessor_cls=noop_preprocessor.NoOpPreprocessor) mock_input_generator_train = mocks.MockInputGenerator( batch_size=_BATCH_SIZE) mock_input_generator_eval = mocks.MockInputGenerator(batch_size=1) fake_hook_builder = FakeHookBuilder() train_eval.train_eval_model( t2r_model=mock_t2r_model, input_generator_train=mock_input_generator_train, input_generator_eval=mock_input_generator_eval, max_train_steps=_MAX_TRAIN_STEPS, model_dir=model_dir, train_hook_builders=[fake_hook_builder], eval_hook_builders=[fake_hook_builder], eval_steps=_EVAL_STEPS, eval_throttle_secs=_EVAL_THROTTLE_SECS, create_exporters_fn=train_eval.create_default_exporters) self.assertTrue(fake_hook_builder.hook_mock.begin.called) # We ensure that both numpy and tf_example inference models are exported. best_exporter_numpy_path = os.path.join(model_dir, 'export', 'best_exporter_numpy', '*') numpy_model_paths = sorted(tf.io.gfile.glob(best_exporter_numpy_path)) # This mock network converges nicely which is why we have several best # models, by default we keep the best 5 and the latest one is always the # best. self.assertLen(numpy_model_paths, 5) best_exporter_tf_example_path = os.path.join( model_dir, 'export', 'best_exporter_tf_example', '*') tf_example_model_paths = sorted( tf.io.gfile.glob(best_exporter_tf_example_path)) # This mock network converges nicely which is why we have several best # models, by default we keep the best 5 and the latest one is always the # best. self.assertLen(tf_example_model_paths, 5) # We test both saved models within one test since the bulk of the time # is spent training the model in the firstplace. # Verify that the serving estimator does exactly the same as the normal # estimator with all the parameters. estimator_predict = tf.estimator.Estimator( model_fn=mock_t2r_model.model_fn, config=tf.estimator.RunConfig(model_dir=model_dir)) prediction_ref = estimator_predict.predict( input_fn=mock_input_generator_eval.create_dataset_input_fn( mode=tf.estimator.ModeKeys.EVAL)) # Now we can load our exported estimator graph with the numpy feed_dict # interface, there are no dependencies on the model_fn or preprocessor # anymore. # We load the latest model since it had the best eval performance. numpy_predictor_fn = tf.contrib.predictor.from_saved_model( numpy_model_paths[-1]) features, labels = mock_input_generator_eval.create_numpy_data() ref_error = self._compute_total_loss( labels, [val['logit'].flatten() for val in prediction_ref]) numpy_predictions = [] for feature, label in zip(features, labels): predicted = numpy_predictor_fn({'x': feature.reshape(1, -1) })['logit'].flatten() numpy_predictions.append(predicted) # This ensures that we actually achieve perfect classification. if label > 0: self.assertGreater(predicted[0], 0) else: self.assertLess(predicted[0], 0) numpy_error = self._compute_total_loss(labels, numpy_predictions) # Now we can load our exported estimator graph with the tf_example feed_dict # interface, there are no dependencies on the model_fn or preprocessor # anymore. # We load the latest model since it had the best eval performance. tf_example_predictor_fn = tf.contrib.predictor.from_saved_model( tf_example_model_paths[-1]) tf_example_predictions = [] for feature, label in zip(features, labels): # We have to create our serialized tf.Example proto. example = tf.train.Example() example.features.feature[ 'measured_position'].float_list.value.extend(feature) feed_dict = { 'input_example_tensor': np.array(example.SerializeToString()).reshape(1, ) } predicted = tf_example_predictor_fn(feed_dict)['logit'].flatten() tf_example_predictions.append(predicted) # This ensures that we actually achieve perfect classification. if label > 0: self.assertGreater(predicted[0], 0) else: self.assertLess(predicted[0], 0) tf_example_error = self._compute_total_loss(labels, tf_example_predictions) np.testing.assert_almost_equal(tf_example_error, numpy_error) # The exported saved models both have to have the same performance and since # we train on eval on the same fixed dataset the latest and greatest # model error should also be the best. np.testing.assert_almost_equal(ref_error, tf_example_error)
def test_init_from_checkpoint_use_avg_model_params_and_weights(self): """Tests that a simple model trains and exported models are valid.""" gin.bind_parameter('tf.estimator.RunConfig.save_checkpoints_steps', 100) gin.bind_parameter('tf.estimator.RunConfig.keep_checkpoint_max', 10) model_dir = self.create_tempdir().full_path mock_t2r_model = mocks.MockT2RModel( preprocessor_cls=noop_preprocessor.NoOpPreprocessor, use_avg_model_params=True) mock_input_generator_train = mocks.MockInputGenerator( batch_size=_BATCH_SIZE) mock_input_generator = mocks.MockInputGenerator(batch_size=1) mock_input_generator.set_specification_from_model( mock_t2r_model, tf.estimator.ModeKeys.TRAIN) train_eval.train_eval_model( t2r_model=mock_t2r_model, input_generator_train=mock_input_generator_train, max_train_steps=_MAX_TRAIN_STEPS, model_dir=model_dir) init_checkpoint = tf.train.NewCheckpointReader( tf.train.latest_checkpoint(model_dir)) # Verify that the serving estimator does exactly the same as the normal # estimator with all the parameters. initial_estimator_predict = tf.estimator.Estimator( model_fn=mock_t2r_model.model_fn, config=tf.estimator.RunConfig(model_dir=model_dir)) # pylint: disable=g-complex-comprehension initial_predictions = [ prediction['logit'] for prediction in list( initial_estimator_predict.predict( input_fn=mock_input_generator.create_dataset_input_fn( mode=tf.estimator.ModeKeys.EVAL))) ] # The continuous training has its own directory. continue_model_dir = self.create_tempdir().full_path init_from_checkpoint_fn = functools.partial( abstract_model.default_init_from_checkpoint_fn, checkpoint=model_dir) continue_mock_t2r_model = mocks.MockT2RModel( preprocessor_cls=noop_preprocessor.NoOpPreprocessor, init_from_checkpoint_fn=init_from_checkpoint_fn) continue_mock_input_generator_train = mocks.MockInputGenerator( batch_size=_BATCH_SIZE) # Re-initialize the model and train for one step, basically the same # performance as the original model. train_eval.train_eval_model( t2r_model=continue_mock_t2r_model, input_generator_train=continue_mock_input_generator_train, model_dir=continue_model_dir, max_train_steps=_MAX_TRAIN_STEPS) continue_checkpoint = tf.train.NewCheckpointReader( tf.train.latest_checkpoint(continue_model_dir)) for tensor_name, _ in tf.train.list_variables(model_dir): if 'ExponentialMovingAverage' in tensor_name: # These values are replaced by the swapping saver when using the # use_avg_model_params. continue if 'Adam' in tensor_name: # The adam optimizer values are not required. continue if 'global_step' in tensor_name: # The global step will be incremented by 1. continue self.assertAllClose(init_checkpoint.get_tensor(tensor_name), continue_checkpoint.get_tensor(tensor_name), atol=1e-3) # Verify that the serving estimator does exactly the same as the normal # estimator with all the parameters. continue_estimator_predict = tf.estimator.Estimator( model_fn=mock_t2r_model.model_fn, config=tf.estimator.RunConfig(model_dir=continue_model_dir)) continue_predictions = [ prediction['logit'] for prediction in list( continue_estimator_predict.predict( input_fn=mock_input_generator.create_dataset_input_fn( mode=tf.estimator.ModeKeys.EVAL))) ] self.assertTrue( np.allclose(initial_predictions, continue_predictions, atol=1e-2)) # A randomly initialized model estimator with all the parameters. random_estimator_predict = tf.estimator.Estimator( model_fn=mock_t2r_model.model_fn) random_predictions = [ prediction['logit'] for prediction in list( random_estimator_predict.predict( input_fn=mock_input_generator.create_dataset_input_fn( mode=tf.estimator.ModeKeys.EVAL))) ] self.assertFalse( np.allclose(initial_predictions, random_predictions, atol=1e-2))
def test_init_from_checkpoint_use_avg_model_params(self): """Tests that a simple model trains and exported models are valid.""" gin.bind_parameter('tf.estimator.RunConfig.save_checkpoints_steps', 100) gin.bind_parameter('tf.estimator.RunConfig.keep_checkpoint_max', 10) model_dir = self.create_tempdir().full_path mock_t2r_model = mocks.MockT2RModel( preprocessor_cls=noop_preprocessor.NoOpPreprocessor, use_avg_model_params=True) mock_input_generator_train = mocks.MockInputGenerator( batch_size=_BATCH_SIZE) mock_input_generator = mocks.MockInputGenerator(batch_size=1) mock_input_generator.set_specification_from_model( mock_t2r_model, tf.estimator.ModeKeys.TRAIN) train_eval.train_eval_model( t2r_model=mock_t2r_model, input_generator_train=mock_input_generator_train, max_train_steps=_MAX_TRAIN_STEPS, model_dir=model_dir) # Verify that the serving estimator does exactly the same as the normal # estimator with all the parameters. initial_estimator_predict = tf.estimator.Estimator( model_fn=mock_t2r_model.model_fn, config=tf.estimator.RunConfig(model_dir=model_dir)) # pylint: disable=g-complex-comprehension initial_predictions = [ prediction['logit'] for prediction in list( initial_estimator_predict.predict( input_fn=mock_input_generator.create_dataset_input_fn( mode=tf.estimator.ModeKeys.EVAL))) ] # The continuous training has its own directory. continue_model_dir = self.create_tempdir().full_path init_from_checkpoint_fn = functools.partial( abstract_model.default_init_from_checkpoint_fn, checkpoint=model_dir) continue_mock_t2r_model = mocks.MockT2RModel( preprocessor_cls=noop_preprocessor.NoOpPreprocessor, init_from_checkpoint_fn=init_from_checkpoint_fn) continue_mock_input_generator_train = mocks.MockInputGenerator( batch_size=_BATCH_SIZE) # Re-initialize the model and train for one step, basically the same # performance as the original model. train_eval.train_eval_model( t2r_model=continue_mock_t2r_model, input_generator_train=continue_mock_input_generator_train, model_dir=continue_model_dir, max_train_steps=_MAX_TRAIN_STEPS + 1) # Verify that the serving estimator does exactly the same as the normal # estimator with all the parameters. continue_estimator_predict = tf.estimator.Estimator( model_fn=mock_t2r_model.model_fn, config=tf.estimator.RunConfig(model_dir=continue_model_dir)) continue_predictions = [ prediction['logit'] for prediction in list( continue_estimator_predict.predict( input_fn=mock_input_generator.create_dataset_input_fn( mode=tf.estimator.ModeKeys.EVAL))) ] self.assertTrue( np.allclose(initial_predictions, continue_predictions, atol=1e-2)) # A randomly initialized model estimator with all the parameters. random_estimator_predict = tf.estimator.Estimator( model_fn=mock_t2r_model.model_fn) random_predictions = [ prediction['logit'] for prediction in list( random_estimator_predict.predict( input_fn=mock_input_generator.create_dataset_input_fn( mode=tf.estimator.ModeKeys.EVAL))) ] self.assertFalse( np.allclose(initial_predictions, random_predictions, atol=1e-2))