Repository: sql-machine-learning/elasticdl

Issue:
Need to use non-persistent gradient tape
Currently, ElasticDL uses persistent gradient tape in its workers. ElasticDL workers create a persistent gradient tape and assign it to ElasticDL embedding layer before training process. Embedding layers use this gradient tape for the whole training process.

However, in eager mode, persistent gradient tape keeps resource forever and TensorFlow requires user to delete it explicitly. This is because TensorFlow designs persistent gradient tape for the situation that users want to call `tf.gradient` many times, e.g.

```
x = tf.constant(3.0)
with tf.GradientTape(persistent=True) as t:
  t.watch(x)
  y = x * x
  z = y * y
dz_dx = t.gradient(z, x)
dy_dx = t.gradient(y, x)
del t  # Drop the reference to the tape
```

Non-persistent gradient tapes free their resource every time `tape.gradient` is called. So we need to use non-persistent gradient tapes.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/elasticdl/layers/embedding.py b/elasticdl/python/elasticdl/layers/embedding.py
--- a/elasticdl/python/elasticdl/layers/embedding.py
+++ b/elasticdl/python/elasticdl/layers/embedding.py
@@ -161,6 +161,7 @@ def compute_mask(self, inputs, mask=None):
 
     def reset(self):
         self.bet_ids_pair = []
+        self.tape = None
 
     def set_tape(self, tape):
         self.tape = tape
diff --git a/elasticdl/python/worker/worker.py b/elasticdl/python/worker/worker.py
--- a/elasticdl/python/worker/worker.py
+++ b/elasticdl/python/worker/worker.py
@@ -50,13 +50,6 @@ def __init__(
         self._worker_id = worker_id
         self._job_type = job_type
         self._minibatch_size = minibatch_size
-        self._tape = None
-        if (
-            self._job_type == JobType.TRAINING_WITH_EVALUATION
-            or self._job_type == JobType.TRAINING_ONLY
-        ):
-            self._tape = tf.GradientTape(persistent=True)
-
         (
             self._model,
             self._dataset_fn,
@@ -95,7 +88,6 @@ def _init_embedding_layer(self):
         self._embedding_layers = find_layer(self._model, Embedding)
         for layer in self._embedding_layers:
             layer.set_lookup_func(self.lookup_embedding)
-        self._set_tape_for_embedding(self._tape)
         if self._embedding_layers:
             # TODO check that Redis IP/PORT is set
             pass
@@ -316,13 +308,14 @@ def training_process_with_acceleration(self, features, labels):
         return self.training_process_eagerly(features, labels)
 
     def training_process_eagerly(self, features, labels):
-        with self._tape:
+        with tf.GradientTape() as tape:
+            self._set_tape_for_embedding(tape)
             outputs = self._model.call(features, training=True)
             loss = self._loss(outputs, labels)
             # Add regularization loss if any
             if self._model.losses:
                 loss += tf.math.add_n(self._model.losses)
-        grads = self._tape.gradient(loss, self.get_trainable_items())
+        grads = tape.gradient(loss, self.get_trainable_items())
         return loss, grads
 
     @tf.function
@@ -343,14 +336,8 @@ def _run_training_task(self, features, labels):
         return accepted, min_model_version, loss
 
     def _run_evaluation_task(self, features, labels):
-        # reset tape to None for evaluation
-        if self._tape:
-            self._set_tape_for_embedding(None)
         evaluation_metrics = self.evaluation_process(features, labels)
         accepted, _ = self.report_evaluation_metrics(evaluation_metrics)
-        # set tape back for training
-        if self._tape:
-            self._set_tape_for_embedding(self._tape)
         return accepted
 
     def _run_prediction_task(self, features):


Test patch: