Update files

app.py

@@ -7,11 +7,15 @@ import gradio as gr
 
 DATASETS = []
 BENCHMARKS = {
+    # Name: (device, AMP, compile)
     "Parameters": (None, None, None),
-    "CPU Rate with compile": ("cpu", False, True),
-    "CPU Rate without compile": ("cpu", False, False),
-    "CUDA Rate with compile": ("cuda", False, True),
-    "CUDA Rate AMP with compile": ("cuda", True, True),
+    "GPU Memory": (None, None, None),
+    "CPU rate": ("cpu", False, False),
+    "CPU rate with compile": ("cpu", False, True),
+    "CPU rate AMP with compile": ("cpu", True, True),
+    "CUDA rate": ("cuda", False, False),
+    "CUDA rate with compile": ("cuda", False, True),
+    "CUDA rate AMP with compile": ("cuda", True, True),
 }
 
 
@@ -35,11 +39,53 @@ def plot_acc_param(param_compare_results_df: pl.DataFrame, width: int = 1000, he
     return chart.properties(title="Accuray vs Parameter Count", width=width, height=height).configure_scale(zero=False)
 
 
+def plot_acc_memory(memory_compare_results_df: pl.DataFrame, width: int = 900, height: int = 640) -> alt.LayerChart:
+    df = memory_compare_results_df.select(
+        "Model name",
+        "Model type",
+        "Accuracy",
+        "Top-3 accuracy",
+        "Resolution",
+        "Peak GPU memory (MB)",
+        "Parameters (M)",
+        "Pareto frontier (mem)",
+    )
+    base = df.plot.point(
+        x="Peak GPU memory (MB)",
+        y="Accuracy",
+        color="Model type",
+        shape="Resolution:N",
+        tooltip=[
+            "Peak GPU memory (MB)",
+            "Parameters (M)",
+            "Accuracy",
+            "Top-3 accuracy",
+            "Model name",
+            "Model type",
+            "Resolution",
+        ],
+    )
+    text = base.mark_text(align="center", baseline="middle", dy=-10).encode(text="Model name")
+    frontier = df.plot.line(x="Peak GPU memory (MB)", y="Pareto frontier (mem)").mark_line(
+        interpolate="step-after", color="red", strokeWidth=0.3, strokeDash=(2, 2)
+    )
+
+    chart = base + text + frontier
+    return chart.properties(title="Accuray vs GPU Memory", width=width, height=height).configure_scale(zero=False)
+
+
 def plot_acc_rate(rate_compare_results_df: pl.DataFrame, width: int = 1000, height: int = 680) -> alt.LayerChart:
-    device = rate_compare_results_df["device"][0]
-    compiled = rate_compare_results_df["compile"][0]
-    batch_size = rate_compare_results_df["batch_size"][0]
-    amp = rate_compare_results_df["amp"][0]
+    if len(rate_compare_results_df) > 0:
+        device = rate_compare_results_df["device"][0]
+        compiled = rate_compare_results_df["compile"][0]
+        batch_size = rate_compare_results_df["batch_size"][0]
+        amp = rate_compare_results_df["amp"][0]
+    else:
+        device = ""
+        compiled = ""
+        batch_size = ""
+        amp = ""
+
     df = rate_compare_results_df.select(
         "Model name",
         "Model type",
@@ -100,7 +146,9 @@ def update_data(
         param_compare_results_df = param_compare_results_df.with_columns(
             pl.col("Accuracy").cum_max().alias("Pareto frontier (p)")
         )
-        param_compare_results_df = param_compare_results_df.drop("Samples / sec", "device", "ms / sample")
+        param_compare_results_df = param_compare_results_df.drop(
+            "Samples / sec", "device", "ms / sample", "Peak GPU memory (MB)"
+        )
         chart = plot_acc_param(param_compare_results_df)
 
         x_max = param_compare_results_df["Parameters (M)"].quantile(0.9)
@@ -108,18 +156,40 @@ def update_data(
         chart.layer[0].encoding.x.scale = alt.Scale(domain=[x_min, x_max], type=x_scale_type)
         output_df = param_compare_results_df
 
+    # Peak memory
+    elif benchmark == "GPU Memory":
+        memory_compare_results_df = compare_results_df.drop_nulls(subset=["Peak GPU memory (MB)"])
+        memory_compare_results_df = memory_compare_results_df.unique(subset=["Model name"]).sort(
+            "Peak GPU memory (MB)", descending=False
+        )
+        memory_compare_results_df = memory_compare_results_df.with_columns(
+            pl.col("Accuracy").cum_max().alias("Pareto frontier (mem)")
+        )
+        memory_compare_results_df = memory_compare_results_df.drop("Samples / sec", "device", "ms / sample")
+        chart = plot_acc_memory(memory_compare_results_df)
+        x_max = memory_compare_results_df["Peak GPU memory (MB)"].quantile(0.9)
+        x_min = memory_compare_results_df["Peak GPU memory (MB)"].quantile(0.1)
+        chart.layer[0].encoding.x.scale = alt.Scale(domain=[x_min, x_max], type=x_scale_type)
+        output_df = memory_compare_results_df
+
     # Rate
     else:
         (device, amp_enabled, compiled) = BENCHMARKS[benchmark]
-        df = compare_results_df.filter(device=device, amp=amp_enabled, compile=compiled)
+        df = compare_results_df.drop_nulls(subset=["ms / sample"])
+        df = df.filter(device=device, amp=amp_enabled, compile=compiled)
         device_compare_results_df = df.unique(subset=["Model name"]).sort("ms / sample", descending=False)
+        device_compare_results_df = device_compare_results_df.drop("Peak GPU memory (MB)")
         device_compare_results_df = device_compare_results_df.with_columns(
             pl.col("Accuracy").cum_max().alias("Pareto frontier (ms)")
         )
         chart = plot_acc_rate(device_compare_results_df)
 
-        x_max = device_compare_results_df["ms / sample"].quantile(0.95) * 1.04
-        x_min = device_compare_results_df["ms / sample"].min() * 0.96
+        x_max = device_compare_results_df["ms / sample"].quantile(0.95)
+        x_min = device_compare_results_df["ms / sample"].min()
+        if x_max is not None and x_min is not None:
+            x_max = x_max * 1.04
+            x_min = x_min * 0.96
+
         chart.layer[0].encoding.x.scale = alt.Scale(domain=[x_min, x_max], type=x_scale_type)
         output_df = device_compare_results_df