Merge branch 'main' into visual-abstract+poster

.gitignore

@@ -2,6 +2,8 @@
 *.txt
 *.csv
 !requirements.txt
+torchviz-output/
+torchview-output/
 
 # generated docs
 docs_src/_build/

experiments/best_results_to_latex.py

@@ -0,0 +1,116 @@
+"""Simple script to make a latex table from best results"""
+
+import argparse
+
+import pandas as pd
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--input",
+    type=str,
+    required=True,
+    help="Path to best_results<params>.csv generated by get_best_results",
+)
+
+
+args = parser.parse_args()
+
+df = pd.read_csv(args.input)
+
+df = df.set_index("model")
+df = df[df["case"] != "Conv"]
+df = df[df["case"] != "SurgicalLast"]
+
+df["memsave"] = df.index.str.startswith("memsave_")
+badi = df.index.map(
+    lambda x: x.split("memsave_", 1)[1] if x.startswith("memsave") else x
+)
+badi.name = "model_clean"
+df2 = df.reset_index().set_index(badi).sort_index()
+divs = df2[(df2["case"] == "All") & (~df2["memsave"])]
+df2["Scaled M"] = df2["Memory Usage (GB)"] / divs["Memory Usage (GB)"]
+df2["Scaled T"] = df2["Time Taken (s)"] / divs["Time Taken (s)"]
+
+df2["Memory [GiB]"] = df2.apply(
+    lambda x: f"{x['Memory Usage (GB)']:.2f} ({x['Scaled M']:.2f})", axis=1
+)
+df2["Time [s]"] = df2.apply(
+    lambda x: f"{x['Time Taken (s)']:.2f} ({x['Scaled T']:.2f})", axis=1
+)
+
+
+def _highlight(group, col_sort, col_bold):
+    for c_s, c_b in zip(col_sort, col_bold):
+        min_idx = group[c_s].argmin()
+        group[c_b] = [
+            f"\\textbf{{{group.iloc[i][c_b]}}}" if i == min_idx else group.iloc[i][c_b]
+            for i in range(len(group.index))
+        ]
+    return group
+
+
+df2 = df2.groupby(["model_clean", "case"]).apply(
+    _highlight, ["Memory Usage (GB)"], ["Memory [GiB]"]
+)
+# .apply(_highlight, ['Memory Usage (GB)', 'Time Taken (s)'], ['Memory [GiB]', 'Time [s]'])
+
+names = {
+    "bert": "BERT",
+    "bart": "BART",
+    "roberta": "RoBERTa",
+    "gpt2": "GPT-2",
+    "t5": "T5 \\cite{JMLR_t5}",
+    "flan-t5": "FLAN-T5",
+    "mistral-7b": "Mistral-7B \\cite{jiang2023mistral}",
+    "transformer": "Transformer \\cite{NIPS2017_3f5ee243_vaswaniattention}",
+    "llama3-8b": "LLaMa3-8B \\cite{touvron2023llama}",
+    "phi3-4b": "Phi3-4B \\cite{gunasekar2023textbooksPhi}",
+    # Conv
+    "deeplabv3_resnet101": "DeepLabv3 (RN101) \\cite{deeplabv3_chen2017rethinking}",
+    "efficientnet_v2_l": "EfficientNetv2-L \\cite{efficientnet_TanL19,efficientnetv2_TanL21}",
+    "fcn_resnet101": "FCN (RN101) \\cite{fcn}",
+    "mobilenet_v3_large": "MobileNetv3-L \\cite{mobilenetv3}",
+    "resnext101_64x4d": "ResNeXt101-64x4d \\cite{resnext_cvpr_XieGDTH17}",
+    "fasterrcnn_resnet50_fpn_v2": "Faster-RCNN (RN101) \\cite{faster_rcnn_RenHGS15}",
+    "ssdlite320_mobilenet_v3_large": "SSDLite (MobileNetv3-L) \\cite{mobilenetv2_Sandler_2018_CVPR}",
+    "vgg16": "VGG-16 \\cite{vgg_SimonyanZ14a}",
+}
+
+# import ipdb; ipdb.set_trace()
+df2 = df2[df2.index.isin(names.keys(), level=0)]
+
+
+def _format_name(n):
+    if n.startswith("memsave_"):
+        mname = n.split("memsave_", 1)[1]
+        return f"{names[mname]} + MemSave"
+    return names[n]
+
+
+ni = df2["model"].apply(_format_name)
+df2 = df2.set_index(ni).sort_index().drop(
+    columns=[
+        "model",
+        "memsave",
+        "Memory Usage (GB)",
+        "Time Taken (s)",
+        "Scaled M",
+        "Scaled T",
+    ]
+)  # fmt: skip
+
+df2_p = df2.pivot_table(
+    index="model", columns="case", values=df2.columns[1:], aggfunc=lambda x: x
+)
+
+short_index = df2_p.index.map(lambda t: "+ MemSave" if "+ MemSave" in t else t)
+df2_p = df2_p.set_index(short_index)
+
+latex_str = df2_p.to_latex(na_rep="-", multicolumn_format="c")
+final_str = ""
+for line in latex_str.split("\n"):
+    add_line = line + "\n"
+    if line.startswith("+ MemSave"):
+        add_line += "\\midrule\n"
+    final_str += add_line
+print(final_str)

experiments/get_best_results.py

@@ -7,7 +7,7 @@
 
 import pandas as pd
 
-from experiments.util.collect_results import case_mapping
+from experiments.util.collect_results import case_inv_mapping
 
 
 def main(base_dir: str):
@@ -16,17 +16,25 @@ def main(base_dir: str):
     Args:
         base_dir (str): The base results dir
     """
-    for device, arch in product(["cuda", "cpu"], ["linear", "conv"]):
+    # Don't recognize None as NaN
+    custom_na_values = pd._libs.parsers.STR_NA_VALUES - {"None"}
+    for device, arch in product(["cuda", "cpu"], ["linear", "conv", "transformer"]):
         # usage stats
         df = None
         idx_col = ["model", "case"]
         for fname in glob(os.path.join(base_dir, f"usage_stats-{arch}-{device}-*.csv")):
             with open(fname) as f:
-                f.readline()
-                temp_df = pd.read_csv(f, index_col=idx_col)
+                # f.readline()
+                temp_df = pd.read_csv(
+                    f,
+                    index_col=idx_col,
+                    header=1,
+                    na_values=custom_na_values,
+                    keep_default_na=False,
+                )
             df = temp_df if df is None else pd.concat([df, temp_df])
         if df is not None:
-            df = df.rename(index=case_mapping, level=1)
+            df = df.rename(index=case_inv_mapping, level=1)
             df["Memory Usage (GB)"] = df["Memory Usage (MB)"] / 1024
             df = df.drop(columns=["Memory Usage (MB)"])
             best_results = df.groupby(idx_col).min()

experiments/paper_demo.py

@@ -14,7 +14,7 @@
 from experiments.util.models import prefix_in_pairs
 
 estimators = ["time", "memory"]
-estimators = ["memory"]
+# estimators = ["memory"]
 # estimators = ["time"]
 
 # improvements can be either speedups or savings based on context
@@ -27,70 +27,68 @@
 
 # repeat the experiment multiple times (generates multiple files to be aggregated by `get_best_results`)
 n_repeat = 5
+batchnorm_eval = True # BatchNorm in eval mode
 
-# CONV
+# ============== CONV CONFIG ==============
 # Valid choices for models are in models.conv_model_fns
+# models = [
+#     "deepmodel",
+#     "resnet101",
+#     "resnet18",
+#     "vgg16",  # "convnext_base",
+#     "fasterrcnn_resnet50_fpn_v2",
+#     "ssdlite320_mobilenet_v3_large",  # "retinanet_resnet50_fpn_v2",
+#     "deeplabv3_resnet101",
+#     "fcn_resnet101",
+#     "efficientnet_v2_l",
+#     "mobilenet_v3_large",
+#     "resnext101_64x4d",
+# ]
+# models = prefix_in_pairs("memsave_", models)
+# # models = ["resnet101", "memsave_resnet101_conv", "memsave_resnet101_conv+relu+bn", "memsave_resnet101_conv_full"]
+# batch_size = 64
+# input_channels = 3
+# input_HW = 224
+# num_classes = 1000
+# device = "cuda"
+# architecture = "conv"
+# cases = collect_results.select_cases(['All', 'Input', 'Conv', 'Norm'])
+
+# ============== TRANSFORMER CONFIG ==============
+# Valid choices for models are in models.transformer_model_fns
 models = [
-    "deepmodel",
-    "resnet101",
-    "resnet18",
-    "vgg16",  # "convnext_base",
-    "fasterrcnn_resnet50_fpn_v2",
-    "ssdlite320_mobilenet_v3_large",  # "retinanet_resnet50_fpn_v2",
-    "deeplabv3_resnet101",
-    "fcn_resnet101",
-    "efficientnet_v2_l",
-    "mobilenet_v3_large",
-    "resnext101_64x4d",
+    "transformer",
+    "gpt2",
+    "bert",
+    "bart",
+    "roberta",
+    "t5",
+    "flan-t5",
+    # "xlm-roberta",
+    "mistral-7b",
+    "llama3-8b",
+    "phi3-4b",
 ]
-
-# models = ["resnet101", "memsave_resnet101_conv", "memsave_resnet101_conv+relu+bn", "memsave_resnet101_conv_full"]
-# models = ["resnet101", "memsave_resnet101_conv_full"]
-
 models = prefix_in_pairs("memsave_", models)
-# models = ["memsave_resnet101"]
 batch_size = 64
-input_channels = 3
-input_HW = 224
-num_classes = 1000
+input_channels = 2048
+input_HW = 256
+num_classes = 5000
 device = "cuda"
-architecture = "conv"
+architecture = "transformer"
+cases = collect_results.select_cases(["All", "Input", "Norm"])
 
-# LINEAR
+# ============== LINEAR CONFIG ==============
 # Valid choices for models are in models.linear_model_fns
-# models = ['deeplinearmodel']
+# models = ["deeplinearmodel"]
 # models += [f"memsave_{m}" for m in models]  # add memsave versions for each model
-# batch_size = 32768
-# input_channels = 3
-# input_HW = 64
+# batch_size = 1024
+# input_channels = 3  # ignored
+# input_HW = 64  # square of this is passed in estimate.py
 # num_classes = 1000
-# device = 'cuda'
-# architecture = 'linear' # use high batch size
-
-cases = [
-    None,  # ALL
-    [  # INPUT
-        "grad_input",
-        "no_grad_conv_weights",
-        "no_grad_conv_bias",
-        "no_grad_linear_weights",
-        "no_grad_linear_bias",
-        "no_grad_norm_weights",
-        "no_grad_norm_bias",
-    ],
-    [  # CONV
-        "no_grad_linear_weights",
-        "no_grad_linear_bias",
-        "no_grad_norm_weights",
-        "no_grad_norm_bias",
-    ],
-    [  # NORM
-        "no_grad_conv_weights",
-        "no_grad_conv_bias",
-        "no_grad_linear_weights",
-        "no_grad_linear_bias",
-    ],
-]
+# device = "cuda"
+# architecture = "linear"  # use high batch size
+# cases = collect_results.select_cases(["All", "Input", "Linear"])
 
 
 if __name__ == "__main__":
@@ -108,19 +106,30 @@
             cases,
             "results",
         )
+        bn_eval_str = "--bn_eval" if batchnorm_eval else ""
 
         for model in models:
+            B = batch_size
+            if model in prefix_in_pairs("memsave_", ["flan-t5"]):
+                B = 56
+            if model in prefix_in_pairs("memsave_", ["mistral-7b", "phi3-4b"]):
+                B = 16
+            if model in prefix_in_pairs("memsave_", ["llama3-8b"]):
+                B = 8
             for estimate in estimators:
                 outputs = []
 
                 collector.clear_file(estimate)
                 for case in cases:
                     pbar.update()
-                    pbar.set_description(f"{model} {estimate} case {case}")
+                    case_display = collect_results.case_inv_mapping[
+                        collect_results.make_case_str(case)
+                    ]
                     case_str = f"--case {' '.join(case)}" if case is not None else ""
+                    pbar.set_description(f"{model} {estimate} case {case_display}")
                     cmd = (
                         f"python experiments/util/estimate.py --architecture {architecture} --model {model} --estimate {estimate} {case_str} "
-                        + f"--device {device} -B {batch_size} -C_in {input_channels} -HW {input_HW} -n_class {num_classes}"
+                        + f"--device {device} -B {B} -C_in {input_channels} -HW {input_HW} -n_class {num_classes} {bn_eval_str}"
                     )
                     proc = subprocess.run(shlex.split(cmd), capture_output=True)
                     assert proc.stderr in [

experiments/resnet_best_results_to_latex.py

@@ -0,0 +1,39 @@
+"""Simple script to make a latex table from resnet results"""
+
+import pandas as pd
+
+df = pd.read_csv("results/resnet101_only/best_results-conv-cpu-usage_stats.csv")
+df = df.set_index("model")
+df = df.drop(columns=["Scaled M", "Scaled T"])
+df = df.drop("memsave_resnet101_conv+relu+bn")
+df = df[df["case"] != "SurgicalLast"]
+df = df[df["case"] != "Conv"]
+
+mem_div = df[df["case"] == "All"].loc["resnet101", "Memory Usage (GB)"]
+time_div = df[df["case"] == "All"].loc["resnet101", "Time Taken (s)"]
+df["Scaled M"] = df["Memory Usage (GB)"] / mem_div
+df["Scaled T"] = df["Time Taken (s)"] / time_div
+
+df["Memory [GiB]"] = df.apply(
+    lambda x: f"{x['Memory Usage (GB)']:.2f} ({x['Scaled M']:.2f})", axis=1
+)
+df["Time [s]"] = df.apply(
+    lambda x: f"{x['Time Taken (s)']:.2f} ({x['Scaled T']:.2f})", axis=1
+)
+
+df = df.drop(columns=["Scaled M", "Scaled T", "Memory Usage (GB)", "Time Taken (s)"])
+df_p = df.pivot_table(
+    index="model", columns="case", values=df.columns[1:], aggfunc=lambda x: x
+)
+
+labels = {
+    "resnet101": "Default ResNet-101",
+    "memsave_resnet101_conv": "+ swap Convolution",
+    "memsave_resnet101_conv_full": "+ swap BatchNorm, ReLU",
+}
+
+df_p = df_p.rename(index=labels)
+df_p = df_p.sort_index(ascending=False)
+
+print(df_p["Memory [GiB]"].to_latex())
+print(df_p["Time [s]"].to_latex())