apply lr schedule to hypernets

1 files changed, 15 insertions, 4 deletions

diff --git a/modules/hypernetworks/hypernetwork.py b/modules/hypernetworks/hypernetwork.py
index 5608e799..470659df 100644
--- a/modules/hypernetworks/hypernetwork.py
+++ b/modules/hypernetworks/hypernetwork.py
@@ -14,6 +14,7 @@ import torch
 from torch import einsum

 from einops import rearrange, repeat

 import modules.textual_inversion.dataset

+from modules.textual_inversion.learn_schedule import LearnSchedule

 

 

 class HypernetworkModule(torch.nn.Module):

@@ -202,8 +203,6 @@ def train_hypernetwork(hypernetwork_name, learn_rate, data_root, log_directory,
     for weight in weights:

         weight.requires_grad = True

 

-    optimizer = torch.optim.AdamW(weights, lr=learn_rate)

-

     losses = torch.zeros((32,))

 

     last_saved_file = "<none>"

@@ -213,12 +212,24 @@ def train_hypernetwork(hypernetwork_name, learn_rate, data_root, log_directory,
     if ititial_step > steps:

         return hypernetwork, filename

 

+    schedules = iter(LearnSchedule(learn_rate, steps, ititial_step))

+    (learn_rate, end_step) = next(schedules)

+    print(f'Training at rate of {learn_rate} until step {end_step}')

+

+    optimizer = torch.optim.AdamW(weights, lr=learn_rate)

+

     pbar = tqdm.tqdm(enumerate(ds), total=steps - ititial_step)

     for i, (x, text, cond) in pbar:

         hypernetwork.step = i + ititial_step

 

-        if hypernetwork.step > steps:

-            break

+        if hypernetwork.step > end_step:

+            try:

+                (learn_rate, end_step) = next(schedules)

+            except Exception:

+                break

+            tqdm.tqdm.write(f'Training at rate of {learn_rate} until step {end_step}')

+            for pg in optimizer.param_groups:

+                pg['lr'] = learn_rate

 

         if shared.state.interrupted:

             break