Learning rate tuning ??

neverix · Sep 2, 2024 · 5a40ca9 · 5a40ca9
1 parent 90fd37f
commit 5a40ca9
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/lr_search.png b/lr_search.png
diff --git a/src/jif/__main__.py b/src/jif/__main__.py
@@ -259,9 +259,11 @@ def main(*args, **kwargs):
         log_dict = train(*args, **kwargs | {"n_steps": 2_000, "lr": lr, "size": "small"},
                          quiet=True, profile=False)
         losses[model].append(log_dict["loss"])
+
         for model, color in zip(models, colors):
             plt.plot(lrs_sampled[model], losses[model], label=model, c=color)
-            plt.scatter(lrs_sampled[model], losses[model], label=model, marker="x", c=color)
+            plt.scatter(lrs_sampled[model], losses[model], marker="x", c=color)
+        plt.xlim(lrs[0], lrs[-1])
         plt.xlabel("Learning rate")
         plt.ylabel("Final loss")
         plt.legend()
@@ -270,4 +272,5 @@ def main(*args, **kwargs):
 
 
 if __name__ == "__main__":
+    # fire.Fire(train); exit()
     fire.Fire(main)