heavyball 0.24.1__tar.gz → 0.24.2__tar.gz

{heavyball-0.24.1 → heavyball-0.24.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 0.24.1
+Version: 0.24.2
 Summary: Efficient optimizers
 Home-page: https://github.com/clashluke/heavyball
 Author: Lucas Nestler

{heavyball-0.24.1 → heavyball-0.24.2}/heavyball/utils.py

@@ -1129,3 +1129,12 @@ def merge_group(group, *tensors):
         append_or_extend(out, dim_merger(t, group['max_size_triangular'] if 'max_size_triangular' in group else group[
             'max_precond_dim'], group.get('split', False)))
     return out
+
+def hook_optimizer_into_model(model, optimizer, *args, **kwargs):
+    def _step(p: Tensor, o: torch.optim.Optimizer):
+        o.step()
+        o.zero_grad()
+
+
+    for p in model.parameters():
+        p.register_post_accumulate_grad_hook(functools.partial(_step, o=optimizer([p], *args, **kwargs)))

{heavyball-0.24.1 → heavyball-0.24.2}/heavyball.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: heavyball
-Version: 0.24.1
+Version: 0.24.2
 Summary: Efficient optimizers
 Home-page: https://github.com/clashluke/heavyball
 Author: Lucas Nestler

{heavyball-0.24.1 → heavyball-0.24.2}/heavyball.egg-info/SOURCES.txt

@@ -33,6 +33,7 @@ test/test_channels_last.py
 test/test_closure.py
 test/test_ema.py
 test/test_foreach.py
+test/test_hook.py
 test/test_mars.py
 test/test_memory.py
 test/test_merge.py

{heavyball-0.24.1 → heavyball-0.24.2}/setup.py

@@ -10,7 +10,7 @@ setuptools.setup(
     name='heavyball',
     license='BSD',
     description='Efficient optimizers',
-    version='0.24.1',
+    version='0.24.2',
     long_description=README,
     url='https://github.com/clashluke/heavyball',
     packages=setuptools.find_packages(),

heavyball-0.24.2/test/test_hook.py

@@ -0,0 +1,51 @@
+import os
+
+os.environ["TORCH_LOGS"] = "+recompiles"
+
+import heavyball
+import heavyball.utils
+import pytest
+import torch
+from benchmark.utils import get_optim
+from heavyball.utils import clean, set_torch, hook_optimizer_into_model
+from torch import nn
+from torch._dynamo import config
+
+heavyball.utils.compile_mode = 'default'
+config.cache_size_limit = 128
+
+
+@pytest.mark.parametrize("opt", heavyball.__all__)
+@pytest.mark.parametrize("size,depth", [(128, 1)])
+def test_foreach(opt, size, depth: int, iterations: int = 128, outer_iterations: int = 1):
+    set_torch()
+    opt = getattr(heavyball, opt)
+
+    peaks = []
+    losses = []
+
+    for use_hook in [False, True]:
+        torch.manual_seed(0x2131290)
+        peaks.append([])
+        losses.append([])
+
+        for i in range(outer_iterations):
+            model = nn.Sequential(*[nn.Linear(size, size) for _ in range(depth)]).cuda()
+
+            if use_hook:
+                hook_optimizer_into_model(model, opt, lr=1e-3, weight_decay=1e-4, warmup_steps=16)
+            else:
+                o = get_optim(opt, model.parameters(), lr=1e-3, weight_decay=1e-4, warmup_steps=16)
+            for _ in range(iterations):
+                loss = model(torch.randn((1024, size), device='cuda')).square().mean()
+                loss.backward()
+                if not use_hook:
+                    o.step()
+                    o.zero_grad()
+                losses[-1].append(loss.detach())
+
+            clean()
+
+    for i, (l0, l1) in enumerate(zip(*losses)):
+        print(i, l0.item(), l1.item())
+        assert torch.allclose(l0.float(), l1.float(), rtol=0.1)