Merge branch 'master' into tensorboard

1 files changed, 30 insertions, 18 deletions

diff --git a/modules/textual_inversion/learn_schedule.py b/modules/textual_inversion/learn_schedule.py
index 2062726a..f63fc72f 100644
--- a/modules/textual_inversion/learn_schedule.py
+++ b/modules/textual_inversion/learn_schedule.py
@@ -4,30 +4,37 @@ import tqdm
 class LearnScheduleIterator:

     def __init__(self, learn_rate, max_steps, cur_step=0):

         """

-        specify learn_rate as "0.001:100, 0.00001:1000, 1e-5:10000" to have lr of 0.001 until step 100, 0.00001 until 1000, 1e-5:10000 until 10000

+        specify learn_rate as "0.001:100, 0.00001:1000, 1e-5:10000" to have lr of 0.001 until step 100, 0.00001 until 1000, and 1e-5 until 10000

         """

 

         pairs = learn_rate.split(',')

         self.rates = []

         self.it = 0

         self.maxit = 0

-        for i, pair in enumerate(pairs):

-            tmp = pair.split(':')

-            if len(tmp) == 2:

-                step = int(tmp[1])

-                if step > cur_step:

-                    self.rates.append((float(tmp[0]), min(step, max_steps)))

-                    self.maxit += 1

-                    if step > max_steps:

+        try:

+            for i, pair in enumerate(pairs):

+                if not pair.strip():

+                    continue

+                tmp = pair.split(':')

+                if len(tmp) == 2:

+                    step = int(tmp[1])

+                    if step > cur_step:

+                        self.rates.append((float(tmp[0]), min(step, max_steps)))

+                        self.maxit += 1

+                        if step > max_steps:

+                            return

+                    elif step == -1:

+                        self.rates.append((float(tmp[0]), max_steps))

+                        self.maxit += 1

                         return

-                elif step == -1:

+                else:

                     self.rates.append((float(tmp[0]), max_steps))

                     self.maxit += 1

                     return

-            else:

-                self.rates.append((float(tmp[0]), max_steps))

-                self.maxit += 1

-                return

+            assert self.rates

+        except (ValueError, AssertionError):

+            raise Exception('Invalid learning rate schedule. It should be a number or, for example, like "0.001:100, 0.00001:1000, 1e-5:10000" to have lr of 0.001 until step 100, 0.00001 until 1000, and 1e-5 until 10000.')

+

 

     def __iter__(self):

         return self

@@ -51,14 +58,19 @@ class LearnRateScheduler:
 

         self.finished = False

 

-    def apply(self, optimizer, step_number):

-        if step_number <= self.end_step:

-            return

+    def step(self, step_number):

+        if step_number < self.end_step:

+            return False

 

         try:

             (self.learn_rate, self.end_step) = next(self.schedules)

-        except Exception:

+        except StopIteration:

             self.finished = True

+            return False

+        return True

+

+    def apply(self, optimizer, step_number):

+        if not self.step(step_number):

             return

 

         if self.verbose: