Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Optimizer is not respecting "trainable" attribute of variables. #14

Open
rubenverhack opened this issue Jan 25, 2021 · 0 comments
Open

Comments

@rubenverhack
Copy link

The current version does not respect untrainable variables. It can be fixed by placing a simple if-statement. However, I'm not sure if this is the best place. Therefore I'm not suggesting it as a PR, but report the issue here.

        eagerly_outside_functions = ops.executing_eagerly_outside_functions()
        update_ops = []
        with ops.name_scope(name or self._name, skip_on_eager=True):
            for grad, hess, var in grads_hessian_and_vars:

                # FIX UNTRAINABLE
                if var.trainable:
                    def _assume_mirrored(grad, hess):
                        if isinstance(grad, ds_values.PerReplica):
                            return ds_values.Mirrored(grad.values), ds_values.Mirrored(hess.values)
                        return grad, hess

                    grad, hess = nest.map_structure(_assume_mirrored, grad, hess)
                    # Colocate the update with variables to avoid unnecessary communication
                    # delays. See b/136304694.
                    with distribution.extended.colocate_vars_with(var):
                        with ops.name_scope("update" if eagerly_outside_functions else
                                    "update_" + var.op.name, skip_on_eager=True):
                            update_ops.extend(distribution.extended.update(
                                    var, apply_grad_to_update_var, args=(grad, hess), group=False))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant