TFX Trainer 组件未将模型导出到文件系统的问题

问题
首先，我使用的是 TFX 0.21.2 版和 Tensorflow 2.1 版。

我建立了一个管道，主要以芝加哥出租车为例。执行 Trainer 组件时，我可以在日志中看到以下内容：

信息培训完成。模型写入 /root/airflow/tfx/pipelines/fish/Trainer/Model/9/serving_model_dir

检查上面的目录时，它是空的。我错过了什么？

这是我的 DAG 定义文件（忽略 import 语句）：

1. _pipeline_name = 'fish'
   
2. _airflow_config = AirflowPipelineConfig(airflow_dag_config = {
   
3.     'schedule_interval': None,
   
4.     'start_date': datetime.datetime(2019, 1, 1),
   
5. })
   
6. _project_root = os.path.join(os.environ['HOME'], 'airflow')
   
7. _data_root = os.path.join(_project_root, 'data', 'fish_data')
   
8. _module_file = os.path.join(_project_root, 'dags', 'fishUtils.py')
   
9. _serving_model_dir = os.path.join(_project_root, 'serving_model', _pipeline_name)
   
10. _tfx_root = os.path.join(_project_root, 'tfx')
    
11. _pipeline_root = os.path.join(_tfx_root, 'pipelines', _pipeline_name)
    
12. _metadata_path = os.path.join(_tfx_root, 'metadata', _pipeline_name,
    
13.                               'metadata.db')
    
14. 
    
15. 
    
16. def _create_pipeline(pipeline_name: Text, pipeline_root: Text, data_root: Text,
    
17.                      module_file: Text, serving_model_dir: Text,
    
18.                      metadata_path: Text,
    
19.                      direct_num_workers: int) -> pipeline.Pipeline:
    
20. 
    
21.     examples = external_input(data_root)
    
22.     example_gen = CsvExampleGen(input=examples)
    
23. 
    
24.     statistics_gen = StatisticsGen(examples=example_gen.outputs['examples'])
    
25. 
    
26.     infer_schema = SchemaGen(
    
27.       statistics=statistics_gen.outputs['statistics'],
    
28.       infer_feature_shape=False)
    
29. 
    
30.     validate_stats = ExampleValidator(
    
31.       statistics=statistics_gen.outputs['statistics'],
    
32.       schema=infer_schema.outputs['schema'])
    
33. 
    
34.     trainer = Trainer(
    
35.     examples=example_gen.outputs['examples'], schema=infer_schema.outputs['schema'],
    
36.     module_file=_module_file, train_args= trainer_pb2.TrainArgs(num_steps=10000),
    
37.     eval_args= trainer_pb2.EvalArgs(num_steps=5000))
    
38. 
    
39.     model_validator = ModelValidator(
    
40.       examples=example_gen.outputs['examples'],
    
41.       model=trainer.outputs['model'])
    
42. 
    
43.     pusher = Pusher(
    
44.       model=trainer.outputs['model'],
    
45.       model_blessing=model_validator.outputs['blessing'],
    
46.       push_destination=pusher_pb2.PushDestination(
    
47.         filesystem=pusher_pb2.PushDestination.Filesystem(
    
48.           base_directory=_serving_model_dir)))
    
49. 
    
50.     return pipeline.Pipeline(
    
51.       pipeline_name=_pipeline_name,
    
52.       pipeline_root=_pipeline_root,
    
53.       components=[
    
54.           example_gen,
    
55.           statistics_gen,
    
56.           infer_schema,
    
57.           validate_stats,
    
58.           trainer,
    
59.           model_validator,
    
60.           pusher],
    
61.       enable_cache=True,
    
62.       metadata_connection_config=metadata.sqlite_metadata_connection_config(
    
63.           metadata_path),
    
64.       beam_pipeline_args=['--direct_num_workers=%d' % direct_num_workers]
    
65.   )
    
66. 
    
67. runner = AirflowDagRunner(config = _airflow_config)
    
68. DAG = runner.run(
    
69.     _create_pipeline(
    
70.         pipeline_name=_pipeline_name,
    
71.         pipeline_root=_pipeline_root,
    
72.         data_root=_data_root,
    
73.         module_file=_module_file,
    
74.         serving_model_dir=_serving_model_dir,
    
75.         metadata_path=_metadata_path,
    
76.         # 0 means auto-detect based on on the number of CPUs available during
    
77.         # execution time.
    
78.         direct_num_workers=0))

复制代码

这是我的模块文件：

1. _DENSE_FLOAT_FEATURE_KEYS = ['length']
   
2. 
   
3. real_valued_columns = [tf.feature_column.numeric_column('length')]
   
4. 
   
5. def _eval_input_receiver_fn():
   
6. 
   
7.   serialized_tf_example = tf.compat.v1.placeholder(
   
8.       dtype=tf.string, shape=[None], name='input_example_tensor')
   
9. 
   
10.   features = tf.io.parse_example(
    
11.       serialized=serialized_tf_example,
    
12.       features={
    
13.           'length': tf.io.FixedLenFeature([], tf.float32),
    
14.           'label': tf.io.FixedLenFeature([], tf.int64),
    
15.       })
    
16. 
    
17.   receiver_tensors = {'examples': serialized_tf_example}
    
18. 
    
19.   return tfma.export.EvalInputReceiver(
    
20.       features={'length' : features['length']},
    
21.       receiver_tensors=receiver_tensors,
    
22.       labels= features['label'],
    
23.       )
    
24. 
    
25. def parser(serialized_example):
    
26. 
    
27.   features = tf.io.parse_single_example(
    
28.       serialized_example,
    
29.       features={
    
30.           'length': tf.io.FixedLenFeature([], tf.float32),
    
31.           'label': tf.io.FixedLenFeature([], tf.int64),
    
32.       })
    
33.   return ({'length' : features['length']}, features['label'])
    
34. 
    
35. def _input_fn(filenames):
    
36.   # TFRecordDataset doesn't directly accept paths with wildcards
    
37.   filenames = tf.data.Dataset.list_files(filenames)
    
38.   dataset = tf.data.TFRecordDataset(filenames, 'GZIP')
    
39.   dataset = dataset.map(parser)
    
40.   dataset = dataset.shuffle(2000)
    
41.   dataset = dataset.batch(40)
    
42.   dataset = dataset.repeat(10)
    
43. 
    
44.   return dataset
    
45. 
    
46. def trainer_fn(trainer_fn_args, schema):
    
47. 
    
48.     estimator = tf.estimator.LinearClassifier(feature_columns=real_valued_columns)
    
49. 
    
50.     train_input_fn = lambda: _input_fn(trainer_fn_args.train_files)
    
51. 
    
52.     train_spec = tf.estimator.TrainSpec(
    
53.       train_input_fn,
    
54.       max_steps=trainer_fn_args.train_steps)
    
55. 
    
56.     eval_input_fn = lambda: _input_fn(trainer_fn_args.eval_files)
    
57. 
    
58.     eval_spec = tf.estimator.EvalSpec(
    
59.       eval_input_fn,
    
60.       steps=trainer_fn_args.eval_steps,
    
61.       name='fish-eval')
    
62. 
    
63.     receiver_fn = lambda: _eval_input_receiver_fn()
    
64. 
    
65.     return {
    
66.       'estimator': estimator,
    
67.       'train_spec': train_spec,
    
68.       'eval_spec': eval_spec,
    
69.       'eval_input_receiver_fn': receiver_fn
    
70.   }

复制代码

在此先感谢您的帮助！

回答
为遇到与我相同问题的任何人发布解决方案。

模型没有写入文件系统的原因是估计器需要一个配置参数来知道在哪里写入模型。

以下对 trainer_fn 函数的修改应该可以解决问题。

1. run_config = tf.estimator.RunConfig(save_checkpoints_steps=999, keep_checkpoint_max=1)  
   
2. 
   
3. run_config = run_config.replace(model_dir=trainer_fn_args.serving_model_dir)
   
4. 
   
5. estimator=tf.estimator.LinearClassifier(feature_columns=real_valued_columns,config=run_config)

复制代码