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(_): ...@@ -339,8 +339,9 @@ def run_ncf(_):
loss *= (1.0 / (batch_size*strategy.num_replicas_in_sync)) loss *= (1.0 / (batch_size*strategy.num_replicas_in_sync))
grads = tape.gradient(loss, keras_model.trainable_variables) grads = tape.gradient(loss, keras_model.trainable_variables)
optimizer.apply_gradients(list(zip(grads, # Converting gradients to dense form helps in perf on GPU for NCF
keras_model.trainable_variables))) grads = neumf_model.sparse_to_dense_grads(list(zip(grads, keras_model.trainable_variables)))
optimizer.apply_gradients(grads)
return loss return loss
per_replica_losses = strategy.experimental_run(step_fn, per_replica_losses = strategy.experimental_run(step_fn,
......
...@@ -45,7 +45,7 @@ from official.recommendation import stat_utils ...@@ -45,7 +45,7 @@ from official.recommendation import stat_utils
from official.utils.logs import mlperf_helper 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. """Convert sparse gradients to dense gradients.
All sparse gradients, which are represented as instances of tf.IndexedSlices, All sparse gradients, which are represented as instances of tf.IndexedSlices,
...@@ -135,7 +135,7 @@ def neumf_model_fn(features, labels, mode, params): ...@@ -135,7 +135,7 @@ def neumf_model_fn(features, labels, mode, params):
tvars = tf.compat.v1.trainable_variables() tvars = tf.compat.v1.trainable_variables()
gradients = optimizer.compute_gradients( gradients = optimizer.compute_gradients(
loss, tvars, colocate_gradients_with_ops=True) loss, tvars, colocate_gradients_with_ops=True)
gradients = _sparse_to_dense_grads(gradients) gradients = sparse_to_dense_grads(gradients)
minimize_op = optimizer.apply_gradients( minimize_op = optimizer.apply_gradients(
gradients, global_step=global_step, name="train") gradients, global_step=global_step, name="train")
update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS) 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