def model_fn(features, labels, mode, params):
"""
This is a function for creating a computational tensorflow graph.
The function is in format required by tf.estimator.
"""
is_training = mode == tf.estimator.ModeKeys.TRAIN
logits = shufflenet(features['images'],
is_training,
num_classes=params['num_classes'],
depth_multiplier=params['depth_multiplier'])
predictions = {
'probabilities': tf.nn.softmax(logits, axis=1),
'classes': tf.argmax(logits, axis=1, output_type=tf.int32)
}
if mode == tf.estimator.ModeKeys.PREDICT:
export_outputs = tf.estimator.export.PredictOutput({
name: tf.identity(tensor, name)
for name, tensor in predictions.items()
})
return tf.estimator.EstimatorSpec(
mode,
predictions=predictions,
export_outputs={'outputs': export_outputs})
with tf.name_scope('weight_decay'):
add_weight_decay(params['weight_decay'])
regularization_loss = tf.losses.get_regularization_loss()
with tf.name_scope('cross_entropy'):
losses = tf.nn.sparse_softmax_cross_entropy_with_logits(
labels=labels['labels'], logits=logits)
loss = tf.reduce_mean(losses, axis=0)
tf.losses.add_loss(loss)
total_loss = tf.losses.get_total_loss(add_regularization_losses=True)
tf.summary.scalar('cross_entropy_loss', loss)
tf.summary.scalar('regularization_loss', regularization_loss)
if mode == tf.estimator.ModeKeys.EVAL:
eval_metric_ops = {
'accuracy':
tf.metrics.accuracy(labels['labels'], predictions['classes']),
'top5_accuracy':
tf.metrics.mean(
tf.to_float(
tf.nn.in_top_k(predictions=predictions['probabilities'],
targets=labels['labels'],
k=5)))
}
return tf.estimator.EstimatorSpec(mode,
loss=total_loss,
eval_metric_ops=eval_metric_ops)
assert mode == tf.estimator.ModeKeys.TRAIN
with tf.variable_scope('learning_rate'):
global_step = tf.train.get_global_step()
learning_rate = tf.train.polynomial_decay(
params['initial_learning_rate'],
global_step,
params['decay_steps'],
params['end_learning_rate'],
power=1.0) # linear decay
tf.summary.scalar('learning_rate', learning_rate)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops), tf.variable_scope('optimizer'):
optimizer = tf.train.MomentumOptimizer(learning_rate,
momentum=MOMENTUM,
use_nesterov=USE_NESTEROV)
grads_and_vars = optimizer.compute_gradients(total_loss)
train_op = optimizer.apply_gradients(grads_and_vars, global_step)
for g, v in grads_and_vars:
tf.summary.histogram(v.name[:-2] + '_hist', v)
tf.summary.histogram(v.name[:-2] + '_grad_hist', g)
with tf.name_scope('evaluation_ops'):
train_accuracy = tf.reduce_mean(tf.to_float(
tf.equal(labels['labels'], predictions['classes'])),
axis=0)
train_top5_accuracy = tf.reduce_mean(tf.to_float(
tf.nn.in_top_k(predictions=predictions['probabilities'],
targets=labels['labels'],
k=5)),
axis=0)
tf.summary.scalar('train_accuracy', train_accuracy)
tf.summary.scalar('train_top5_accuracy', train_top5_accuracy)
with tf.control_dependencies([train_op]), tf.name_scope('ema'):
ema = tf.train.ExponentialMovingAverage(decay=MOVING_AVERAGE_DECAY,
num_updates=global_step)
train_op = ema.apply(tf.trainable_variables())
return tf.estimator.EstimatorSpec(mode, loss=total_loss, train_op=train_op)
def model_fn(features, labels, mode, params):
"""
This is a function for creating a computational tensorflow graph.
The function is in format required by tf.estimator.
"""
global INIT_FLAG
is_training = mode == tf.estimator.ModeKeys.TRAIN
logits = shufflenet(
features['images'], is_training,
num_classes=params['num_classes'],
depth_multiplier=params['depth_multiplier']
)
predictions = {
'probabilities': tf.nn.softmax(logits, axis=1),
'classes': tf.argmax(logits, axis=1, output_type=tf.int32)
}
if mode == tf.estimator.ModeKeys.PREDICT:
export_outputs = tf.estimator.export.PredictOutput({
name: tf.identity(tensor, name)
for name, tensor in predictions.items()
})
return tf.estimator.EstimatorSpec(
mode, predictions=predictions,
export_outputs={'outputs': export_outputs}
)
init_fn=None
if INIT_FLAG:
exclude =['global_step:0']#,'classifier/biases:0','classifier/weights:0'] # ['ShuffleNetV2/Conv1/weights:0','global_step:0','classifier/biases:0','classifier/weights:0'
#]
all_variables = tf.contrib.slim.get_variables_to_restore()
varialbes_to_use=[]
num_params=0
for v in all_variables:
shape=v.get_shape()
num_params+= reduce(mul, [dim.value for dim in shape], 1)
if v.name not in exclude:
varialbes_to_use.append(v)
init_fn = tf.contrib.framework.assign_from_checkpoint_fn(tf.train.latest_checkpoint('models/jiu-huan_0.33_9967'), varialbes_to_use, ignore_missing_vars=True)
print('***********************params: ', num_params)
with tf.name_scope('weight_decay'):
add_weight_decay(params['weight_decay'])
regularization_loss = tf.losses.get_regularization_loss()
with tf.name_scope('cross_entropy'):
losses = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=labels['labels'], logits=logits)
loss = tf.reduce_mean(losses, axis=0)
tf.losses.add_loss(loss)
total_loss = tf.losses.get_total_loss(add_regularization_losses=True)
tf.summary.scalar('cross_entropy_loss', loss)
tf.summary.scalar('regularization_loss', regularization_loss)
if mode == tf.estimator.ModeKeys.EVAL:
eval_metric_ops = {
'accuracy': tf.metrics.accuracy(labels['labels'], predictions['classes']),
'top5_accuracy': tf.metrics.mean(tf.to_float(tf.nn.in_top_k(
predictions=predictions['probabilities'], targets=labels['labels'], k=5
)))
}
return tf.estimator.EstimatorSpec(mode, loss=total_loss, eval_metric_ops=eval_metric_ops)
assert mode == tf.estimator.ModeKeys.TRAIN
with tf.variable_scope('learning_rate'):
global_step = tf.train.get_global_step()
learning_rate = tf.train.polynomial_decay(
params['initial_learning_rate'], global_step,
params['decay_steps'], params['end_learning_rate'],
power=1.0
) # linear decay
tf.summary.scalar('learning_rate', learning_rate)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
with tf.control_dependencies(update_ops), tf.variable_scope('optimizer'):
optimizer = tf.train.MomentumOptimizer(learning_rate, momentum=MOMENTUM, use_nesterov=USE_NESTEROV)
grads_and_vars = optimizer.compute_gradients(total_loss)
train_op = optimizer.apply_gradients(grads_and_vars, global_step)
for g, v in grads_and_vars:
tf.summary.histogram(v.name[:-2] + '_hist', v)
tf.summary.histogram(v.name[:-2] + '_grad_hist', g)
with tf.name_scope('evaluation_ops'):
train_accuracy = tf.reduce_mean(tf.to_float(tf.equal(
labels['labels'], predictions['classes']
)), axis=0)
train_top5_accuracy = tf.reduce_mean(tf.to_float(tf.nn.in_top_k(
predictions=predictions['probabilities'], targets=labels['labels'], k=5
)), axis=0)
tf.summary.scalar('train_accuracy', train_accuracy)
tf.summary.scalar('train_top5_accuracy', train_top5_accuracy)
with tf.control_dependencies([train_op]), tf.name_scope('ema'):
ema = tf.train.ExponentialMovingAverage(decay=MOVING_AVERAGE_DECAY, num_updates=0)
train_op = ema.apply(tf.trainable_variables())
if INIT_FLAG:
INIT_FLAG=False
return tf.estimator.EstimatorSpec(mode, loss=total_loss, train_op=train_op, training_hooks=[RestoreHook(init_fn)])
else:
return tf.estimator.EstimatorSpec(mode, loss=total_loss, train_op=train_op)
image = cv2.resize(image, (224, 224), cv2.INTER_LINEAR)
Image.fromarray(image)
# # Load a trained model
# In[ ]:
tf.reset_default_graph()
raw_images = tf.placeholder(tf.uint8, [None, 224, 224, 3])
images = tf.to_float(raw_images)/255.0
logits = shufflenet(images, is_training=False, depth_multiplier='0.5')
probabilities = tf.nn.softmax(logits, axis=1)
ema = tf.train.ExponentialMovingAverage(decay=0.995)
variables_to_restore = ema.variables_to_restore()
saver = tf.train.Saver(variables_to_restore)
# # Predict
# In[ ]:
with tf.Session() as sess:
saver.restore(sess, 'run00/model.ckpt-1331064')
feed_dict = {raw_images: np.expand_dims(image, axis=0)}