Skip to content

chore(optim): wrap torch.autograd.grad() with torch.enable_grad() context #220

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.

Sign up for GitHub

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

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

chore(optim): wrap torch.autograd.grad() with torch.enable_grad() context #220

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
6c5304a
chore(optim): wrap `torch.autograd.grad()` with `torch.enable_grad()`…
XuehaiPan May 21, 2024
5f9b918
fix(optim): move `apply_updates` in grad context
XuehaiPan May 21, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 10 additions & 9 deletions torchopt/optim/func/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,15 +87,16 @@ def step(
if inplace is None:
inplace = self.inplace

# Step parameter only
grads = torch.autograd.grad(loss, params, create_graph=True, allow_unused=True)
updates, self.optim_state = self.impl.update(
grads,
self.optim_state,
params=params,
inplace=inplace,
)
return apply_updates(params, updates, inplace=inplace)
with torch.enable_grad():
# Step parameters only
grads = torch.autograd.grad(loss, params, create_graph=True, allow_unused=True)
updates, self.optim_state = self.impl.update(
grads,
self.optim_state,
params=params,
inplace=inplace,
)
return apply_updates(params, updates, inplace=inplace)

def state_dict(self) -> OptState:
"""Extract the references of the optimizer states.
Expand Down
30 changes: 15 additions & 15 deletions torchopt/optim/meta/base.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,32 +66,32 @@ def step(self, loss: torch.Tensor) -> None: # pylint: disable=too-many-locals
loss (torch.Tensor): The loss that is used to compute the gradients to the network
parameters.
"""
# Step parameter only
for i, (param_container, state) in enumerate(
zip(self.param_containers_groups, self.state_groups),
):
flat_params: TupleOfTensors
flat_params, container_treespec = pytree.tree_flatten_as_tuple(param_container) # type: ignore[arg-type]

if isinstance(state, UninitializedState):
state = self.impl.init(flat_params)
grads = torch.autograd.grad(
loss,
flat_params,
create_graph=True,
allow_unused=True,
)
updates, new_state = self.impl.update(
grads,
state,
params=flat_params,
inplace=False,
)
self.state_groups[i] = new_state
flat_new_params = apply_updates(flat_params, updates, inplace=False)

with torch.enable_grad():
# Step parameters only
grads = torch.autograd.grad(loss, flat_params, create_graph=True, allow_unused=True)
updates, new_state = self.impl.update(
grads,
state,
params=flat_params,
inplace=False,
)
flat_new_params = apply_updates(flat_params, updates, inplace=False)

new_params: ModuleTensorContainers = pytree.tree_unflatten( # type: ignore[assignment]
container_treespec,
flat_new_params,
)

self.state_groups[i] = new_state
for container, new_param in zip(param_container, new_params):
container.update(new_param)

Expand Down