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)
