[REF] Randomly initialize nets and buffers, support specifying mode

plutonium-239 · Jul 4, 2024 · a354d9b · a354d9b
1 parent 2c5afa1
commit a354d9b
Show file tree

Hide file tree

Showing 4 changed files with 203 additions and 171 deletions.
diff --git a/experiments/visual_abstract/gather_data.py b/experiments/visual_abstract/gather_data.py
@@ -5,8 +5,7 @@
 
 from pandas import DataFrame
 
-HERE = path.abspath(__file__)
-HEREDIR = path.dirname(HERE)
+HEREDIR = path.dirname(path.abspath(__file__))
 RAWDATADIR = path.join(HEREDIR, "raw")
 DATADIR = path.join(HEREDIR, "gathered")
 makedirs(RAWDATADIR, exist_ok=True)
@@ -15,32 +14,33 @@
 max_num_layers = 10
 requires_grads = ["all", "none", "4", "4+"]
 implementations = ["torch", "ours"]
-architectures = ["linear", "conv", "norm_eval"]
-architectures = ["norm_eval"]
-architectures = ["linear"]
+architectures = ["linear", "conv", "bn"]
+modes = ["eval", "train"]
 
 if __name__ == "__main__":
-    for implementation, requires_grad, arch in product(
-        implementations, requires_grads, architectures
+    for implementation, requires_grad, architecture, mode in product(
+        implementations, requires_grads, architectures, modes
     ):
         if implementation == "ours" and requires_grad != "4":
             continue
+        if mode == "eval" and architecture != "bn":
+            continue
 
-        layers = list(range(1, max_num_layers + 1))
         peakmems = []
+        layers = list(range(1, max_num_layers + 1))
         for num_layers in layers:
-            with open(
-                path.join(
-                    RAWDATADIR,
-                    f"peakmem_implementation_{arch}_{implementation}_num_layers_{num_layers}_requires_grad_{requires_grad}.txt",
-                ),
-                "r",
-            ) as f:
+            readpath = path.join(
+                RAWDATADIR,
+                f"peakmem_{architecture}_mode_{mode}_implementation_{implementation}"
+                + f"_num_layers_{num_layers}_requires_grad_{requires_grad}.txt",
+            )
+            with open(readpath, "r") as f:
                 peakmems.append(float(f.read()))
 
         df = DataFrame({"num_layers": layers, "peakmem": peakmems})
         savepath = path.join(
             DATADIR,
-            f"peakmem_implementation_{arch}_{implementation}_requires_grad_{requires_grad}.csv",
+            f"peakmem_{architecture}_mode_{mode}_implementation_{implementation}"
+            + f"_requires_grad_{requires_grad}.csv",
         )
         df.to_csv(savepath, index=False)
diff --git a/experiments/visual_abstract/generate_data.py b/experiments/visual_abstract/generate_data.py
@@ -12,12 +12,10 @@
 
 max_num_layers = 10
 requires_grads = ["all", "none", "4", "4+"]
-# requires_grads = ["4+"]
 implementations = ["torch", "ours"]
-# implementations = ["ours"]
-architectures = ["linear", "conv", "norm_eval"]
-architectures = ["norm_eval"]
-architectures = ["linear"]
+architectures = ["linear", "conv", "bn"]
+modes = ["eval", "train"]
+skip_existing = True
 
 
 def _run(cmd: List[str]):
@@ -41,19 +39,23 @@ def _run(cmd: List[str]):
 
 
 if __name__ == "__main__":
-    for implementation, requires_grad, arch in product(
-        implementations, requires_grads, architectures
+    for implementation, requires_grad, architecture, mode in product(
+        implementations, requires_grads, architectures, modes
     ):
         if implementation == "ours" and requires_grad != "4":
             continue
+        if mode == "eval" and architecture != "bn":
+            continue
         for num_layers in range(1, max_num_layers + 1):
             _run(
                 [
                     "python",
                     SCRIPT,
                     f"--implementation={implementation}",
-                    f"--architecture={arch}",
+                    f"--architecture={architecture}",
                     f"--num_layers={num_layers}",
                     f"--requires_grad={requires_grad}",
+                    f"--mode={mode}",
                 ]
+                + (["--skip_existing"] if skip_existing else [])
             )
diff --git a/experiments/visual_abstract/plot_data.py b/experiments/visual_abstract/plot_data.py
@@ -1,14 +1,13 @@
 """Visualize memory consumpion."""
 
+from itertools import product
 from os import path
 
 from matplotlib import pyplot as plt
 from pandas import read_csv
 from tueplots import bundles
 
-HERE = path.abspath(__file__)
-HEREDIR = path.dirname(HERE)
-
+HEREDIR = path.dirname(path.abspath(__file__))
 DATADIR = path.join(HEREDIR, "gathered")
 
 requires_grads = ["all", "none", "4+", "4"]
@@ -33,57 +32,61 @@
     "4": "dashdot",
     "4 (ours)": "dotted",
 }
-architectures = ["linear", "conv", "norm_eval"]
-architectures = ["norm_eval"]
-architectures = ["conv"]
+architectures = ["linear", "conv", "bn"]
+modes = ["train", "eval"]
+
+if __name__ == "__main__":
+    for architecture, mode in product(architectures, modes):
+        if mode == "eval" and architecture != "bn":
+            continue
 
-for arch in architectures:
-    with plt.rc_context(bundles.icml2024()):
-        plt.rcParams.update({"figure.figsize": (3.25, 2.5)})
-        fig, ax = plt.subplots()
-        ax.set_xlabel("Number of layers")
-        ax.set_ylabel("Peak memory [MiB]")
+        with plt.rc_context(bundles.icml2024()):
+            # plt.rcParams.update({"figure.figsize": (3.25, 2.5)})
+            fig, ax = plt.subplots()
+            ax.set_xlabel("Number of layers")
+            ax.set_ylabel("Peak memory [MiB]")
 
-        markerstyle = {"markersize": 3.5, "fillstyle": "none"}
+            markerstyle = {"markersize": 3.5, "fillstyle": "none"}
 
-        # visualize PyTorch's behavior
-        implementation = "torch"
+            # visualize PyTorch's behavior
+            implementation = "torch"
 
-        for requires_grad in requires_grads:
-            df = read_csv(
-                path.join(
+            for requires_grad in requires_grads:
+                readpath = path.join(
                     DATADIR,
-                    f"peakmem_implementation_{arch}_{implementation}_requires_grad_{requires_grad}.csv",
+                    f"peakmem_{architecture}_mode_{mode}_implementation_{implementation}"
+                    + f"_requires_grad_{requires_grad}.csv",
                 )
+                df = read_csv(readpath)
+                ax.plot(
+                    df["num_layers"],
+                    df["peakmem"],
+                    label=legend_entries[requires_grad],
+                    marker=markers[requires_grad],
+                    linestyle=linestyles[requires_grad],
+                    **markerstyle,
+                )
+
+            # visualize our layer's behavior
+            implementation, requires_grad = "ours", "4"
+            key = f"{requires_grad} ({implementation})"
+            readpath = path.join(
+                DATADIR,
+                f"peakmem_{architecture}_mode_{mode}_implementation_{implementation}"
+                + f"_requires_grad_{requires_grad}.csv",
             )
+            df = read_csv(readpath)
             ax.plot(
                 df["num_layers"],
                 df["peakmem"],
-                label=legend_entries[requires_grad],
-                marker=markers[requires_grad],
-                linestyle=linestyles[requires_grad],
+                label=legend_entries[key],
+                marker=markers[key],
+                linestyle=linestyles[key],
                 **markerstyle,
             )
 
-        # visualize our layer's behavior
-        implementation, requires_grad = "ours", "4"
-        key = f"{requires_grad} ({implementation})"
-        df = read_csv(
-            path.join(
-                DATADIR,
-                f"peakmem_implementation_{arch}_{implementation}_requires_grad_{requires_grad}.csv",
+            plt.legend()
+            plt.savefig(
+                path.join(HEREDIR, f"visual_abstract_{architecture}_{mode}.pdf"),
+                bbox_inches="tight",
             )
-        )
-        ax.plot(
-            df["num_layers"],
-            df["peakmem"],
-            label=legend_entries[key],
-            marker=markers[key],
-            linestyle=linestyles[key],
-            **markerstyle,
-        )
-
-        plt.legend()
-        plt.savefig(
-            path.join(HEREDIR, f"visual_abstract_{arch}.pdf"), bbox_inches="tight"
-        )