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', 3)
    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))
    initial_estimator_predict = tf.estimator.Estimator(model_fn=mock_t2r_model.model_fn, config=tf.estimator.RunConfig(model_dir=model_dir))
    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)))]
    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)
    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):
            continue
        if ('Adam' in tensor_name):
            continue
        if ('global_step' in tensor_name):
            continue
        self.assertAllClose(init_checkpoint.get_tensor(tensor_name), continue_checkpoint.get_tensor(tensor_name), atol=0.001)
    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=0.1))
    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=0.01))
Tests that a simple model trains and exported models are valid.
def test_init_from_checkpoint_use_avg_model_params_and_weights(self):
    
    gin.bind_parameter('tf.estimator.RunConfig.save_checkpoints_steps', 100)
    gin.bind_parameter('tf.estimator.RunConfig.keep_checkpoint_max', 3)
    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))
    initial_estimator_predict = tf.estimator.Estimator(model_fn=mock_t2r_model.model_fn, config=tf.estimator.RunConfig(model_dir=model_dir))
    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)))]
    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)
    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):
            continue
        if ('Adam' in tensor_name):
            continue
        if ('global_step' in tensor_name):
            continue
        self.assertAllClose(init_checkpoint.get_tensor(tensor_name), continue_checkpoint.get_tensor(tensor_name), atol=0.001)
    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=0.1))
    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=0.01))