Promote hard fine-tune

config.json

@@ -3,7 +3,7 @@
   "dim": 256,
   "embed_dim": 256,
   "heads": 8,
-  "max_crossings": 64,
+  "max_crossings": 128,
   "target_mean": [
     4.885252952575684,
     17.482948303222656

knot2vec.py

@@ -182,11 +182,11 @@ class Knot2Vec(nn.Module, PyTorchModelHubMixin):
         return self.run(pds, batch_size)["embedding"]
 
 
-def scramble(link, max_crossings):
+def scramble(link, max_crossings, max_moves=40):
     """Another diagram of the same knot, from random Reidemeister moves."""
     while True:
         view = link.copy()
-        view.backtrack(random.randint(2, 40))
+        view.backtrack(random.randint(2, max_moves))
         if random.random() < 0.5:
             view.simplify("basic")
         if 0 < len(view.crossings) <= max_crossings:
@@ -199,13 +199,10 @@ def diagrams(job):
     random.seed(f"{args.seed}:{name}")
     link = snappy.Link(pd)
     rows = [{"name": name, "pd": pd, "split": "train"}] if seen else []
+    splits = ["train"] * (args.views if seen else 0) + ["test"] * args.test_views
     rows += [
-        {"name": name, "pd": scramble(link, args.max_crossings), "split": "train"}
-        for _ in range(args.views if seen else 0)
-    ]
-    rows += [
-        {"name": name, "pd": scramble(link, args.max_crossings), "split": "test"}
-        for _ in range(args.test_views)
+        {"name": name, "pd": scramble(link, args.max_crossings, args.max_moves), "split": split}
+        for split in splits
     ]
     for row in rows:
         gauss(row["pd"])
@@ -307,6 +304,12 @@ def train(args):
         target_mean=targets.mean(0).tolist(),
         target_std=(targets.std(0) + 1e-6).tolist(),
     ).to(device)
+    if args.init:
+        # Copy the weights; embeddings for longer diagrams keep their fresh rows.
+        weights = model.state_dict()
+        for key, value in Knot2Vec.from_pretrained(args.init).state_dict().items():
+            weights[key][: len(value)] = value
+        model.load_state_dict(weights)
     targets = torch.tensor((targets - targets.mean(0)) / (targets.std(0) + 1e-6), device=device)
     signatures = torch.tensor(
         [(knots[n]["signature"] + MAX_SIGNATURE) // 2 for n in names], device=device
@@ -389,6 +392,7 @@ def train(args):
                 "learning_rate": args.lr,
                 "temperature": args.temperature,
                 "aux_weight": args.aux_weight,
+                "init": args.init,
                 "dim": args.dim,
                 "depth": args.depth,
                 "parameters": sum(p.numel() for p in model.parameters()),
@@ -433,6 +437,8 @@ def evaluate(args):
     index = {n: i for i, n in enumerate(knots["name"])}
     gallery = model.embed(knots["pd"]).to(device)
     test = table(args.data, "diagrams", "test")
+    longest = model.positions.num_embeddings // 2
+    test = test.filter(lambda row: len(row["pd"]) <= longest)
     if args.limit:
         test = test.select(random.Random(0).sample(range(len(test)), min(args.limit, len(test))))
     out = model.run(test["pd"])
@@ -441,7 +447,8 @@ def evaluate(args):
     truth = {key: torch.tensor(knots[key])[gold] for key in
              ("seen", "crossing_number", "signature", "determinant", "volume")}
     hits1, hits5 = top[:, 0] == gold, (top == gold[:, None]).any(1)
-    result = {"model": args.model, "gallery": len(index), "chance_top1": round(1 / len(index), 6)}
+    result = {"model": args.model, "data": args.data, "max_crossings": longest,
+              "gallery": len(index), "chance_top1": round(1 / len(index), 6)}
     for group, mask in (("seen", truth["seen"]), ("unseen", ~truth["seen"])):
         if not mask.any():
             continue
@@ -467,7 +474,7 @@ def evaluate(args):
     Path(args.output).mkdir(parents=True, exist_ok=True)
     knot_map(gallery, knots, Path(args.output) / "map.png")
     if args.push:
-        push_result(args.repo, "eval", result, args.revision)
+        push_result(args.repo, args.result, result, args.revision)
         HfApi().upload_file(
             path_or_fileobj=Path(args.output) / "map.png",
             path_in_repo="results/map.png",
@@ -570,6 +577,7 @@ def main():
     data_parser.add_argument("--test-views", type=int, default=4, help="test diagrams per knot")
     data_parser.add_argument("--holdout", type=float, default=0.1, help="share of unseen knots")
     data_parser.add_argument("--max-crossings", type=int, default=64)
+    data_parser.add_argument("--max-moves", type=int, default=40, help="Reidemeister moves")
     data_parser.add_argument("--seed", type=int, default=0)
     data_parser.add_argument("--output", default="out/data")
     data_parser.add_argument("--repo", default=DATA)
@@ -587,6 +595,7 @@ def main():
     train_parser.add_argument("--dim", type=int, default=256)
     train_parser.add_argument("--depth", type=int, default=6)
     train_parser.add_argument("--max-crossings", type=int, default=64)
+    train_parser.add_argument("--init", default=None, help="model repo or folder to start from")
     train_parser.add_argument("--seed", type=int, default=0)
     train_parser.add_argument("--output", default="out/model")
     train_parser.add_argument("--repo", default=REPO)
@@ -600,6 +609,7 @@ def main():
     eval_parser.add_argument("--limit", type=int, default=None, help="number of test diagrams")
     eval_parser.add_argument("--repo", default=REPO, help="where --push stores the result")
     eval_parser.add_argument("--revision", default=None, help="model branch")
+    eval_parser.add_argument("--result", default="eval", help="results/<result>.json for --push")
     eval_parser.add_argument("--output", default="out/eval", help="where the knot map goes")
     eval_parser.add_argument("--push", action="store_true")
     eval_parser.set_defaults(run=evaluate)

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:786f8b9cb77a3b430caaa584312e3db9a5939c4f480c5b038963634b642d205a
-size 19708148
+oid sha256:98ac2a4175b287b9ddf642ab689e897c7b56ac3df721c78609bbff43a408840e
+size 19904756

results/eval-hard.json

@@ -0,0 +1,46 @@
+{
+ "model": "jgalego/knot2vec",
+ "data": "jgalego/knot2vec-diagrams-hard",
+ "max_crossings": 128,
+ "gallery": 12965,
+ "chance_top1": 7.7e-05,
+ "seen": {
+  "n": 46676,
+  "top1": 0.773,
+  "top5": 0.916,
+  "signature_accuracy": 0.88,
+  "determinant_within_10pct": 0.238,
+  "volume_mae": 1.39,
+  "top1_by_crossings": {
+   "3": 1.0,
+   "5": 0.75,
+   "6": 0.875,
+   "7": 0.786,
+   "8": 0.583,
+   "9": 0.71,
+   "10": 0.618,
+   "11": 0.719,
+   "12": 0.739,
+   "13": 0.787
+  }
+ },
+ "unseen": {
+  "n": 5184,
+  "top1": 0.771,
+  "top5": 0.91,
+  "signature_accuracy": 0.849,
+  "determinant_within_10pct": 0.23,
+  "volume_mae": 1.552,
+  "top1_by_crossings": {
+   "4": 0.5,
+   "5": 0.75,
+   "6": 1.0,
+   "8": 0.417,
+   "9": 0.5,
+   "10": 0.515,
+   "11": 0.726,
+   "12": 0.741,
+   "13": 0.787
+  }
+ }
+}

results/eval-hard64.json

@@ -0,0 +1,46 @@
+{
+ "model": "jgalego/knot2vec",
+ "data": "jgalego/knot2vec-diagrams-hard",
+ "max_crossings": 64,
+ "gallery": 12965,
+ "chance_top1": 7.7e-05,
+ "seen": {
+  "n": 40047,
+  "top1": 0.84,
+  "top5": 0.962,
+  "signature_accuracy": 0.91,
+  "determinant_within_10pct": 0.253,
+  "volume_mae": 1.29,
+  "top1_by_crossings": {
+   "3": 1.0,
+   "5": 0.667,
+   "6": 1.0,
+   "7": 0.808,
+   "8": 0.656,
+   "9": 0.769,
+   "10": 0.691,
+   "11": 0.789,
+   "12": 0.806,
+   "13": 0.854
+  }
+ },
+ "unseen": {
+  "n": 4412,
+  "top1": 0.84,
+  "top5": 0.958,
+  "signature_accuracy": 0.881,
+  "determinant_within_10pct": 0.243,
+  "volume_mae": 1.45,
+  "top1_by_crossings": {
+   "4": 0.5,
+   "5": 0.75,
+   "6": 1.0,
+   "8": 0.5,
+   "9": 0.526,
+   "10": 0.642,
+   "11": 0.793,
+   "12": 0.82,
+   "13": 0.854
+  }
+ }
+}

results/eval.json

@@ -1,44 +1,46 @@
 {
  "model": "jgalego/knot2vec",
+ "data": "jgalego/knot2vec-diagrams",
+ "max_crossings": 128,
  "gallery": 12965,
  "chance_top1": 7.7e-05,
  "seen": {
   "n": 46676,
-  "top1": 0.881,
-  "top5": 0.976,
-  "signature_accuracy": 0.917,
-  "determinant_within_10pct": 0.267,
-  "volume_mae": 1.237,
+  "top1": 0.874,
+  "top5": 0.978,
+  "signature_accuracy": 0.926,
+  "determinant_within_10pct": 0.264,
+  "volume_mae": 1.21,
   "top1_by_crossings": {
-   "3": 1.0,
-   "5": 0.75,
+   "3": 0.75,
+   "5": 0.5,
    "6": 1.0,
    "7": 0.821,
-   "8": 0.833,
-   "9": 0.756,
-   "10": 0.769,
-   "11": 0.844,
-   "12": 0.864,
-   "13": 0.889
+   "8": 0.861,
+   "9": 0.727,
+   "10": 0.731,
+   "11": 0.84,
+   "12": 0.851,
+   "13": 0.884
   }
  },
  "unseen": {
   "n": 5184,
-  "top1": 0.88,
+  "top1": 0.867,
   "top5": 0.976,
-  "signature_accuracy": 0.891,
-  "determinant_within_10pct": 0.248,
-  "volume_mae": 1.39,
+  "signature_accuracy": 0.896,
+  "determinant_within_10pct": 0.246,
+  "volume_mae": 1.373,
   "top1_by_crossings": {
    "4": 0.75,
    "5": 0.75,
-   "6": 0.75,
+   "6": 1.0,
    "8": 0.917,
-   "9": 0.65,
+   "9": 0.8,
    "10": 0.779,
-   "11": 0.806,
-   "12": 0.86,
-   "13": 0.893
+   "11": 0.782,
+   "12": 0.852,
+   "13": 0.878
   }
  }
 }

results/probes.json

@@ -4,45 +4,45 @@
   "by_crossings": {
    "0": {
     "n": 334,
-    "top1": 0.988
+    "top1": 0.979
    },
    "11": {
     "n": 30554,
-    "top1": 0.946
+    "top1": 0.935
    },
    "21": {
     "n": 7414,
-    "top1": 0.87
+    "top1": 0.85
    },
    "31": {
     "n": 7375,
-    "top1": 0.794
+    "top1": 0.777
    },
    "41": {
     "n": 6183,
-    "top1": 0.67
+    "top1": 0.702
    }
   },
   "by_writhe_shift": {
    "0": {
     "n": 25800,
-    "top1": 0.938
+    "top1": 0.927
    },
    "1": {
     "n": 18850,
-    "top1": 0.848
+    "top1": 0.838
    },
    "3": {
     "n": 6298,
-    "top1": 0.773
+    "top1": 0.774
    },
    "6": {
     "n": 900,
-    "top1": 0.711
+    "top1": 0.731
    },
    "11": {
     "n": 12,
-    "top1": 0.583
+    "top1": 0.75
    }
   }
  },
@@ -52,41 +52,41 @@
    "kept": 1.0,
    "mean_crossings": 32.6,
    "mean_writhe_shift": 3.5,
-   "top1": 0.894
+   "top1": 0.868
   },
   "r2_only": {
    "n": 2592,
    "kept": 1.0,
    "mean_crossings": 32.6,
    "mean_writhe_shift": 0.0,
-   "top1": 0.877
+   "top1": 0.855
   },
   "moves_40": {
    "n": 2592,
    "kept": 1.0,
    "mean_crossings": 48.8,
    "mean_writhe_shift": 2.8,
-   "top1": 0.61
+   "top1": 0.676
   },
   "moves_60": {
    "n": 2591,
    "kept": 1.0,
    "mean_crossings": 60.1,
    "mean_writhe_shift": 3.3,
-   "top1": 0.273
+   "top1": 0.576
   },
   "moves_80": {
    "n": 138,
    "kept": 0.053,
    "mean_crossings": 62.1,
    "mean_writhe_shift": 3.8,
-   "top1": 0.188
+   "top1": 0.514
   }
  },
  "small": {
   "up_to_7": {
    "knots": 14,
-   "top1": 0.874,
+   "top1": 0.873,
    "by_growth": {
     "1": {
      "n": 1555,
@@ -94,39 +94,39 @@
     },
     "2": {
      "n": 561,
-     "top1": 0.829
+     "top1": 0.809
     },
     "4": {
      "n": 595,
-     "top1": 0.713
+     "top1": 0.718
     },
     "8": {
      "n": 89,
-     "top1": 0.674
+     "top1": 0.742
     }
    },
-   "errors_named_up_to_7": 0.328,
-   "errors_named_11_plus": 0.441
+   "errors_named_up_to_7": 0.361,
+   "errors_named_11_plus": 0.392
   },
   "12_13": {
    "knots": 14,
-   "top1": 0.907,
+   "top1": 0.892,
    "by_growth": {
     "1": {
      "n": 1875,
-     "top1": 0.966
+     "top1": 0.953
     },
     "2": {
      "n": 881,
-     "top1": 0.799
+     "top1": 0.778
     },
     "4": {
      "n": 44,
-     "top1": 0.568
+     "top1": 0.636
     }
    },
    "errors_named_up_to_7": 0.0,
-   "errors_named_11_plus": 0.992
+   "errors_named_11_plus": 0.997
   }
  }
 }

results/train.json

@@ -1,19 +1,20 @@
 {
- "step": 19999,
- "loss": 0.6226,
- "contrastive": 0.1426,
- "invariants": 0.48,
- "pair_top1": 0.961,
- "minutes": 101.8,
- "data": "jgalego/knot2vec-diagrams",
- "steps": 20000,
+ "step": 9999,
+ "loss": 1.1689,
+ "contrastive": 0.5045,
+ "invariants": 0.6644,
+ "pair_top1": 0.869,
+ "minutes": 92.2,
+ "data": "jgalego/knot2vec-diagrams-hard",
+ "steps": 10000,
  "batch_size": 512,
- "learning_rate": 0.0003,
+ "learning_rate": 0.0001,
  "temperature": 0.05,
  "aux_weight": 1.0,
+ "init": "jgalego/knot2vec",
  "dim": 256,
  "depth": 6,
- "parameters": 4924943,
- "runtime_s": 6109,
+ "parameters": 4974095,
+ "runtime_s": 5533,
  "device": "NVIDIA A10G"
 }