Unverified Commit 44ff121d authored by nnigania's avatar nnigania Committed by GitHub
Browse files

NCF CTL Perf optimization to convert gradients from sparse to dense (#7102)

* borrowing a tf1.x optimization which converts gradients from sparse to dense for better perf

* cleanup after code review
parent 5afa9569
......@@ -339,8 +339,9 @@ def run_ncf(_):
loss *= (1.0 / (batch_size*strategy.num_replicas_in_sync))
grads = tape.gradient(loss, keras_model.trainable_variables)
optimizer.apply_gradients(list(zip(grads,
keras_model.trainable_variables)))
# Converting gradients to dense form helps in perf on GPU for NCF
grads = neumf_model.sparse_to_dense_grads(list(zip(grads, keras_model.trainable_variables)))
optimizer.apply_gradients(grads)
return loss
per_replica_losses = strategy.experimental_run(step_fn,
......
......@@ -45,7 +45,7 @@ from official.recommendation import stat_utils
from official.utils.logs import mlperf_helper
def _sparse_to_dense_grads(grads_and_vars):
def sparse_to_dense_grads(grads_and_vars):
"""Convert sparse gradients to dense gradients.
All sparse gradients, which are represented as instances of tf.IndexedSlices,
......@@ -135,7 +135,7 @@ def neumf_model_fn(features, labels, mode, params):
tvars = tf.compat.v1.trainable_variables()
gradients = optimizer.compute_gradients(
loss, tvars, colocate_gradients_with_ops=True)
gradients = _sparse_to_dense_grads(gradients)
gradients = sparse_to_dense_grads(gradients)
minimize_op = optimizer.apply_gradients(
gradients, global_step=global_step, name="train")
update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment