added a few small tools

Author: ChadLin9596

src/py_utils/automation.py

@@ -7,3 +7,34 @@ def tsp_submit(*cmds):
     tsp_cmd.extend(cmds)
     result = subprocess.run(tsp_cmd, capture_output=True, text=True)
     return int(result.stdout.strip())
+
+
+def _form_row_header(N):
+
+    row_0 = ["|"] * (N + 1)
+    row_0 = (" " * 3).join(row_0)
+    row_1 = ["|"] * (N + 1)
+    row_1 = " --- ".join(row_1)
+
+    return "\n".join([row_0, row_1])
+
+
+def form_in_md_rows(*datas):
+
+    N = len(datas[0])
+    for data in datas:
+        assert N == len(data)
+
+    header = _form_row_header(N)
+
+    prefix = "| "
+    suffix = " |"
+
+    rows = []
+    for data in datas:
+        row = map(str, data)
+        row = " | ".join(row)
+        row = prefix + row + suffix
+        rows.append(row)
+
+    return "\n".join([header] + rows)

src/py_utils/utils_torch.py

@@ -541,3 +541,14 @@ def segmented_softmax(logits, groups, n_groups):
     denominator = denominator[groups] + 1e-12
 
     return exp_logits / denominator
+
+
+def count_model_parameters(model, trainable_only=False):
+
+    num_params = []
+    for p in model.parameters():
+        if trainable_only and not p.requires_grad:
+            continue
+        num_params.append(p.numel())
+
+    return sum(num_params)