add use of exponentially moving average
@@ -40,7 +40,8 @@ _logger.setLevel(logging.INFO)
def train(sess, loss, x_train, y_train,
init_all=True, evaluate=None, feed=None, args=None,
rng=None, var_list=None, fprop_args=None, optimizer=None,
- devices=None, x_batch_preprocessor=None):
+ devices=None, x_batch_preprocessor=None, use_ema=False,
+ ema_decay=.998):
"""
Run (optionally multi-replica, synchronous) training to minimize `loss`
:param sess: TF session to use when training the graph
@@ -69,6 +70,10 @@ def train(sess, loss, x_train, y_train,
Takes a single tensor containing an x_train batch as input
Returns a single tensor containing an x_train batch as output
Called to preprocess the data before passing the data to the Loss
+ :param use_ema: bool
+ If true, uses an exponential moving average of the model parameters
+ :param ema_decay: float
+ The decay parameter for EMA, if EMA is used
:return: True if model trained
"""
args = _ArgsWrapper(args or {})
@@ -124,6 +129,21 @@ def train(sess, loss, x_train, y_train,
with tf.control_dependencies(tf.get_collection(tf.GraphKeys.UPDATE_OPS)):
train_step = optimizer.apply_gradients(grad)
+
+ if use_ema:
+ ema = tf.train.ExponentialMovingAverage(decay=ema_decay)
+ with tf.control_dependencies([train_step]):
+ train_step = ema.apply(var_list)
+ avg_params = [ema.average(param) for param in var_list]
+ tmp_params = [tf.Variable(param, trainable=False) for param in var_list]
+ param_to_tmp = [tf.assign(tmp, param) for tmp, param in safe_zip(tmp_params, var_list)]
+ with tf.control_dependencies(param_to_tmp):
+ avg_to_param = [tf.assign(param, avg) for param, avg in safe_zip(var_list, avg_params)]
+ with tf.control_dependencies(param_to_tmp):
+ tmp_to_avg = [tf.assign(avg, tmp) for avg, tmp in safe_zip(avg_params, tmp_params)]
+ swap = tmp_to_avg
+
+
batch_size = args.batch_size
assert batch_size % num_devices == 0
@@ -197,8 +217,6 @@ def train(sess, loss, x_train, y_train,
# Compute batch start and end indices
start = batch * batch_size
end = (batch + 1) * batch_size
- # start, end = batch_indices(
- # batch, len(x_train), args.batch_size)
# Perform one training step
feed_dict = {}
@@ -227,7 +245,13 @@ def train(sess, loss, x_train, y_train,
_logger.info("Epoch " + str(epoch) + " took " +
str(cur - prev) + " seconds")
if evaluate is not None:
+ if use_ema:
+ sess.run(swap)
evaluate()
+ if use_ema:
+ sess.run(swap)
+ if use_ema:
+ sess.run(swap)
return True
