Modifying Adam to support complex numbers as 2d real numbers #80279

zaxtax · 2022-06-25T11:09:53Z

This commit addresses issues in #65711

facebook-github-bot · 2022-06-25T11:09:59Z

🔗 Helpful links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/80279
📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓Need help or want to give feedback on the CI? Visit our office hours

✅ No Failures (0 Pending)

As of commit 3e6f0a0 (more details on the Dr. CI page):

Expand to see more

💚 💚 Looks good so far! There are no failures yet. 💚 💚

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

albanD

Code change looks ok but CI failures are real and need to be fixed.

albanD · 2022-07-21T09:23:03Z

torch/optim/adam.py

@@ -260,6 +260,12 @@ def _single_tensor_adam(params: List[Tensor],
            grad = grad.add(param, alpha=weight_decay)

        # Decay the first and second moment running average coefficient
+        if torch.is_complex(param):


The comment above needs to be moved down

albanD · 2022-07-21T09:23:18Z

torch/optim/adam.py

@@ -307,6 +313,13 @@ def _single_tensor_adam(params: List[Tensor],

            param.addcdiv_(exp_avg, denom, value=-step_size)

+        if torch.is_complex(param):


This doesn't do anything right?

Oops. I see my error.

albanD · 2022-07-21T09:24:43Z

test/test_optim.py

+        a1_imag = a1.imag.clone().detach()
+        a1_real.requires_grad_()
+        a1_imag.requires_grad_()
+        a2 = torch.complex(a1_real, a1_imag)


This a2 is never used?

test/test_optim.py

albanD · 2022-07-21T09:43:57Z

test/test_optim.py

+            f(a1).backward()
+            f(a2).backward()
+
+            assert(torch.allclose(a1.grad.real, a1_real.grad))


Please use the builtin asserts on self. to play nicely with the test suite in general.

zaxtax · 2022-07-25T18:04:55Z

@albanD does this look alright?

albanD · 2022-07-25T19:06:19Z

test/test_optim.py

@@ -320,6 +320,34 @@ def _test_complex_optimizer(self, optimizer_constructor):

            self.assertEqual(torch.view_as_real(complex_param), real_param)

+    def _test_complex_2d(self, optimizer_constructor, f=None):


The f argument is not actually needed right?

I want to reuse this function for all tests involving the complex optimisers.

albanD · 2022-07-25T19:06:46Z

test/test_optim.py

+        a1 = torch.randn(2, dtype=torch.complex64, requires_grad=True)
+        a1_real = a1.real.clone().detach()
+        a1_imag = a1.imag.clone().detach()
+        a1_real.requires_grad_()


nit: this can be inlined above if you want.

albanD · 2022-07-25T19:07:27Z

test/test_optim.py

+            f(a1).backward()
+            f(a2).backward()
+
+            self.assertTrue(torch.allclose(a1.grad.real, a1_real.grad))


Isn't self.assertEqual() working here? It is doing a close check by default IIRC

albanD · 2022-07-25T19:07:45Z

torch/optim/adam.py

@@ -307,6 +314,13 @@ def _single_tensor_adam(params: List[Tensor],

            param.addcdiv_(exp_avg, denom, value=-step_size)

+        if is_complex_param:
+            grad = torch.view_as_complex(grad)


This is still not needed right?

albanD · 2022-07-25T19:07:55Z

torch/optim/adam.py

@@ -404,4 +424,5 @@ def _multi_tensor_adam(params: List[Tensor],
            torch._foreach_div_(exp_avg_sq_sqrt, bias_correction2_sqrt)
            denom = torch._foreach_add(exp_avg_sq_sqrt, eps)

-        torch._foreach_addcdiv_(params, exp_avgs, denom, step_size)
+        torch._foreach_addcdiv_(params_, exp_avgs, denom, step_size)
+    params = [torch.view_as_complex(x) if torch.is_complex(params[i]) else x for i, x in enumerate(params_)]


Not needed either.

params is the original params_ has all complex tensors converted to reals.

When is params actually modified inplace?

params is never updated inplace. params_ replaces it wholesale

So if it is never updated inplace, there is no need to restore it here and this line does nothing?

I might be misunderstanding something, but doesn't params need to hold the updated values. I do all the computation in params_ but that is only locally defined. For the update step to change the parameters, I assumed I have to specifically change params. Certainly, if I don't the unit test fails for me.

I've pushed changes that remove this line since I think I misunderstood the semantics of view_as_real and view_as_complex

Yes, the content of these lists are modified inplace!

albanD

Change looks good.
Small nit about test organization but good otherwise.

albanD · 2022-07-27T14:56:39Z

test/test_optim.py

@@ -566,27 +592,14 @@ def test_adam(self):
                 lambda opt: ReduceLROnPlateau(opt)],
                constructor_accepts_maximize=True
            )
+            self._test_complex_2d(optimizer)


I'm a bit surprised we're rolling out a custom test here. Why can't _test_complex_optimizer() be re-used?

They don't test quite the same things, but I can refactor these two functions into a single one. Is this a blocker for the PR being merged?

zaxtax · 2022-07-27T18:38:17Z

@pytorchbot merge

pytorchmergebot · 2022-07-27T18:39:34Z

@pytorchbot successfully started a merge job. Check the current status here

github-actions · 2022-07-27T18:40:13Z

Hey @zaxtax.
You've committed this PR, but it does not have both a 'release notes: ...' and 'topics: ...' label. Please add one of each to the PR. The 'release notes: ...' label should represent the part of PyTorch that this PR changes (fx, autograd, distributed, etc) and the 'topics: ...' label should represent the kind of PR it is (not user facing, new feature, bug fix, perf improvement, etc). The list of valid labels can be found here for the 'release notes: ...' and here for the 'topics: ...'.
For changes that are 'topic: not user facing' there is no need for a release notes label.

…#80279) Summary: This commit addresses issues in #65711 Pull Request resolved: #80279 Approved by: https://github.com/albanD Test Plan: contbuild & OSS CI, see https://hud.pytorch.org/commit/pytorch/pytorch/f9ef363982136f45dfb2bd4205c545cb17e59afd Reviewed By: osalpekar Differential Revision: D38227584 Pulled By: osalpekar fbshipit-source-id: 48fbb9187124fe7d337e464f41f34ccc0d8b927b

zaxtax requested a review from albanD as a code owner June 25, 2022 11:09

facebook-github-bot added the cla signed label Jun 25, 2022

pytorchbot added the open source label Jun 25, 2022

mruberry added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Jun 27, 2022

zaxtax mentioned this pull request Jun 27, 2022

Handle complex optimization in AdamW by treating complex numbers as 2D real numbers #80280

Closed

albanD reviewed Jun 29, 2022

View reviewed changes

zaxtax force-pushed the adam-2d-complex branch 2 times, most recently from 63dbc12 to 92b5082 Compare June 29, 2022 14:51

zaxtax force-pushed the adam-2d-complex branch from 92b5082 to 63ab89c Compare July 21, 2022 04:15

albanD reviewed Jul 21, 2022

View reviewed changes

zaxtax force-pushed the adam-2d-complex branch 4 times, most recently from b7cd8cf to 019ff97 Compare July 25, 2022 00:49

albanD reviewed Jul 25, 2022

View reviewed changes

Modifying Adam to support complex numbers as 2d real numbers

3e6f0a0

zaxtax force-pushed the adam-2d-complex branch from 019ff97 to 3e6f0a0 Compare July 27, 2022 14:36

albanD approved these changes Jul 27, 2022

View reviewed changes

pytorchmergebot added the Merged label Jul 27, 2022

pytorchmergebot closed this in f9ef363 Jul 27, 2022

		@@ -307,6 +313,13 @@ def _single_tensor_adam(params: List[Tensor],

		param.addcdiv_(exp_avg, denom, value=-step_size)

		if torch.is_complex(param):

		@@ -320,6 +320,34 @@ def _test_complex_optimizer(self, optimizer_constructor):

		self.assertEqual(torch.view_as_real(complex_param), real_param)

		def _test_complex_2d(self, optimizer_constructor, f=None):

Modifying Adam to support complex numbers as 2d real numbers #80279

Modifying Adam to support complex numbers as 2d real numbers #80279

Uh oh!

Conversation

zaxtax commented Jun 25, 2022

Uh oh!

facebook-github-bot commented Jun 25, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful links

✅ No Failures (0 Pending)

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zaxtax commented Jul 25, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zaxtax Jul 27, 2022 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zaxtax commented Jul 27, 2022

Uh oh!

pytorchmergebot commented Jul 27, 2022

Uh oh!

github-actions bot commented Jul 27, 2022

Uh oh!

Uh oh!

facebook-github-bot commented Jun 25, 2022 •

edited

Loading

zaxtax Jul 27, 2022 •

edited

Loading