Add another trainer

Author: JasonGross

gbmi/exp_modular_arithmetic/run_train_clocks_faster.py

@@ -0,0 +1,54 @@
+# %%
+from tqdm.auto import tqdm
+
+from gbmi.exp_modular_arithmetic import SEEDS
+from gbmi.exp_modular_arithmetic.train import FAST_CLOCK_CONFIG, train_or_load_model
+from gbmi.training_tools.logging import ModelMatrixLoggingOptions
+from gbmi.utils import set_params
+
+with tqdm(SEEDS, desc="Seed") as pbar:
+    for seed in pbar:
+        pbar.set_postfix({"seed": seed})
+        runtime, model = train_or_load_model(
+            set_params(
+                FAST_CLOCK_CONFIG,
+                {
+                    "seed": seed,
+                },
+                post_init=True,
+            ),
+            # force="load",
+            # force="train",
+        )
+
+
+# %%
+import shutil
+from pathlib import Path
+
+import torch
+
+base = Path(".").resolve()
+wandbs = (base / "artifacts").glob("*/*.pth")
+total = len(list((base / "artifacts").glob("*/*.pth")))
+model_base = base / "models"
+model_base.mkdir(exist_ok=True, parents=True)
+# %%
+with tqdm(wandbs, total=total) as pbar:
+    for p in pbar:
+        total -= 1
+        cache = torch.load(p, map_location="cpu")
+        pbar.set_postfix(
+            {
+                "seed": cache["run_config"]["seed"],
+                "orig_name": p.name,
+                "suffix_drop": "-".join(p.name.split("-")[-6:]),
+            }
+        )
+        seed = cache["run_config"]["seed"]
+        shutil.copy(
+            p,
+            model_base / f"{'-'.join(p.name.split('-')[:-6])}-{seed}{p.suffix}",
+        )
+    # break
+# %%

gbmi/exp_modular_arithmetic/train.py

@@ -187,6 +187,20 @@ def log_softmax(self, x: Tensor, **kwargs) -> Tensor:
     checkpoint_every=(500, "epochs"),
 )
 
+FAST_CLOCK_CONFIG = Config(
+    experiment=ModularArithmetic(
+        p=113,
+        training_ratio=0.8,
+        logging_options=ModelMatrixLoggingOptions.none(),
+    ),
+    seed=0,
+    deterministic=False,
+    train_for=(10000, "epochs"),
+    log_every_n_steps=1,
+    validate_every=(100, "epochs"),
+    checkpoint_every=(2000, "epochs"),
+)
+
 PIZZA_CONFIG = Config(
     experiment=ModularArithmetic(
         p=59,