Add support for custom optimizer for model_train in utils_tf
@@ -66,7 +66,8 @@ def initialize_uninitialized_global_variables(sess):
def model_train(sess, x, y, predictions, X_train, Y_train, save=False,
predictions_adv=None, init_all=True, evaluate=None,
- feed=None, args=None, rng=None, var_list=None):
+ feed=None, args=None, rng=None, var_list=None,
+ optimizer=None):
"""
Train a TF graph
:param sess: TF session to use when training the graph
@@ -93,6 +94,7 @@ def model_train(sess, x, y, predictions, X_train, Y_train, save=False,
and 'filename'
:param rng: Instance of numpy.random.RandomState
:param var_list: Optional list of parameters to train.
+ :param optimizer: Optimizer to be used for training
:return: True if model trained
"""
args = _ArgsWrapper(args or {})
@@ -109,13 +111,15 @@ def model_train(sess, x, y, predictions, X_train, Y_train, save=False,
if rng is None:
rng = np.random.RandomState()
+ if optimizer is None:
+ optimizer = tf.train.AdamOptimizer(learning_rate=args.learning_rate)
+
# Define loss
loss = model_loss(y, predictions)
if predictions_adv is not None:
loss = (loss + model_loss(y, predictions_adv)) / 2
- train_step = tf.train.AdamOptimizer(learning_rate=args.learning_rate)
- train_step = train_step.minimize(loss, var_list=var_list)
+ train_step = optimizer.minimize(loss, var_list=var_list)
with sess.as_default():
if hasattr(tf, "global_variables_initializer"):
