def _model_fn(features, labels, mode, params):
model_fn = MODELS[FLAGS.model].model
global_step = tf.train.get_or_create_global_step()
if FLAGS.num_gpus > 0 and mode == learn.ModeKeys.TRAIN:
split_features = {k: tf.split(v, FLAGS.num_gpus)
for k, v in features.iteritems()}
split_labels = {k: tf.split(v, FLAGS.num_gpus)
for k, v in labels.iteritems()}
grads = []
predictions = collections.defaultdict(list)
losses = []
opt = ops.create_optimizer(
params.optimizer, params.learning_rate, params.decay_steps)
for i in range(FLAGS.num_gpus):
with tf.device(tf.DeviceSpec(device_type="GPU", device_index=i)):
with tf.name_scope("tower_%d" % i):
with tf.variable_scope(tf.get_variable_scope(), reuse=i > 0):
device_features = {k: v[i] for k, v in split_features.iteritems()}
device_labels = {k: v[i] for k, v in split_labels.iteritems()}
device_predictions, device_loss = model_fn(
device_features, device_labels, mode, params)
for k, v in device_predictions.iteritems():
predictions[k].append(v)
if device_loss is not None:
losses.append(device_loss)
device_grads = opt.compute_gradients(device_loss)
grads.append(device_grads)
grads = ops.average_gradients(grads)
train_op = opt.apply_gradients(grads, global_step=global_step)
for k, v in predictions.iteritems():
predictions[k] = tf.concat(v, axis=0)
loss = tf.add_n(losses) if losses else None
else:
with tf.device(tf.DeviceSpec(device_type="GPU", device_index=0)):
predictions, loss = model_fn(features, labels, mode, params)
train_op = None
if mode == learn.ModeKeys.TRAIN:
opt = ops.create_optimizer(
params.optimizer, params.learning_rate, params.decay_steps)
train_op = opt.minimize(loss, global_step=global_step)
tf.summary.scalar("loss/loss", loss)
return tf.contrib.learn.ModelFnOps(
mode=mode,
predictions=predictions,
loss=loss,
train_op=train_op)
def _model_fn(features, labels, mode, params):
model_fn = MODELS[FLAGS.model].model_fn
global_step = tf.train.get_or_create_global_step()
if FLAGS.num_gpus > 0 and mode == learn.ModeKeys.TRAIN:
split_features = {k: tf.split(v, FLAGS.num_gpus)
for k, v in features.iteritems()}
split_labels = {k: tf.split(v, FLAGS.num_gpus)
for k, v in labels.iteritems()}
grads = []
predictions = collections.defaultdict(list)
losses = []
opt = ops.create_optimizer(
params.optimizer, params.learning_rate, params.decay_steps)
for i in range(FLAGS.num_gpus):
with tf.device(tf.DeviceSpec(device_type='GPU', device_index=i)):
with tf.name_scope('tower_%d' % i):
with tf.variable_scope(tf.get_variable_scope(), reuse=i > 0):
device_features = {k: v[i] for k, v in split_features.iteritems()}
device_labels = {k: v[i] for k, v in split_labels.iteritems()}
device_predictions, device_loss = model_fn(
device_features, device_labels, mode, params)
for k, v in device_predictions.iteritems():
predictions[k].append(v)
if device_loss is not None:
losses.append(device_loss)
device_grads = opt.compute_gradients(device_loss)
grads.append(device_grads)
grads = ops.average_gradients(grads)
train_op = opt.apply_gradients(grads, global_step=global_step)
for k, v in predictions.iteritems():
predictions[k] = tf.concat(v, axis=0)
loss = tf.add_n(losses) if losses else None
else:
with tf.device(tf.DeviceSpec(device_type='GPU', device_index=0)):
predictions, loss = model_fn(features, labels, mode, params)
train_op = None
if mode == learn.ModeKeys.TRAIN:
opt = ops.create_optimizer(
params.optimizer, params.learning_rate, params.decay_steps)
train_op = opt.minimize(loss, global_step=global_step)
tf.summary.scalar('loss/loss', loss)
return tf.contrib.learn.ModelFnOps(
mode=mode,
predictions=predictions,
loss=loss,
train_op=train_op)
def _model_fn(features, labels, mode, params):
model_fn = MODELS[FLAGS.model].model
global_step = tf.train.get_or_create_global_step()
opt = ops.create_optimizer(params.optimizer, params.learning_rate,
params.momentum, params.warmup_steps,
params.decay_steps, params.decay_rate)
if FLAGS.num_gpus:
split_features = {
k: tf.split(v, FLAGS.num_gpus)
for k, v in features.iteritems()
}
split_labels = {
k: tf.split(v, FLAGS.num_gpus)
for k, v in labels.iteritems()
}
grads = []
predictions = collections.defaultdict(list)
losses = []
for i in range(FLAGS.num_gpus):
with tf.device(tf.DeviceSpec(device_type="GPU",
device_index=i)):
with tf.name_scope("tower_%d" % i) as name_scope:
with tf.variable_scope(tf.get_variable_scope(),
reuse=i > 0):
device_features = {
k: v[i]
for k, v in split_features.iteritems()
}
device_labels = {
k: v[i]
for k, v in split_labels.iteritems()
}
device_predictions, device_loss, device_metrics = model_fn(
device_features, device_labels, mode, params)
tf.summary.scalar("loss/main", device_loss)
if i == 0:
metrics = device_metrics
update_ops = tf.get_collection(
tf.GraphKeys.UPDATE_OPS, name_scope)
reg_losses = tf.get_collection(
tf.GraphKeys.REGULARIZATION_LOSSES,
name_scope)
tf.summary.scalar("loss/regularization",
tf.add_n(reg_losses))
device_loss = tf.add_n([device_loss] +
reg_losses)
for k, v in device_predictions.iteritems():
predictions[k].append(v)
if device_loss is not None:
losses.append(device_loss)
if mode == tf.estimator.ModeKeys.TRAIN:
with tf.control_dependencies(update_ops):
device_grads = opt.compute_gradients(
device_loss)
grads.append(device_grads)
with tf.device(tf.DeviceSpec(device_type="GPU", device_index=0)):
if mode == tf.estimator.ModeKeys.TRAIN:
grads = ops.average_gradients(grads)
train_op = opt.apply_gradients(grads,
global_step=global_step)
else:
train_op = None
for k, v in predictions.iteritems():
predictions[k] = tf.concat(v, axis=0)
loss = tf.add_n(losses) if losses else None
else:
predictions, loss, metrics = model_fn(features, labels, mode,
params)
tf.summary.scalar("loss/main", loss)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
reg_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
tf.summary.scalar("loss/regularization", tf.add_n(reg_losses))
loss = tf.add_n([loss] + reg_losses)
if mode == tf.estimator.ModeKeys.TRAIN:
with tf.control_dependencies(update_ops):
train_op = opt.minimize(loss, global_step=global_step)
else:
train_op = None
if mode == tf.estimator.ModeKeys.TRAIN:
opts = tf.profiler.ProfileOptionBuilder(
).trainable_variables_parameter()
stats = tf.profiler.profile(tf.get_default_graph(), options=opts)
print("Total parameters:", stats.total_parameters)
return tf.estimator.EstimatorSpec(mode=mode,
predictions=predictions,
loss=loss,
train_op=train_op,
eval_metric_ops=metrics)