def create_model_fn(features, labels, mode, params=None): """Constructs the classification model. Modifed from: https://github.com/tensorflow/models/blob/master/research/ object_detection/model_lib.py. Args: features: A 4-D float32 tensor with shape [batch_size, height, width, channels] representing a batch of images. (Support dict) labels: A 1-D int32 tensor with shape [batch_size] representing the labels of each image. (Support dict) mode: Mode key for tf.estimator.ModeKeys. params: Parameter dictionary passed from the estimator. Returns: An `EstimatorSpec` the encapsulates the model and its serving configurations. """ params = params or {} loss, acc, train_op, export_outputs = None, None, None, None is_training = mode == tf.estimator.ModeKeys.TRAIN cls_model = model.Model(is_training=is_training, num_classes=FLAGS.num_classes) preprocessed_inputs = cls_model.preprocess(features.get('image')) prediction_dict = cls_model.predict(preprocessed_inputs) postprocessed_dict = cls_model.postprocess(prediction_dict) if mode == tf.estimator.ModeKeys.TRAIN: if FLAGS.checkpoint_path: init_variables_from_checkpoint() if mode in (tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL): loss_dict = cls_model.loss(prediction_dict, labels) loss = loss_dict['loss'] classes = postprocessed_dict['classes'] acc = tf.reduce_mean(tf.cast(tf.equal(classes, labels), 'float')) tf.summary.scalar('loss', loss) tf.summary.scalar('accuracy', acc) scaffold = None if mode == tf.estimator.ModeKeys.TRAIN: global_step = tf.train.get_or_create_global_step() learning_rate = configure_learning_rate(FLAGS.decay_steps, global_step) optimizer = tf.train.MomentumOptimizer(learning_rate=learning_rate, momentum=0.9) train_op = slim.learning.create_train_op(loss, optimizer, summarize_gradients=True) keep_checkpoint_every_n_hours = FLAGS.keep_checkpoint_every_n_hours saver = tf.train.Saver( sharded=True, keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours, save_relative_paths=True) tf.add_to_collection(tf.GraphKeys.SAVERS, saver) scaffold = tf.train.Scaffold(saver=saver) eval_metric_ops = None if mode == tf.estimator.ModeKeys.EVAL: accuracy = tf.metrics.accuracy(labels=labels, predictions=classes) eval_metric_ops = {'Accuracy': accuracy} if mode == tf.estimator.ModeKeys.PREDICT: export_output = exporter._add_output_tensor_nodes(postprocessed_dict) export_outputs = { tf.saved_model.signature_constants.PREDICT_METHOD_NAME: tf.estimator.export.PredictOutput(export_output) } return tf.estimator.EstimatorSpec(mode=mode, predictions=prediction_dict, loss=loss, train_op=train_op, eval_metric_ops=eval_metric_ops, export_outputs=export_outputs, scaffold=scaffold)
def create_model_fn(features, labels, mode, params=None): """Constructs the classification model. Modifed from: https://github.com/tensorflow/models/blob/master/research/ object_detection/model_lib.py. Args: features: A 4-D float32 tensor with shape [batch_size, height, width, channels] representing a batch of images. (Support dict) labels: A 1-D int32 tensor with shape [batch_size] representing the labels of each image. (Support dict) mode: Mode key for tf.estimator.ModeKeys. params: Parameter dictionary passed from the estimator. Returns: An `EstimatorSpec` the encapsulates the model and its serving configurations. """ # 定义网络参数 params = params or {} # 定义网络衡量参数 loss, acc, train_op, export_outputs = None, None, None, None # 根据mode传参进行匹配,mode 指定训练模式,可以取 (TRAIN, EVAL, PREDICT)三者之一 is_training = mode == tf.estimator.ModeKeys.TRAIN # 根据传入的feature读取图片数据 cls_model = model.Model(is_training=is_training, num_classes=FLAGS.num_classes) #预处理,获得输出值,获得预测值 preprocessed_inputs = cls_model.preprocess(features.get('image')) prediction_dict, top_conv, norm_grads_cam = cls_model.predict( preprocessed_inputs) postprocessed_dict = cls_model.postprocess(prediction_dict) # train if mode == tf.estimator.ModeKeys.TRAIN: #进行对应预训练模型的加载 if FLAGS.checkpoint_path: # checkpoint_exclude_scopes = 'resnet_v1_50/conv1,resnet_v1_50/block1' # 指定一些层不加载参数 init_variables_from_checkpoint() if mode in (tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL): loss_dict = cls_model.loss(prediction_dict, labels) loss = loss_dict['loss'] classes = postprocessed_dict['classes'] add_loss = cls_model.add_loss_of_variance(classes, top_conv) add_loss = add_loss * 0.01 loss = tf.add(loss, add_loss) acc = tf.reduce_mean(tf.cast(tf.equal(classes, labels), 'float')) tf.summary.scalar('loss', loss) tf.summary.scalar('accuracy', acc) scaffold = None # train if mode == tf.estimator.ModeKeys.TRAIN: # 设定步数,设定学习率等等超参数 global_step = tf.train.get_or_create_global_step() learning_rate = configure_learning_rate(FLAGS.decay_steps, global_step) optimizer = tf.train.MomentumOptimizer(learning_rate=learning_rate, momentum=0.9) # 冻结层设置,指定一些层不训练 # scopes_to_freeze = 'resnet_v1_50/block1,resnet_v1_50/block2/unit_1' vars_to_train = get_trainable_variables() train_op = slim.learning.create_train_op( loss, optimizer, variables_to_train=vars_to_train, summarize_gradients=True) # 多少时间保存一次模型 keep_checkpoint_every_n_hours = FLAGS.keep_checkpoint_every_n_hours saver = tf.train.Saver( sharded=True, keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours, save_relative_paths=True) tf.add_to_collection(tf.GraphKeys.SAVERS, saver) scaffold = tf.train.Scaffold(saver=saver) eval_metric_ops = None if mode == tf.estimator.ModeKeys.EVAL: accuracy = tf.metrics.accuracy(labels=labels, predictions=classes) eval_metric_ops = {'Eval_Accuracy': accuracy} if mode == tf.estimator.ModeKeys.PREDICT: export_output = exporter._add_output_tensor_nodes(postprocessed_dict) export_outputs = { tf.saved_model.signature_constants.PREDICT_METHOD_NAME: tf.estimator.export.PredictOutput(export_output) } # 返回这个实例化对象 return tf.estimator.EstimatorSpec(mode=mode, predictions=prediction_dict, loss=loss, train_op=train_op, eval_metric_ops=eval_metric_ops, export_outputs=export_outputs, scaffold=scaffold)
def create_model_fn(features, labels, mode, params=None): """Constructs the classification model. Modifed from: https://github.com/tensorflow/models/blob/master/research/ object_detection/model_lib.py. Args: features: A 4-D float32 tensor with shape [batch_size, height, width, channels] representing a batch of images. (Support dict) 可以是一个张量,也可以是由张量组成的一个字典; labels: A 1-D int32 tensor with shape [batch_size] representing the labels of each image. (Support dict) 可以是一个张量,也可以是由张量组成的一个字典; mode: Mode key for tf.estimator.ModeKeys. 指定训练模式,可以取 (TRAIN, EVAL, PREDICT)三者之一; params: Parameter dictionary passed from the estimator. 是一个(可要可不要的)字典,指定其它超参数。 Returns: An `EstimatorSpec` the encapsulates the model and its serving configurations. """ params = params or {} loss, acc, train_op, export_outputs = None, None, None, None is_training = mode == tf.estimator.ModeKeys.TRAIN cls_model = model.Model(is_training=is_training, num_classes=FLAGS.num_classes) preprocessed_inputs = cls_model.preprocess(features.get('image')) prediction_dict = cls_model.predict(preprocessed_inputs) postprocessed_dict = cls_model.postprocess(prediction_dict) if mode == tf.estimator.ModeKeys.TRAIN: if FLAGS.checkpoint_path: init_variables_from_checkpoint() if mode in (tf.estimator.ModeKeys.TRAIN, tf.estimator.ModeKeys.EVAL): loss_dict = cls_model.loss(prediction_dict, labels) loss = loss_dict['loss'] classes = postprocessed_dict['classes'] acc = tf.reduce_mean(tf.cast(tf.equal(classes, labels), 'float')) tf.summary.scalar('loss', loss) tf.summary.scalar('accuracy', acc) scaffold = None if mode == tf.estimator.ModeKeys.TRAIN: global_step = tf.train.get_or_create_global_step() learning_rate = configure_learning_rate(FLAGS.decay_steps, global_step) optimizer = tf.train.MomentumOptimizer(learning_rate=learning_rate, momentum=0.9) train_op = slim.learning.create_train_op(loss, optimizer, summarize_gradients=True) keep_checkpoint_every_n_hours = FLAGS.keep_checkpoint_every_n_hours saver = tf.train.Saver( sharded=True, keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours, save_relative_paths=True) tf.add_to_collection(tf.GraphKeys.SAVERS, saver) scaffold = tf.train.Scaffold(saver=saver) eval_metric_ops = None if mode == tf.estimator.ModeKeys.EVAL: accuracy = tf.metrics.accuracy(labels=labels, predictions=classes) eval_metric_ops = {'Accuracy': accuracy} if mode == tf.estimator.ModeKeys.PREDICT: export_output = exporter._add_output_tensor_nodes(postprocessed_dict) export_outputs = { tf.saved_model.signature_constants.PREDICT_METHOD_NAME: tf.estimator.export.PredictOutput(export_output) } # tf.estimator.EstimatorSpec( # mode, 指定当前是处于训练、验证还是预测状态 # predictions=None, 预测的一个张量,或者是由张量组成的一个字典 # loss=None, 损失张量 # train_op=None, 指定优化操作 # eval_metric_ops=None, 指定各种评估度量的字典,这个字典的值必须是如下两种形式: Metric 类的实例; 调用某个评估度量函数的结果对 (metric_tensor, update_op); # export_outputs=None, 用于模型保存,描述了导出到 SavedModel 的输出格式 # training_chief_hooks=None, # training_hooks=None, # scaffold=None, 一个 tf.train.Scaffold 对象,可以在训练阶段初始化、保存等时使用 # evaluation_hooks=None, # prediction_hooks=None) return tf.estimator.EstimatorSpec(mode=mode, predictions=prediction_dict, loss=loss, train_op=train_op, eval_metric_ops=eval_metric_ops, export_outputs=export_outputs, scaffold=scaffold)