Skip to content

Commit 0468a8d

Browse files
Jonathan Shencopybara-github
authored andcommitted
Default to wrapping optimizer.Apply in learner variable_scope.
optimizer.Apply can create variables. The variables created by different optimizers can conflict with the same name in eager mode (there is no variable name auto-uniquification in eager mode). Thus manually wrapping them in a variable scope is necessary. Also delete an unused function. PiperOrigin-RevId: 415344496
1 parent 9ddf63a commit 0468a8d

File tree

1 file changed

+4
-21
lines changed

1 file changed

+4
-21
lines changed

lingvo/core/learner.py

Lines changed: 4 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -224,30 +224,13 @@ def Apply(self, metrics, vmap, gradient_mask=None, gradient_adjuster=None):
224224
else:
225225
lr_or_callable = lr
226226

227-
var_update_op = tf.group(
228-
[tpu_emb_update_op,
229-
self.optimizer.Apply(lr_or_callable, var_grads)])
227+
with self._SelfVariableScope():
228+
var_update_op = tf.group(
229+
[tpu_emb_update_op,
230+
self.optimizer.Apply(lr_or_callable, var_grads)])
230231

231232
return losses, var_update_op, eval_metrics
232233

233-
def ComputeActivationGradients(self, activations, activations_grad, vmap):
234-
p = self.params
235-
vmap = self.GetTrainableVariables(vmap)
236-
237-
for v in vmap.Flatten():
238-
tf.logging.info('%s: bprop variable: %s', p.name, v.name)
239-
return self.optimizer.ComputeGradients(
240-
activations,
241-
vmap,
242-
p.grad_aggregation_method,
243-
p.colocate_gradients_with_ops,
244-
p.gate_gradients,
245-
compute_gradients_fn=self._CustomComputeGradientsFn(),
246-
skip_zero_gradients=p.skip_zero_gradients,
247-
skip_none_gradients=False,
248-
activations_grad=activations_grad,
249-
is_activations=True)
250-
251234
def ComputeLosses(self, metrics):
252235
p = self.params
253236

0 commit comments

Comments
 (0)