Merge pull request #40 from lienhua34/dist-base

调整支持原生 TensorFlow API 的 TaaS 分布式模型训练任务代码模板

Author: lienhua34

caicloud.tensorflow/caicloud/clever/examples/dist-tf-template.py renamed to caicloud.tensorflow/caicloud/clever/templates/template-raw.py

@@ -8,6 +8,16 @@
 from caicloud.clever.tensorflow import dist_base
 from caicloud.clever.tensorflow import model_exporter
 
+tf.app.flags.DEFINE_string("export_dir",
+                           "/tmp/mnist/saved_model",
+                           "model export directory path.")
+tf.app.flags.DEFINE_string("checkpoint_dir",
+                           "",
+                           "model checkpoint directory path.")
+FLAGS = tf.app.flags.FLAGS
+
+_train_op = None
+
 def model_fn(sync, num_replicas):
     """TensorFlow 模型定义函数。
 
@@ -18,13 +28,33 @@ def model_fn(sync, num_replicas):
     - `sync`：当前是否采用参数同步更新模式。 
     - `num_replicas`：分布式 TensorFlow 的计算节点（worker）个数。
     """
+    global _train_op
 
     # TODO：添加业务模型定义操作。
-
+    # global_step = ...
+    # _train_op = ...
+
+    # 添加模型评估配置：
+    # accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
+    # def accuracy_evalute_fn(session):
+    #     return session.run(accuracy, ...)
+    # model_metric_ops = {
+    #    "accuracy": accuracy_evalute_fn
+    # }
+
+    # 定义模型导出配置
+    # model_export_spec = model_exporter.ModelExportSpec(
+    #    export_dir=FLAGS.export_dir,
+    #    input_tensors={"image": _input_images},
+    #    output_tensors={"logits": logits})
+    
     # model_fn 函数需要返回 ModelFnHandler 对象告知 TaaS 平台所构建的模型的一些信息，
     # 例如 global_step、优化器 Optimizer、模型评估指标以及模型导出的相关配置等等。
     # 详细信息请参考 docs.caicloud.io。
-    return dist_base.ModelFnHandler()
+    return dist_base.ModelFnHandler(
+        global_step = global_step,
+        model_metric_ops = model_metric_ops,
+        model_export_spec = model_export_spec)
 
 def train_fn(session, num_global_step):
     """模型训练的每一轮操作。
@@ -38,21 +68,41 @@ def train_fn(session, num_global_step):
 
     # TODO：添加业务模型训练操作。
 
-    # train_fn 函数返回一个 bool 值，用于告知 TaaS 平台是否要提前终止模型训练。返回 True，
-    # 表示终止训练；否则，TaaS 将继续下一轮训练。
-    # 例如，为了防止训练模型过拟合，在训练过程中定时使用验证数据评测模型效果。如果发现模型
-    # 在训练数据集上的效果有优化，而在验证数据集上的效果却开始劣化，则说明模型可能出现了过
-    # 拟合，此时我们就可以通过返回 True 来告知 TaaS 平台提前终止模型训练。
+    # train_fn 函数返回一个 bool 值，用于告知 TaaS 平台是否要提前终止模型训练。
+    # 返回 True，表示终止训练；否则，TaaS 将继续下一轮训练。
+    # 例如，为了防止训练模型过拟合，在训练过程中定时使用验证数据评测模型效果。当模型效果
+    # 达到预期效果，便可以通过返回 True 来结束模型训练。
     return False
 
 def gen_init_fn():
     """获取自定义初始化函数。
     
     有些情况下，我需要从某个事先训练好的 checkpoint 文件中加载模型的参数。此时，我们需要自
     己实现使用 tf.Saver() 从该 checkpoint 中加载模型参数进行自定义初始化的函数。
+
+    注：如果不需要自定义初始化，可以不提供 gen_init_fn 实现，或者 gen_init_fn 返回 None。
     """
-    return None
 
+    # TODO: 添加自己的处理逻辑
+
+    # 定义 tf.train.Saver 会修改 TensorFlow 的 Graph 结构，
+    # 而当 Base 框架调用自定义初始化函数 init_from_checkpoint 的时候，
+    # TensorFlow 模型的 Graph 结构已经变成 finalized，不再允许修改 Graph 结构。
+    # 所以，这个定义必须放在  init_from_checkpoint 函数外面。
+    saver = tf.train.Saver(tf.trainable_variables())
+
+    def init_from_checkpoint(scaffold, sess):
+        """执行自定义初始化的函数。
+
+        TaaS 平台会优先从设置的日志保存路径中获取最新的 checkpoint 来 restore 模型参数，
+        如果日志保存路径中找不到 checkpoint 文件，才会调用本函数来进行模型初始化。
+
+        本函数必须接收两个参数：
+          - scafford: tf.train.Scaffold 对象；
+          - sess: tf.Session 对象。
+        """
+        saver.restore(sess, checkpoint_path)
+    return init_from_checkpoint
 
 def after_train_hook(session):
     """模型训练操作。