From c20873a44c7a316e15f4780547887f37d7f20cce Mon Sep 17 00:00:00 2001 From: Richard Date: Mon, 13 Jul 2026 11:52:53 +0800 Subject: [PATCH 1/4] Visualization Fixes - centralizing graphing utils for the nth time - make sure INT4 is visualized --- src/QuantAdv.py | 26 +- src/__init__.py | 2 +- src/attack.py | 18 +- src/combine.py | 782 +----------------------------------------------- src/data.py | 335 +++++++++------------ 5 files changed, 171 insertions(+), 992 deletions(-) diff --git a/src/QuantAdv.py b/src/QuantAdv.py index 6c46796..9da27d4 100644 --- a/src/QuantAdv.py +++ b/src/QuantAdv.py @@ -717,7 +717,7 @@ def run_quant_component_ablation(model, loader, name, eps=DEFAULT_EPS): The rows distinguish three interpretations of "testing quantized": weight-only quantization, activation-only quantization, and both together. For each interpretation the experiment reports ordinary hard-round PGD and - a budget-matched STE/BPDA PGD. A large hard-PGD vs STE-PGD gap is evidence + a budget-matched BPDA-PGD. A large PGD vs BPDA-PGD gap is evidence of gradient masking, not evidence that quantization itself improves robustness. @@ -725,7 +725,7 @@ def run_quant_component_ablation(model, loader, name, eps=DEFAULT_EPS): doesn't touch the gradient path back to the input (the conv/linear op is still linear in ``x`` for whatever rounded weight value it has), so ``weight_only`` should show ``frac_zero_grad_hard`` near the background - rate and a small hard-vs-STE gap. Only ``act_only``/``both`` quantize a + rate and a small PGD-vs-BPDA gap. Only ``act_only``/``both`` quantize a tensor that ``x`` actually flows through, so only those can mask. """ configs = [ @@ -757,18 +757,6 @@ def run_quant_component_ablation(model, loader, name, eps=DEFAULT_EPS): g_hard = torch.autograd.grad(loss, x_in)[0].flatten() frac_zero = (g_hard.abs() < GRAD_ZERO_THRESHOLD).float().mean().item() - with ste_mode(model, True): - torch.manual_seed(0) - pgd_ste = make_torchattack( - torchattacks.PGD, - model, - eps=eps, - alpha=PGD_ALPHA, - steps=PGD_STEPS, - random_start=PGD_RANDOM_START, - ) - pgd_ste_acc = accuracy_under_attack(model, loader, pgd_ste) - bpda = run_bpda( model, loader, @@ -785,15 +773,9 @@ def run_quant_component_ablation(model, loader, name, eps=DEFAULT_EPS): "clean_acc": clean_acc, "PGD_acc": pgd_hard_acc, "PGD_hard_acc": pgd_hard_acc, - "PGD_ste_acc": pgd_ste_acc, "BPDA_acc": bpda["BPDA_PGD"], "BPDA_mean": bpda["BPDA_PGD_mean"], "BPDA_std": bpda["BPDA_PGD_std"], - "PGD_minus_STE": ( - pgd_hard_acc - pgd_ste_acc - if pgd_hard_acc is not None and pgd_ste_acc is not None - else None - ), "PGD_minus_BPDA": ( pgd_hard_acc - bpda["BPDA_PGD"] if pgd_hard_acc is not None and bpda["BPDA_PGD"] is not None @@ -1829,9 +1811,7 @@ def save_epsilon_sweep(): def build_reports(): """Combine result files and generate final report artifacts.""" - tables = report_data.combine_all(report_data.DATA_DIR) - report_data.plot_all(tables, report_data.DATA_DIR) - report_data.print_report(tables) + report_data.generate_reports(report_data.DATA_DIR) safe_call(build_reports, "report generation failed", show_traceback=True) print("All done.") diff --git a/src/__init__.py b/src/__init__.py index a7f4c9e..3a9c35a 100644 --- a/src/__init__.py +++ b/src/__init__.py @@ -1,3 +1,3 @@ """ -QuantAdv/src: Active centralized development for ./src_old +QuantAdv/src: Active centralized development """ diff --git a/src/attack.py b/src/attack.py index aac46fe..ff17d8a 100644 --- a/src/attack.py +++ b/src/attack.py @@ -1181,13 +1181,9 @@ def run_epsilon_sweep_for_model( if it tracks ``Random_Noise_acc`` instead of dropping as epsilon grows, is showing you gradient masking, not robustness. - ``PGD_ste_acc`` is the same sweep with the masking bypassed (single - restart, matching ``PGD_acc``'s budget), so you can read the masking - gap directly off this table without cross-referencing ``BPDA_acc``. - ``BPDA_acc`` is kept for compatibility and now uses a restart budget - (``n_restarts=1``) matched to ``PGD_acc``/``PGD_ste_acc`` rather than - the previous mismatched default, so any gap you see is attributable to - the gradient regime and not to a bigger attack budget. + ``BPDA_acc`` uses a restart budget (``n_restarts=1``) matched to + ``PGD_acc``. The sweep therefore reports only vanilla PGD and BPDA-PGD; + it does not execute or report a separate PGD-STE attack. """ if safe_set is None: @@ -1225,14 +1221,6 @@ def run_pgd_sweep(use_ste): "PGD (hard-round) sweep failed", context=context, ) - if is_quant: - safe_set( - row, - "PGD_ste_acc", - lambda: run_pgd_sweep(True), - "PGD (STE) sweep failed", - context=context, - ) safe_set( row, "Random_Noise_acc", diff --git a/src/combine.py b/src/combine.py index ec541af..5c55c76 100644 --- a/src/combine.py +++ b/src/combine.py @@ -1,784 +1,34 @@ #!/usr/bin/env python3 -""" -Standalone result combiner for QuantAdv experiments. - -Reads individual per-model CSV/JSON files from DATA_DIR, writes combined CSVs, -and regenerates the summary PNG plots without loading models or running attacks. +"""Recover partial QuantAdv outputs and regenerate report artifacts. """ from __future__ import annotations import argparse -import json -import math -import os -import re from pathlib import Path -from typing import Iterable - -import numpy as np -import pandas as pd - -import matplotlib - -matplotlib.use("Agg") -import matplotlib.pyplot as plt -import seaborn as sns - -try: - from . import data as shared_data -except ImportError: - import data as shared_data - - -try: - import config as _config -except Exception: - _config = None - - -PROJECT_ROOT = Path(__file__).resolve().parent -DATA_DIR = Path(getattr(_config, "DATA_DIR", PROJECT_ROOT / "data")) - -RESULTS_CSV = Path(getattr(_config, "RESULTS_CSV", DATA_DIR / "accuracyresult.csv")) -SWEEP_CSV = Path(getattr(_config, "SWEEP_CSV", DATA_DIR / "sweepresult.csv")) -PLOT_PNG = Path(getattr(_config, "PLOT_PNG", DATA_DIR / "accuracyplot.png")) - -SWEEP_PLOT_PNG = Path( - getattr(_config, "SWEEP_PLOT_PNG", DATA_DIR / "epsilon_sweep.png") -) -ABLATION_PLOT_PNG = Path( - getattr(_config, "ABLATION_PLOT_PNG", DATA_DIR / "pgd_steps_ablation.png") -) -TRAJECTORY_PLOT_PNG = Path( - getattr(_config, "TRAJECTORY_PLOT_PNG", DATA_DIR / "pgd_trajectory.png") -) -LAYERWISE_PLOT_PNG = Path( - getattr(_config, "LAYERWISE_PLOT_PNG", DATA_DIR / "layerwise_grad_profile.png") -) -COMPONENT_ABLATION_PLOT_PNG = Path( - getattr(_config, "COMPONENT_ABLATION_PLOT_PNG", DATA_DIR / "component_ablation.png") -) -MASKING_SUMMARY_PLOT_PNG = Path( - getattr( - _config, "MASKING_SUMMARY_PLOT_PNG", DATA_DIR / "gradient_masking_summary.png" - ) -) -MARGIN_PLOT_PNG = Path( - getattr(_config, "MARGIN_PLOT_PNG", DATA_DIR / "confidence_margin.png") -) -HEATMAP_PLOT_PNG = Path( - getattr(_config, "HEATMAP_PLOT_PNG", DATA_DIR / "results_heatmap.png") -) - -ABLATION_COMBINED_CSV = DATA_DIR / "ablation_combined.csv" -LAYERWISE_COMBINED_CSV = DATA_DIR / "layerwise_combined.csv" -TRAJECTORY_COMBINED_CSV = DATA_DIR / "trajectory_combined.csv" -COMPONENT_ABLATION_COMBINED_CSV = DATA_DIR / "component_ablation_combined.csv" -MARGIN_COMBINED_CSV = DATA_DIR / "margin_combined.csv" - -PLOT_MAX_ACCURACY = getattr(_config, "PLOT_MAX_ACCURACY", 1.0) -PLOT_DPI = getattr(_config, "PLOT_DPI", 200) -PLOT_BBOX_INCHES = getattr(_config, "PLOT_BBOX_INCHES", "tight") -PLOT_GRID_ALPHA = getattr(_config, "PLOT_GRID_ALPHA", 0.35) -PLOT_LEGEND_FONT_SIZE = getattr(_config, "PLOT_LEGEND_FONT_SIZE", 8) - -SUMMARY_PLOT_FIGSIZE = getattr(_config, "SUMMARY_PLOT_FIGSIZE", (14, 7)) -SUMMARY_XTICK_ROTATION = getattr(_config, "SUMMARY_XTICK_ROTATION", 45) -SUMMARY_GRID_ALPHA = getattr(_config, "SUMMARY_GRID_ALPHA", 0.35) - -SWEEP_PLOT_COLS_MAX = getattr(_config, "SWEEP_PLOT_COLS_MAX", 3) -SWEEP_PLOT_WIDTH = getattr(_config, "SWEEP_PLOT_WIDTH", 6) -SWEEP_PLOT_HEIGHT = getattr(_config, "SWEEP_PLOT_HEIGHT", 4) - -ABLATION_FIGSIZE = getattr(_config, "ABLATION_FIGSIZE", (12, 7)) -TRAJECTORY_FIGSIZE = getattr(_config, "TRAJECTORY_FIGSIZE", (14, 5)) - -LAYERWISE_PLOT_COLS_MAX = getattr(_config, "LAYERWISE_PLOT_COLS_MAX", 3) -LAYERWISE_PLOT_WIDTH = getattr(_config, "LAYERWISE_PLOT_WIDTH", 6) -LAYERWISE_PLOT_HEIGHT = getattr(_config, "LAYERWISE_PLOT_HEIGHT", 4) -LAYERWISE_XTICK_ROTATION = getattr(_config, "LAYERWISE_XTICK_ROTATION", 90) -LAYERWISE_XTICK_FONT_SIZE = getattr(_config, "LAYERWISE_XTICK_FONT_SIZE", 6) - -COMPONENT_ABLATION_COL_WRAP = getattr(_config, "COMPONENT_ABLATION_COL_WRAP", 3) -COMPONENT_ABLATION_HEIGHT = getattr(_config, "COMPONENT_ABLATION_HEIGHT", 4) - -MASKING_SUMMARY_FIGSIZE = getattr(_config, "MASKING_SUMMARY_FIGSIZE", (13, 5)) -MASKING_BASELINE_LINEWIDTH = getattr(_config, "MASKING_BASELINE_LINEWIDTH", 1.0) -MASKING_SCATTER_SIZE = getattr(_config, "MASKING_SCATTER_SIZE", 60) - -MARGIN_PLOT_COLS_MAX = getattr(_config, "MARGIN_PLOT_COLS_MAX", 3) -MARGIN_PLOT_WIDTH = getattr(_config, "MARGIN_PLOT_WIDTH", 5) -MARGIN_PLOT_HEIGHT = getattr(_config, "MARGIN_PLOT_HEIGHT", 4) -MARGIN_HIST_BINS = getattr(_config, "MARGIN_HIST_BINS", 30) -MARGIN_HIST_ALPHA = getattr(_config, "MARGIN_HIST_ALPHA", 0.55) - -HEATMAP_MIN_WIDTH = getattr(_config, "HEATMAP_MIN_WIDTH", 12) -HEATMAP_MIN_HEIGHT = getattr(_config, "HEATMAP_MIN_HEIGHT", 6) -HEATMAP_ROW_HEIGHT = getattr(_config, "HEATMAP_ROW_HEIGHT", 0.45) -HEATMAP_VMIN = getattr(_config, "HEATMAP_VMIN", 0.0) -HEATMAP_VMAX = getattr(_config, "HEATMAP_VMAX", 1.0) -HEATMAP_LINEWIDTHS = getattr(_config, "HEATMAP_LINEWIDTHS", 0.5) - - -def _model_from_prefixed_path(path: Path, prefix: str, suffix: str) -> str: - """Extract the model name encoded in a result filename.""" - name = path.name - if name.startswith(prefix): - name = name[len(prefix) :] - if name.endswith(suffix): - name = name[: -len(suffix)] - return name - - -def _read_csvs(paths: Iterable[Path], model_prefix: str | None = None) -> pd.DataFrame: - """Read and concatenate matching CSV files with model metadata attached.""" - frames = [] - for path in sorted(paths): - try: - df = pd.read_csv(path) - except Exception as exc: - print(f"[WARN] could not read {path}: {exc}") - continue - - if model_prefix and "model" not in df.columns: - df.insert(0, "model", _model_from_prefixed_path(path, model_prefix, ".csv")) - - frames.append(df) - - if not frames: - return pd.DataFrame() - - out = pd.concat(frames, ignore_index=True, sort=False) - if "model" in out.columns: - out = out.drop_duplicates( - subset=[ - c - for c in ["model", "epsilon", "steps", "layer", "config"] - if c in out.columns - ], - keep="last", - ) - return out - - -def _write(df: pd.DataFrame, path: Path) -> pd.DataFrame: - """Create parent directories and write a dataframe to CSV.""" - if df is not None and not df.empty: - path.parent.mkdir(parents=True, exist_ok=True) - df.to_csv(path, index=False) - print(f"wrote {path} ({len(df)} rows)") - return df - - -def _result_files(data_dir: Path) -> list[Path]: - """Collect result-file paths by experiment family.""" - blocked = { - RESULTS_CSV.name, - SWEEP_CSV.name, - ABLATION_COMBINED_CSV.name, - LAYERWISE_COMBINED_CSV.name, - TRAJECTORY_COMBINED_CSV.name, - COMPONENT_ABLATION_COMBINED_CSV.name, - MARGIN_COMBINED_CSV.name, - "defense_summary.csv", - } - files = [] - for path in data_dir.glob("results_*.csv"): - if path.name in blocked: - continue - if path.name.startswith("results_sweep_"): - continue - files.append(path) - return files - - -def combine_scalar_results(data_dir: Path, output: Path = RESULTS_CSV) -> pd.DataFrame: - """Combine per-model scalar result CSV files into one table.""" - files = _result_files(data_dir) - df = _read_csvs(files, model_prefix="results_") - return _write(df, output) - - -def combine_sweeps(data_dir: Path, output: Path = SWEEP_CSV) -> pd.DataFrame: - """Combine epsilon-sweep CSV files into one table.""" - files = [p for p in data_dir.glob("sweep_*.csv") if p.name != output.name] - df = _read_csvs(files, model_prefix="sweep_") - if not df.empty and {"model", "epsilon"}.issubset(df.columns): - df = df.sort_values(["model", "epsilon"]).reset_index(drop=True) - return _write(df, output) - - -def combine_ablation( - data_dir: Path, output: Path = ABLATION_COMBINED_CSV -) -> pd.DataFrame: - """Combine PGD step-ablation CSV files into one table.""" - files = [p for p in data_dir.glob("ablation_*.csv") if p.name != output.name] - df = _read_csvs(files, model_prefix="ablation_") - if not df.empty and {"model", "steps"}.issubset(df.columns): - df = df.sort_values(["model", "steps"]).reset_index(drop=True) - return _write(df, output) - - -def combine_layerwise( - data_dir: Path, output: Path = LAYERWISE_COMBINED_CSV -) -> pd.DataFrame: - """Combine layerwise gradient-profile CSV files into one table.""" - files = [p for p in data_dir.glob("layerwise_*.csv") if p.name != output.name] - df = _read_csvs(files, model_prefix="layerwise_") - return _write(df, output) - - -def combine_component_ablation( - data_dir: Path, output: Path = COMPONENT_ABLATION_COMBINED_CSV -) -> pd.DataFrame: - """Combine quantization component-ablation CSV files into one table.""" - files = [ - p for p in data_dir.glob("component_ablation_*.csv") if p.name != output.name - ] - df = _read_csvs(files, model_prefix="component_ablation_") - return _write(df, output) - - -def combine_trajectories( - data_dir: Path, output: Path = TRAJECTORY_COMBINED_CSV -) -> pd.DataFrame: - """Combine PGD trajectory diagnostic CSV files into one table.""" - rows = [] - for path in sorted(data_dir.glob("trajectory_*.json")): - model = _model_from_prefixed_path(path, "trajectory_", ".json") - try: - with open(path, "r", encoding="utf-8") as f: - traj = json.load(f) - except Exception as exc: - print(f"[WARN] could not read {path}: {exc}") - continue - - grad = traj.get("grad_norm_per_step", []) or [] - move = traj.get("movement_from_random_start_per_step", []) or [] - n = max(len(grad), len(move)) - for i in range(n): - rows.append( - { - "model": model, - "step": i + 1, - "grad_norm_per_step": grad[i] if i < len(grad) else np.nan, - "movement_from_random_start_per_step": ( - move[i] if i < len(move) else np.nan - ), - } - ) - - return _write(pd.DataFrame(rows), output) - - -def combine_margins(data_dir: Path, output: Path = MARGIN_COMBINED_CSV) -> pd.DataFrame: - """Combine confidence-margin diagnostic CSV files into one table.""" - rows = [] - for path in sorted(data_dir.glob("margin_*.json")): - model = _model_from_prefixed_path(path, "margin_", ".json") - try: - with open(path, "r", encoding="utf-8") as f: - margins = json.load(f) - except Exception as exc: - print(f"[WARN] could not read {path}: {exc}") - continue - - for kind, values in ( - ("clean", margins.get("clean_margins", [])), - ("adv", margins.get("adv_margins", [])), - ): - for i, value in enumerate(values): - rows.append({"model": model, "kind": kind, "index": i, "margin": value}) - - return _write(pd.DataFrame(rows), output) - - -def _model_names(*dfs: pd.DataFrame) -> list[str]: - """Return the preferred display order for model names.""" - names = [] - for df in dfs: - if df is not None and not df.empty and "model" in df.columns: - names.extend(df["model"].dropna().astype(str).tolist()) - return list(dict.fromkeys(names)) - - -def plot_summary_results(df_results: pd.DataFrame, output: Path = PLOT_PNG) -> None: - """Plot the headline accuracy and robustness summary metrics.""" - if df_results is None or df_results.empty or "model" not in df_results.columns: - return - - acc_cols = [ - c - for c in [ - "clean_acc", - "FGSM", - "PGD", - "CW", - "DeepFool", - "JSMA", - "AutoAttack", - "Transfer_from_FP32", - "MIM_Transfer", - "UAP_Transfer", - "Surrogate_Transfer", - "Random_Noise", - "BPDA_PGD", - "NES", - "Boundary_acc", - ] - if c in df_results.columns and df_results[c].notna().any() - ] - - if not acc_cols: - return - - df_plot = df_results.melt( - id_vars="model", value_vars=acc_cols, var_name="Attack", value_name="Accuracy" - ) - df_plot = df_plot.dropna(subset=["Accuracy"]) - if df_plot.empty: - return - - plt.figure(figsize=SUMMARY_PLOT_FIGSIZE) - sns.barplot(data=df_plot, x="model", y="Accuracy", hue="Attack") - plt.xticks(rotation=SUMMARY_XTICK_ROTATION, ha="right") - plt.title("Model Accuracy under Various Adversarial Attacks") - plt.ylim(0, PLOT_MAX_ACCURACY) - plt.grid(axis="y", linestyle="--", alpha=SUMMARY_GRID_ALPHA) - plt.tight_layout() - plt.savefig(output, dpi=PLOT_DPI, bbox_inches=PLOT_BBOX_INCHES) - plt.close() - print(f"wrote {output}") +import data as report_data -def plot_epsilon_sweep_curves( - df_sweep: pd.DataFrame, output: Path = SWEEP_PLOT_PNG -) -> None: - """Plot robustness metrics as a function of attack epsilon.""" - if df_sweep is None or df_sweep.empty: - return - value_cols = [ - c for c in ["PGD_acc", "Random_Noise_acc", "BPDA_acc"] if c in df_sweep.columns - ] - if not value_cols: - return - df_long = df_sweep.melt( - id_vars=["model", "epsilon"], - value_vars=value_cols, - var_name="Attack", - value_name="Accuracy", - ) - df_long = df_long.dropna(subset=["Accuracy"]) - if df_long.empty: - return - - models = df_long["model"].dropna().astype(str).unique() - cols = min(SWEEP_PLOT_COLS_MAX, len(models)) - rows = int(math.ceil(len(models) / cols)) - fig, axes = plt.subplots( - rows, - cols, - figsize=(SWEEP_PLOT_WIDTH * cols, SWEEP_PLOT_HEIGHT * rows), - squeeze=False, - ) - - for i, model in enumerate(models): - ax = axes[i // cols][i % cols] - sns.lineplot( - data=df_long[df_long["model"] == model], - x="epsilon", - y="Accuracy", - hue="Attack", - marker="o", - ax=ax, - ) - ax.set_title(model) - ax.set_ylim(0, PLOT_MAX_ACCURACY) - ax.grid(linestyle="--", alpha=PLOT_GRID_ALPHA) - - for j in range(len(models), rows * cols): - axes[j // cols][j % cols].axis("off") - - fig.suptitle("Accuracy vs Perturbation Budget (Epsilon Sweep)") - fig.tight_layout() - fig.savefig(output, dpi=PLOT_DPI, bbox_inches=PLOT_BBOX_INCHES) - plt.close(fig) - print(f"wrote {output}") - - -def plot_pgd_steps_ablation( - df_ablation: pd.DataFrame, output: Path = ABLATION_PLOT_PNG -) -> None: - """Plot accuracy as PGD step count changes.""" - if ( - df_ablation is None - or df_ablation.empty - or not {"steps", "acc", "model"}.issubset(df_ablation.columns) - ): - return - - plt.figure(figsize=ABLATION_FIGSIZE) - sns.lineplot(data=df_ablation, x="steps", y="acc", hue="model", marker="o") - plt.title("PGD Accuracy vs Number of Steps (Gradient Masking Check)") - plt.xlabel("PGD steps") - plt.ylabel("Accuracy") - plt.ylim(0, PLOT_MAX_ACCURACY) - plt.grid(linestyle="--", alpha=PLOT_GRID_ALPHA) - plt.tight_layout() - plt.savefig(output, dpi=PLOT_DPI, bbox_inches=PLOT_BBOX_INCHES) - plt.close() - print(f"wrote {output}") - - -def plot_pgd_trajectory( - df_traj: pd.DataFrame, output: Path = TRAJECTORY_PLOT_PNG -) -> None: - """Plot gradient norms and movement along PGD trajectories.""" - required = { - "model", - "step", - "grad_norm_per_step", - "movement_from_random_start_per_step", - } - if df_traj is None or df_traj.empty or not required.issubset(df_traj.columns): - return - - fig, axes = plt.subplots(1, 2, figsize=TRAJECTORY_FIGSIZE) - for model, group in df_traj.groupby("model", sort=False): - group = group.sort_values("step") - axes[0].plot( - group["step"], group["grad_norm_per_step"], marker="o", label=model - ) - axes[1].plot( - group["step"], - group["movement_from_random_start_per_step"], - marker="o", - label=model, - ) - - axes[0].set_title("Gradient Norm per PGD Step") - axes[0].set_xlabel("Step") - axes[0].set_ylabel("Grad Norm") - axes[0].set_yscale("log") - axes[0].grid(linestyle="--", alpha=PLOT_GRID_ALPHA) - axes[0].legend(fontsize=PLOT_LEGEND_FONT_SIZE) - - axes[1].set_title("Perturbation Movement per PGD Step") - axes[1].set_xlabel("Step") - axes[1].set_ylabel("Linf Movement from Random Start") - axes[1].grid(linestyle="--", alpha=PLOT_GRID_ALPHA) - axes[1].legend(fontsize=PLOT_LEGEND_FONT_SIZE) - - fig.tight_layout() - fig.savefig(output, dpi=PLOT_DPI, bbox_inches=PLOT_BBOX_INCHES) - plt.close(fig) - print(f"wrote {output}") - - -def plot_layerwise_grad_profile( - df_layer: pd.DataFrame, output: Path = LAYERWISE_PLOT_PNG -) -> None: - """Plot per-layer gradient magnitudes for diagnosed models.""" - required = {"model", "layer", "grad_norm_hard", "grad_norm_ste"} - if df_layer is None or df_layer.empty or not required.issubset(df_layer.columns): - return - - models = df_layer["model"].dropna().astype(str).unique() - cols = min(LAYERWISE_PLOT_COLS_MAX, len(models)) - rows = int(math.ceil(len(models) / cols)) - fig, axes = plt.subplots( - rows, - cols, - figsize=(LAYERWISE_PLOT_WIDTH * cols, LAYERWISE_PLOT_HEIGHT * rows), - squeeze=False, - ) - - for i, model in enumerate(models): - df = df_layer[df_layer["model"] == model].reset_index(drop=True) - ax = axes[i // cols][i % cols] - x = np.arange(len(df)) - ax.plot(x, df["grad_norm_hard"], marker="o", label="hard-round") - ax.plot(x, df["grad_norm_ste"], marker="o", label="STE") - ax.set_yscale("log") - ax.set_xticks(x) - ax.set_xticklabels( - df["layer"], - rotation=LAYERWISE_XTICK_ROTATION, - fontsize=LAYERWISE_XTICK_FONT_SIZE, - ) - ax.set_title(model) - ax.set_ylabel("Grad Norm (log)") - ax.legend(fontsize=PLOT_LEGEND_FONT_SIZE) - ax.grid(linestyle="--", alpha=PLOT_GRID_ALPHA) - - for j in range(len(models), rows * cols): - axes[j // cols][j % cols].axis("off") - - fig.suptitle("Layerwise Gradient Norms: Hard-Round vs STE") - fig.tight_layout() - fig.savefig(output, dpi=PLOT_DPI, bbox_inches=PLOT_BBOX_INCHES) - plt.close(fig) - print(f"wrote {output}") - - -def plot_component_ablation( - df_component: pd.DataFrame, output: Path = COMPONENT_ABLATION_PLOT_PNG -) -> None: - """Plot clean, hard-PGD, and BPDA component-ablation accuracies.""" - required = {"model", "config", "clean_acc", "PGD_acc"} - if ( - df_component is None - or df_component.empty - or not required.issubset(df_component.columns) - ): - return - - value_vars = [ - c for c in ["clean_acc", "PGD_acc", "BPDA_acc"] if c in df_component.columns - ] - df_long = df_component.melt( - id_vars=["model", "config"], - value_vars=value_vars, - var_name="Metric", - value_name="Accuracy", - ).dropna(subset=["Accuracy"]) - - if df_long.empty: - return - - g = sns.catplot( - data=df_long, - x="config", - y="Accuracy", - hue="Metric", - col="model", - kind="bar", - col_wrap=COMPONENT_ABLATION_COL_WRAP, - height=COMPONENT_ABLATION_HEIGHT, - sharey=True, - ) - g.set_titles("{col_name}") - g.set(ylim=(0, PLOT_MAX_ACCURACY)) - g.savefig(output, dpi=PLOT_DPI, bbox_inches=PLOT_BBOX_INCHES) - plt.close(g.fig) - print(f"wrote {output}") - - -def plot_gradient_masking_summary( - df_results: pd.DataFrame, output: Path = MASKING_SUMMARY_PLOT_PNG -) -> None: - """Plot metrics that summarize potential gradient masking.""" - if ( - df_results is None - or df_results.empty - or not {"model", "PGD", "AutoAttack"}.issubset(df_results.columns) - ): - return - - df = df_results.dropna(subset=["PGD", "AutoAttack"]).copy() - if df.empty: - return - - df["PGD_minus_AutoAttack"] = df["PGD"] - df["AutoAttack"] - - fig, axes = plt.subplots(1, 2, figsize=MASKING_SUMMARY_FIGSIZE) - sns.barplot(data=df, x="model", y="PGD_minus_AutoAttack", ax=axes[0]) - axes[0].axhline(0, color="black", linewidth=MASKING_BASELINE_LINEWIDTH) - axes[0].set_xticklabels( - axes[0].get_xticklabels(), rotation=SUMMARY_XTICK_ROTATION, ha="right" - ) - axes[0].set_title("PGD - AutoAttack Accuracy Gap") - axes[0].grid(axis="y", linestyle="--", alpha=PLOT_GRID_ALPHA) - - if "frac_zero_grad_hard" in df.columns and df["frac_zero_grad_hard"].notna().any(): - df2 = df.dropna(subset=["frac_zero_grad_hard"]) - sns.scatterplot( - data=df2, - x="frac_zero_grad_hard", - y="PGD_minus_AutoAttack", - hue="model", - s=MASKING_SCATTER_SIZE, - ax=axes[1], - ) - axes[1].set_title("Masking Gap vs Fraction of Zero Gradients") - axes[1].grid(linestyle="--", alpha=PLOT_GRID_ALPHA) - else: - axes[1].axis("off") - - fig.tight_layout() - fig.savefig(output, dpi=PLOT_DPI, bbox_inches=PLOT_BBOX_INCHES) - plt.close(fig) - print(f"wrote {output}") - - -def plot_confidence_margin_diagnostic( - df_margins: pd.DataFrame, output: Path = MARGIN_PLOT_PNG -) -> None: - """Plot clean and adversarial confidence-margin distributions.""" - required = {"model", "kind", "margin"} - if ( - df_margins is None - or df_margins.empty - or not required.issubset(df_margins.columns) - ): - return - - models = df_margins["model"].dropna().astype(str).unique() - cols = min(MARGIN_PLOT_COLS_MAX, len(models)) - rows = int(math.ceil(len(models) / cols)) - fig, axes = plt.subplots( - rows, - cols, - figsize=(MARGIN_PLOT_WIDTH * cols, MARGIN_PLOT_HEIGHT * rows), - squeeze=False, - ) - - for i, model in enumerate(models): - ax = axes[i // cols][i % cols] - group = df_margins[df_margins["model"] == model] - clean = group[group["kind"] == "clean"]["margin"].dropna() - adv = group[group["kind"] == "adv"]["margin"].dropna() - if not clean.empty: - ax.hist( - clean, - bins=MARGIN_HIST_BINS, - alpha=MARGIN_HIST_ALPHA, - label="clean", - density=True, - ) - if not adv.empty: - ax.hist( - adv, - bins=MARGIN_HIST_BINS, - alpha=MARGIN_HIST_ALPHA, - label="PGD-adv", - density=True, - ) - ax.set_title(model) - ax.set_xlabel("Top1 - Top2 Softmax Margin") - ax.legend(fontsize=PLOT_LEGEND_FONT_SIZE) - ax.grid(linestyle="--", alpha=PLOT_GRID_ALPHA) - - for j in range(len(models), rows * cols): - axes[j // cols][j % cols].axis("off") - - fig.suptitle("Confidence Margin: Clean vs PGD-Adversarial") - fig.tight_layout() - fig.savefig(output, dpi=PLOT_DPI, bbox_inches=PLOT_BBOX_INCHES) - plt.close(fig) - print(f"wrote {output}") - - -def plot_results_heatmap( - df_results: pd.DataFrame, output: Path = HEATMAP_PLOT_PNG -) -> None: - """Plot scalar metrics as a model-by-metric heatmap.""" - if df_results is None or df_results.empty or "model" not in df_results.columns: - return - - candidate_cols = [ - "clean_acc", - "FGSM", - "PGD", - "AutoAttack", - "CW", - "DeepFool", - "JSMA", - "Surrogate_Transfer", - "Transfer_from_FP32", - "MIM_Transfer", - "UAP_Transfer", - "Random_Noise", - "BPDA_PGD", - "NES", - "Boundary_acc", - ] - cols = [ - c - for c in candidate_cols - if c in df_results.columns and df_results[c].notna().any() - ] - if not cols: - return - - df_heat = df_results.set_index("model")[cols].astype(float) - plt.figure( - figsize=( - max(HEATMAP_MIN_WIDTH, len(cols)), - max(HEATMAP_MIN_HEIGHT, len(df_heat) * HEATMAP_ROW_HEIGHT), - ) - ) - sns.heatmap( - df_heat, - annot=True, - fmt=".2f", - cmap="RdYlGn", - vmin=HEATMAP_VMIN, - vmax=HEATMAP_VMAX, - linewidths=HEATMAP_LINEWIDTHS, +def parse_args() -> argparse.Namespace: + """Parse recovery-report command-line options.""" + parser = argparse.ArgumentParser( + description="Merge incomplete QuantAdv outputs and regenerate reports." ) - plt.title("Full Results Heatmap: Models vs Attacks") - plt.tight_layout() - plt.savefig(output, dpi=PLOT_DPI, bbox_inches=PLOT_BBOX_INCHES) - plt.close() - print(f"wrote {output}") - - -def combine_all(data_dir: Path) -> dict[str, pd.DataFrame]: - """Combine all available experiment result families.""" - data_dir.mkdir(parents=True, exist_ok=True) - - df_results = combine_scalar_results(data_dir, RESULTS_CSV) - df_sweep = combine_sweeps(data_dir, SWEEP_CSV) - df_ablation = combine_ablation(data_dir, ABLATION_COMBINED_CSV) - df_layer = combine_layerwise(data_dir, LAYERWISE_COMBINED_CSV) - df_component = combine_component_ablation(data_dir, COMPONENT_ABLATION_COMBINED_CSV) - df_traj = combine_trajectories(data_dir, TRAJECTORY_COMBINED_CSV) - df_margins = combine_margins(data_dir, MARGIN_COMBINED_CSV) - - if df_results.empty and RESULTS_CSV.exists(): - df_results = pd.read_csv(RESULTS_CSV) - if df_sweep.empty and SWEEP_CSV.exists(): - df_sweep = pd.read_csv(SWEEP_CSV) - - return { - "results": df_results, - "sweep": df_sweep, - "ablation": df_ablation, - "layerwise": df_layer, - "component_ablation": df_component, - "trajectory": df_traj, - "margins": df_margins, - } - - -def plot_all(dfs: dict[str, pd.DataFrame]) -> None: - """Generate all summary plots from combined result tables.""" - plot_summary_results(dfs["results"], PLOT_PNG) - plot_epsilon_sweep_curves(dfs["sweep"], SWEEP_PLOT_PNG) - plot_pgd_steps_ablation(dfs["ablation"], ABLATION_PLOT_PNG) - plot_pgd_trajectory(dfs["trajectory"], TRAJECTORY_PLOT_PNG) - plot_layerwise_grad_profile(dfs["layerwise"], LAYERWISE_PLOT_PNG) - plot_component_ablation(dfs["component_ablation"], COMPONENT_ABLATION_PLOT_PNG) - plot_gradient_masking_summary(dfs["results"], MASKING_SUMMARY_PLOT_PNG) - plot_confidence_margin_diagnostic(dfs["margins"], MARGIN_PLOT_PNG) - plot_results_heatmap(dfs["results"], HEATMAP_PLOT_PNG) + parser.add_argument("--data-dir", type=Path, default=Path(report_data.DATA_DIR)) + parser.add_argument("--no-plots", action="store_true") + parser.add_argument("--quiet", action="store_true", help="Do not print table sizes.") + return parser.parse_args() def main() -> None: - """Run the command-line entry point for this module.""" - parser = argparse.ArgumentParser( - description="Merge incomplete QuantAdv outputs and regenerate reports." + """Run partial-output recovery using the shared report implementation.""" + args = parse_args() + report_data.generate_reports( + args.data_dir, + plots=not args.no_plots, + summary=not args.quiet, ) - parser.add_argument("--data-dir", type=Path, default=shared_data.DATA_DIR) - parser.add_argument("--no-plots", action="store_true") - args = parser.parse_args() - dfs = shared_data.combine_all(args.data_dir) - if not args.no_plots: - shared_data.plot_all(dfs, args.data_dir) - shared_data.print_report(dfs) if __name__ == "__main__": diff --git a/src/data.py b/src/data.py index 744fda1..9171a91 100644 --- a/src/data.py +++ b/src/data.py @@ -12,8 +12,9 @@ import math import os import re +from dataclasses import dataclass from pathlib import Path -from typing import Iterable +from typing import Callable, Iterable import numpy as np import pandas as pd @@ -83,75 +84,6 @@ def _write(df: pd.DataFrame, path: Path) -> pd.DataFrame: return df -def _result_files(data_dir: Path) -> list[Path]: - """Collect result-file paths by experiment family.""" - blocked = { - _name(RESULTS_CSV), - _name(SWEEP_CSV), - _name(ABLATION_COMBINED_CSV), - _name(LAYERWISE_COMBINED_CSV), - _name(TRAJECTORY_COMBINED_CSV), - _name(COMPONENT_ABLATION_COMBINED_CSV), - _name(MARGIN_COMBINED_CSV), - "defense_summary.csv", - } - files = [] - for path in data_dir.glob("results_*.csv"): - if path.name in blocked: - continue - if path.name.startswith("results_sweep_"): - continue - files.append(path) - return files - - -def combine_scalar_results(data_dir: Path, output: Path = RESULTS_CSV) -> pd.DataFrame: - """Combine per-model scalar result CSV files into one table.""" - files = _result_files(data_dir) - df = _read_csvs(files, model_prefix="results_") - return _write(df, output) - - -def combine_sweeps(data_dir: Path, output: Path = SWEEP_CSV) -> pd.DataFrame: - """Combine epsilon-sweep CSV files into one table.""" - files = [p for p in data_dir.glob("sweep_*.csv") if p.name != output.name] - df = _read_csvs(files, model_prefix="sweep_") - if not df.empty and {"model", "epsilon"}.issubset(df.columns): - df = df.sort_values(["model", "epsilon"]).reset_index(drop=True) - return _write(df, output) - - -def combine_ablation( - data_dir: Path, output: Path = ABLATION_COMBINED_CSV -) -> pd.DataFrame: - """Combine PGD step-ablation CSV files into one table.""" - files = [p for p in data_dir.glob("ablation_*.csv") if p.name != output.name] - df = _read_csvs(files, model_prefix="ablation_") - if not df.empty and {"model", "steps"}.issubset(df.columns): - df = df.sort_values(["model", "steps"]).reset_index(drop=True) - return _write(df, output) - - -def combine_layerwise( - data_dir: Path, output: Path = LAYERWISE_COMBINED_CSV -) -> pd.DataFrame: - """Combine layerwise gradient-profile CSV files into one table.""" - files = [p for p in data_dir.glob("layerwise_*.csv") if p.name != output.name] - df = _read_csvs(files, model_prefix="layerwise_") - return _write(df, output) - - -def combine_component_ablation( - data_dir: Path, output: Path = COMPONENT_ABLATION_COMBINED_CSV -) -> pd.DataFrame: - """Combine quantization component-ablation CSV files into one table.""" - files = [ - p for p in data_dir.glob("component_ablation_*.csv") if p.name != output.name - ] - df = _read_csvs(files, model_prefix="component_ablation_") - return _write(df, output) - - def combine_trajectories( data_dir: Path, output: Path = TRAJECTORY_COMBINED_CSV ) -> pd.DataFrame: @@ -181,7 +113,10 @@ def combine_trajectories( } ) - return _write(pd.DataFrame(rows), output) + combined = _merge_frames( + [_read_if_present(Path(output)), pd.DataFrame(rows)], ["model", "step"] + ) + return _write(combined, Path(output)) def combine_margins(data_dir: Path, output: Path = MARGIN_COMBINED_CSV) -> pd.DataFrame: @@ -203,16 +138,11 @@ def combine_margins(data_dir: Path, output: Path = MARGIN_COMBINED_CSV) -> pd.Da for i, value in enumerate(values): rows.append({"model": model, "kind": kind, "index": i, "margin": value}) - return _write(pd.DataFrame(rows), output) - - -def _model_names(*dfs: pd.DataFrame) -> list[str]: - """Return the preferred display order for model names.""" - names = [] - for df in dfs: - if df is not None and not df.empty and "model" in df.columns: - names.extend(df["model"].dropna().astype(str).tolist()) - return list(dict.fromkeys(names)) + combined = _merge_frames( + [_read_if_present(Path(output)), pd.DataFrame(rows)], + ["model", "kind", "index"], + ) + return _write(combined, Path(output)) def plot_summary_results(df_results: pd.DataFrame, output: Path = PLOT_PNG) -> None: @@ -273,7 +203,7 @@ def plot_epsilon_sweep_curves( return value_cols = [ c - for c in ["PGD_acc", "PGD_ste_acc", "Random_Noise_acc", "BPDA_acc"] + for c in ["PGD_acc", "Random_Noise_acc", "BPDA_acc"] if c in df_sweep.columns ] if not value_cols: @@ -442,14 +372,7 @@ def plot_layerwise_grad_profile( def plot_component_ablation( df_component: pd.DataFrame, output: Path = COMPONENT_ABLATION_PLOT_PNG ) -> None: - """Plot clean, hard-PGD, STE-PGD, and BPDA component-ablation accuracies. - - ``PGD_ste_acc`` is a budget-matched STE companion to ``PGD_acc`` - (``PGD_hard_acc``) added alongside ``BPDA_acc``. Expect ``weight_only`` - to show little gap between ``PGD_acc`` and ``PGD_ste_acc`` (weight - rounding doesn't sit on the gradient path back to the input) and - ``act_only``/``both`` to show a large one (activation rounding does). - """ + """Plot clean, vanilla-PGD, and BPDA-PGD component-ablation accuracies.""" required = {"model", "config", "clean_acc", "PGD_acc"} if ( df_component is None @@ -460,7 +383,7 @@ def plot_component_ablation( value_vars = [ c - for c in ["clean_acc", "PGD_acc", "PGD_ste_acc", "BPDA_acc"] + for c in ["clean_acc", "PGD_acc", "BPDA_acc"] if c in df_component.columns ] df_long = df_component.melt( @@ -665,11 +588,11 @@ def add_sweep_masking_metrics(df_sweep: pd.DataFrame) -> pd.DataFrame: """ if df_sweep is None or df_sweep.empty: return df_sweep - if {"PGD_acc", "PGD_ste_acc"}.issubset(df_sweep.columns): + if {"PGD_acc", "BPDA_acc"}.issubset(df_sweep.columns): df_sweep = df_sweep.copy() - df_sweep["PGD_masking_gap"] = pd.to_numeric( + df_sweep["PGD_minus_BPDA"] = pd.to_numeric( df_sweep["PGD_acc"], errors="coerce" - ) - pd.to_numeric(df_sweep["PGD_ste_acc"], errors="coerce") + ) - pd.to_numeric(df_sweep["BPDA_acc"], errors="coerce") return df_sweep @@ -776,7 +699,11 @@ def _merge_frames(frames: Iterable[pd.DataFrame], keys: list[str]) -> pd.DataFra def _combine_csv_family( - data_dir: Path, output: Path, patterns: Iterable[str], keys: list[str] + data_dir: Path, + output: Path, + patterns: Iterable[str], + keys: list[str], + model_prefix: str | None = None, ) -> pd.DataFrame: """Combine a family of CSV files using filename patterns and keys.""" paths = { @@ -785,10 +712,7 @@ def _combine_csv_family( for path in data_dir.glob(pattern) if path.resolve() != output.resolve() } - frames = [ - _read_if_present(output), - *(_read_if_present(path) for path in sorted(paths)), - ] + frames = [_read_if_present(output), _read_csvs(sorted(paths), model_prefix)] return _write(_merge_frames(frames, keys), output) @@ -965,102 +889,127 @@ def plot_defense_comparison(df: pd.DataFrame, output: Path = DEFENSE_PLOT_PNG) - plt.close(grid.fig) -def combine_all(data_dir: Path) -> dict[str, pd.DataFrame]: - """Combine all available experiment result families.""" - data_dir = Path(data_dir) - data_dir.mkdir(parents=True, exist_ok=True) +@dataclass(frozen=True) +class CsvFamily: + """Describe one family of partial CSV artifacts and its combined table.""" - result_path = data_dir / Path(RESULTS_CSV).name - df_results = add_paired_tests( - add_derived_metrics( - _combine_csv_family( - data_dir, - result_path, - ("results_*.csv", "accuracyresult*.csv"), - ["model"], - ) - ) - ) - _write(df_results, result_path) - df_sweep = _combine_csv_family( - data_dir, - data_dir / Path(SWEEP_CSV).name, - ("sweep_*.csv", "sweepresult*.csv"), - ["model", "epsilon"], - ) - df_sweep = add_sweep_masking_metrics(df_sweep) - df_ablation = combine_ablation( - data_dir, data_dir / Path(ABLATION_COMBINED_CSV).name - ) - df_layer = combine_layerwise(data_dir, data_dir / Path(LAYERWISE_COMBINED_CSV).name) - df_component = combine_component_ablation( - data_dir, data_dir / Path(COMPONENT_ABLATION_COMBINED_CSV).name - ) - df_traj = combine_trajectories( - data_dir, data_dir / Path(TRAJECTORY_COMBINED_CSV).name - ) - df_margins = combine_margins(data_dir, data_dir / Path(MARGIN_COMBINED_CSV).name) - df_performance = _combine_csv_family( + output: str + patterns: tuple[str, ...] + keys: tuple[str, ...] + transform: Callable[[pd.DataFrame], pd.DataFrame] | None = None + model_prefix: str | None = None + + +def combine_csv_family(data_dir: Path, family: CsvFamily) -> pd.DataFrame: + """Combine one configured CSV family, preserving partial aggregate data.""" + output = Path(data_dir) / _name(family.output) + frame = _combine_csv_family( data_dir, - data_dir / Path(PERFORMANCE_CSV).name, - ("performance_metrics*.csv",), - ["model"], + output, + family.patterns, + list(family.keys), + family.model_prefix, ) - df_dither = _combine_csv_family( - data_dir, - data_dir / Path(CHAOTIC_DITHER_SWEEP_CSV).name, + if family.transform is not None: + frame = family.transform(frame) + _write(frame, output) + return frame + +CSV_FAMILIES = { + "results": CsvFamily( + RESULTS_CSV, + ("results_*.csv", "accuracyresult*.csv"), + ("model",), + lambda frame: add_paired_tests(add_derived_metrics(frame)), + "results_", + ), + "sweep": CsvFamily( + SWEEP_CSV, + ("sweep_*.csv", "sweepresult*.csv"), + ("model", "epsilon"), + add_sweep_masking_metrics, + "sweep_", + ), + "ablation": CsvFamily( + ABLATION_COMBINED_CSV, + ("ablation_*.csv",), + ("model", "steps"), + model_prefix="ablation_", + ), + "layerwise": CsvFamily( + LAYERWISE_COMBINED_CSV, + ("layerwise_*.csv",), + ("model", "layer"), + model_prefix="layerwise_", + ), + "component_ablation": CsvFamily( + COMPONENT_ABLATION_COMBINED_CSV, + ("component_ablation_*.csv",), + ("model", "config"), + model_prefix="component_ablation_", + ), + "performance": CsvFamily( + PERFORMANCE_CSV, ("performance_metrics*.csv",), ("model",) + ), + "dither": CsvFamily( + CHAOTIC_DITHER_SWEEP_CSV, ("chaotic_dither_sweep*.csv",), - ["model", "bits", "dither_amplitude"], - ) - df_chunk = _combine_csv_family( - data_dir, - data_dir / Path(CHUNK_COMBINED_CSV).name, - ("chunk_quant_*.csv",), - ["model", "chunk_id"], - ) - df_defense = _read_if_present(data_dir / Path(DEFENSE_SUMMARY_CSV).name) - - return { - "results": df_results, - "sweep": df_sweep, - "ablation": df_ablation, - "layerwise": df_layer, - "component_ablation": df_component, - "trajectory": df_traj, - "margins": df_margins, - "performance": df_performance, - "dither": df_dither, - "chunk_quant": df_chunk, - "defense": df_defense, - } + ("model", "bits", "dither_amplitude"), + ), + "chunk_quant": CsvFamily( + CHUNK_COMBINED_CSV, ("chunk_quant_*.csv",), ("model", "chunk_id") + ), +} +def combine_all(data_dir: Path) -> dict[str, pd.DataFrame]: + """Combine all available experiment result families.""" + data_dir = Path(data_dir) + data_dir.mkdir(parents=True, exist_ok=True) + tables = { + name: combine_csv_family(data_dir, family) + for name, family in CSV_FAMILIES.items() + } + tables.update( + trajectory=combine_trajectories( + data_dir, data_dir / _name(TRAJECTORY_COMBINED_CSV) + ), + margins=combine_margins(data_dir, data_dir / _name(MARGIN_COMBINED_CSV)), + defense=_read_if_present(data_dir / _name(DEFENSE_SUMMARY_CSV)), + ) + return tables + + +@dataclass(frozen=True) +class PlotSpec: + """Connect a combined table to a plotting function and output filename.""" + + table: str + function: Callable[[pd.DataFrame, Path], None] + output: str + +PLOT_SPECS = ( + PlotSpec("results", plot_summary_results, PLOT_PNG), + PlotSpec("sweep", plot_epsilon_sweep_curves, SWEEP_PLOT_PNG), + PlotSpec("ablation", plot_pgd_steps_ablation, ABLATION_PLOT_PNG), + PlotSpec("trajectory", plot_pgd_trajectory, TRAJECTORY_PLOT_PNG), + PlotSpec("layerwise", plot_layerwise_grad_profile, LAYERWISE_PLOT_PNG), + PlotSpec( + "component_ablation", plot_component_ablation, COMPONENT_ABLATION_PLOT_PNG + ), + PlotSpec("results", plot_gradient_masking_summary, MASKING_SUMMARY_PLOT_PNG), + PlotSpec("margins", plot_confidence_margin_diagnostic, MARGIN_PLOT_PNG), + PlotSpec("results", plot_results_heatmap, HEATMAP_PLOT_PNG), + PlotSpec("performance", plot_performance, PERFORMANCE_PLOT_PNG), + PlotSpec("dither", plot_dither_sweep, DITHER_PLOT_PNG), + PlotSpec("chunk_quant", plot_chunk_quantization_attacks, CHUNK_QUANT_PLOT_PNG), + PlotSpec("defense", plot_defense_comparison, DEFENSE_PLOT_PNG), +) def plot_all(dfs: dict[str, pd.DataFrame], output_dir: Path = DATA_DIR) -> None: """Generate all summary plots from combined result tables.""" output_dir = Path(output_dir) - plot_summary_results(dfs["results"], output_dir / Path(PLOT_PNG).name) - plot_epsilon_sweep_curves(dfs["sweep"], output_dir / _name(SWEEP_PLOT_PNG)) - plot_pgd_steps_ablation(dfs["ablation"], output_dir / _name(ABLATION_PLOT_PNG)) - plot_pgd_trajectory(dfs["trajectory"], output_dir / _name(TRAJECTORY_PLOT_PNG)) - plot_layerwise_grad_profile( - dfs["layerwise"], output_dir / _name(LAYERWISE_PLOT_PNG) - ) - plot_component_ablation( - dfs["component_ablation"], output_dir / _name(COMPONENT_ABLATION_PLOT_PNG) - ) - plot_gradient_masking_summary( - dfs["results"], output_dir / _name(MASKING_SUMMARY_PLOT_PNG) - ) - plot_confidence_margin_diagnostic( - dfs["margins"], output_dir / _name(MARGIN_PLOT_PNG) - ) - plot_results_heatmap(dfs["results"], output_dir / _name(HEATMAP_PLOT_PNG)) - plot_performance(dfs["performance"], output_dir / _name(PERFORMANCE_PLOT_PNG)) - plot_dither_sweep(dfs["dither"], output_dir / _name(DITHER_PLOT_PNG)) - plot_chunk_quantization_attacks( - dfs["chunk_quant"], output_dir / _name(CHUNK_QUANT_PLOT_PNG) - ) - plot_defense_comparison(dfs["defense"], output_dir / _name(DEFENSE_PLOT_PNG)) + for spec in PLOT_SPECS: + spec.function(dfs.get(spec.table, pd.DataFrame()), output_dir / _name(spec.output)) overview_dir = output_dir / "visualizations" for name, frame in dfs.items(): @@ -1079,4 +1028,16 @@ def print_report(dfs: dict[str, pd.DataFrame]) -> None: for name, frame in dfs.items() ] ).to_string(index=False) - ) \ No newline at end of file + ) + + +def generate_reports( + data_dir: Path = DATA_DIR, *, plots: bool = True, summary: bool = True +) -> dict[str, pd.DataFrame]: + """Recover partial tables and optionally render and summarize all reports.""" + tables = combine_all(Path(data_dir)) + if plots: + plot_all(tables, Path(data_dir)) + if summary: + print_report(tables) + return tables From 8492e747f9d9e23769ef07fd3bf5d35674f0f054 Mon Sep 17 00:00:00 2001 From: Richard Date: Mon, 13 Jul 2026 12:57:31 +0800 Subject: [PATCH 2/4] Interoperability | Visualization - changed dataset to use CIFAR-100 - changed models to use three from TorchCV: ResNet56, WRN28_10, DenseNet100 - generalized more variables into config.py - changed visualizations to remove any old PGD-STE plotting, changed PGD ablation graph to show BPDA and FP32 stats --- README.md | 16 ++++----- src/QuantAdv.py | 96 +++++++++++++++++++++++++++---------------------- src/attack.py | 30 ++++++++-------- src/config.py | 68 ++++++++++++++++++++++++++--------- src/data.py | 22 ++++++++---- src/defense.py | 16 ++++----- 6 files changed, 151 insertions(+), 97 deletions(-) diff --git a/README.md b/README.md index cd546ae..728ce74 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,5 @@ # QuantAdv -Quantized models introduce discrete rounding operations into the computational graph, which may produce either genuine robustnes against *inference-time evasion attacks* (coarser weight representation changing the decision boundary geometry) or gradient masking (rounding causing zero gradients that blind attacks). We systematically evaluate several models (ResNet20, ResNet56 MobileNetV2, VGG16_BN, ShuffleNetV2, and RepVGG_A0) across four quantization variants (FP32, int8 PTQ, int4 PTQ, int8 QAT) on the CIFAR-10 image dataset using a layered attack suite (FGSM, PGD, AutoAttack, transfer attacks from FP32, and BPDA-corrected PGD) to test these two explanations. +Quantized models introduce discrete rounding operations into the computational graph, which may produce either genuine robustness against *inference-time evasion attacks* (coarser weight representation changing the decision boundary geometry) or gradient masking (rounding causing zero gradients that blind attacks). The current configuration evaluates pretrained TorchCV ResNet56, WRN-28-10, and DenseNet-100 models on CIFAR-100 across FP32, PTQ, and QAT variants using a layered attack suite. Dataset construction, preprocessing, class count, and TorchCV model identifiers are selected centrally in `src/config.py`. ## Setup @@ -22,15 +22,13 @@ curl -O https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz or set Download=true To run -`python src/run_experiment.py` -`python src/combine_results` -`python single/QuantAdv.py` +`python src/QuantAdv.py` -Parallelized (obsolete) -`python archive/launcher.py` +To graph incomplete results +`python src/combine.py` -To combine incomplete results -`python single/combine.py` +Obsolete run +`python src_old/launcher.py` > **Notice:** You may need to adjust pathing or move the scripts to root for obsolete files. @@ -44,4 +42,4 @@ Results are in ./data \ No newline at end of file +

--> diff --git a/src/QuantAdv.py b/src/QuantAdv.py index 9da27d4..e18fb4e 100644 --- a/src/QuantAdv.py +++ b/src/QuantAdv.py @@ -4,8 +4,6 @@ import torch import torch.nn as nn import torch.nn.functional as F -import torchvision -import torchvision.transforms as T import numpy as np import pandas as pd import copy @@ -43,7 +41,7 @@ """QuantAdv experiment runner for quantization and adversarial robustness. This module builds FP32, post-training quantized, quantization-aware trained, -and optional defended CIFAR-10 models, then evaluates them with white-box, +and optional defended image classifiers, then evaluates them with white-box, transfer, adaptive, black-box, and diagnostic attacks. Quantized layers are fake-quantized float modules: they simulate integer rounding during forward passes while optionally using straight-through gradients for attacks and QAT. @@ -66,7 +64,7 @@ def defense_summary_csv_path(): def check_environment(): - """Validate runtime packages and local CIFAR-10 data before a full run.""" + """Validate runtime packages and configured local data before a full run.""" missing = [ pkg for pkg in ("torchattacks", "autoattack", "pytorchcv", "torchao") @@ -78,36 +76,27 @@ def check_environment(): ) print("All required packages are available.") print("device:", device) - if not os.path.isdir(CIFAR10_DIR): - raise FileNotFoundError(f"Expected extracted CIFAR-10 at {CIFAR10_DIR!r}") + if not DATASET_DOWNLOAD and not os.path.isdir(DATASET_DIR): + raise FileNotFoundError( + f"Expected extracted {DATASET_NAME} data at {DATASET_DIR!r}" + ) def get_dataloaders( batch_size=DEFAULT_BATCH_SIZE, eval_n=DEFAULT_EVAL_N, finetune_n=DEFAULT_FINETUNE_N ): - """Build CIFAR-10 fine-tuning and evaluation loaders from fixed subsets.""" - transform_train = T.Compose( - [ - T.RandomCrop(CIFAR_IMAGE_SIZE, padding=CIFAR_RANDOM_CROP_PADDING), - T.RandomHorizontalFlip(), - T.ToTensor(), - T.Normalize(mean=CIFAR_MEAN_VALUES, std=CIFAR_STD_VALUES), - ] - ) - transform_test = T.Compose( - [T.ToTensor(), T.Normalize(mean=CIFAR_MEAN_VALUES, std=CIFAR_STD_VALUES)] + """Build fine-tuning and evaluation loaders for the configured dataset.""" + train_full = DATASET_CLASS( + root=DATASET_ROOT, + download=DATASET_DOWNLOAD, + transform=DATASET_TRAIN_TRANSFORM, + **DATASET_TRAIN_KWARGS, ) - train_full = torchvision.datasets.CIFAR10( - root=PROJECT_ROOT, - train=True, - download=CIFAR_DOWNLOAD, - transform=transform_train, - ) - test_full = torchvision.datasets.CIFAR10( - root=PROJECT_ROOT, - train=False, - download=CIFAR_DOWNLOAD, - transform=transform_test, + test_full = DATASET_CLASS( + root=DATASET_ROOT, + download=DATASET_DOWNLOAD, + transform=DATASET_TEST_TRANSFORM, + **DATASET_TEST_KWARGS, ) finetune_subset = torch.utils.data.Subset(train_full, list(range(finetune_n))) eval_subset = torch.utils.data.Subset(test_full, list(range(eval_n))) @@ -132,15 +121,19 @@ def get_dataloaders( def load_pretrained(arch_key): - """Load a supported pretrained CIFAR-10 architecture onto the active device.""" - if arch_key != "ResNet56": - raise ValueError(f"Unsupported architecture {arch_key!r}; expected 'ResNet56'.") + """Load a configured pretrained TorchCV architecture onto the active device.""" if ptcv_get_model is None: raise ImportError( "Missing package 'pytorchcv'. Install via: pip install -r requirements.txt" ) - model_name = PRETRAINED_NAMES[arch_key] - model = ptcv_get_model(model_name, pretrained=True, root=PYTORCHCV_MODEL_DIR) + try: + model_name = PRETRAINED_NAMES[arch_key] + except KeyError as exc: + raise ValueError( + f"Unknown architecture {arch_key!r}; configured choices are " + f"{tuple(PRETRAINED_NAMES)}" + ) from exc + model = ptcv_get_model(model_name, **PRETRAINED_MODEL_KWARGS) return model.to(device).eval() @@ -599,7 +592,7 @@ def forward(self, x): kwargs["align_corners"] = COMPRESS_IMAGE_ALIGN_CORNERS pixels = F.interpolate(pixels, size=(self.size, self.size), **kwargs) pixels = F.interpolate( - pixels, size=(CIFAR_IMAGE_SIZE, CIFAR_IMAGE_SIZE), **kwargs + pixels, size=(DATASET_IMAGE_SIZE, DATASET_IMAGE_SIZE), **kwargs ) if self.bits is not None: pixels = ImageCompressionSTE.apply(pixels, self.bits) @@ -1057,6 +1050,33 @@ def run_suite(model, loader, name, fp32_ref=None, eps=DEFAULT_EPS): context=name, defaults=adaptive_defaults, ) + + if any( + name == f"{arch_key}_{variant}" + for arch_key in PRETRAINED_NAMES + for variant in PGD_ABLATION_VARIANTS + ): + + def save_pgd_ablation(): + """Persist PGD and BPDA step-ablation diagnostics.""" + rows = [] + for attack_name, use_ste in (("PGD", False), ("BPDA_PGD", True)): + ablation = pgd_steps_ablation( + model, loader, eps=eps, use_ste=use_ste + ) + rows.extend( + { + "model": name, + "attack": attack_name, + "steps": steps, + "acc": accuracy, + } + for steps, accuracy in ablation.items() + ) + pd.DataFrame(rows).to_csv(csv_path(name, "ablation"), index=False) + + safe_call(save_pgd_ablation, "pgd_steps_ablation failed", context=name) + if count_quant_layers(model) > 0: safe_update_vectors( results, @@ -1130,14 +1150,6 @@ def run_suite(model, loader, name, fp32_ref=None, eps=DEFAULT_EPS): defaults={"NES": None}, ) - def save_pgd_ablation(): - """Persist PGD step-ablation diagnostics for the current model.""" - ablation = pgd_steps_ablation(model, loader, eps=eps) - pd.DataFrame( - [{"model": name, "steps": k, "acc": v} for k, v in ablation.items()] - ).to_csv(csv_path(name, "ablation"), index=False) - - safe_call(save_pgd_ablation, "pgd_steps_ablation failed", context=name) if RUN_PGD_TRAJECTORY: def save_trajectory(): diff --git a/src/attack.py b/src/attack.py index ff17d8a..6de9740 100644 --- a/src/attack.py +++ b/src/attack.py @@ -78,13 +78,13 @@ def ste_mode(model, flag): def normalize_pixels(x): - """Normalize pixel-space CIFAR tensors with the experiment constants.""" - return (x - CIFAR_MEAN.to(x.device)) / CIFAR_STD.to(x.device) + """Normalize pixel-space tensors with the configured dataset constants.""" + return (x - DATASET_MEAN.to(x.device)) / DATASET_STD.to(x.device) def denormalize_inputs(x): - """Map normalized CIFAR tensors back to pixel space.""" - return x * CIFAR_STD.to(x.device) + CIFAR_MEAN.to(x.device) + """Map normalized dataset tensors back to pixel space.""" + return x * DATASET_STD.to(x.device) + DATASET_MEAN.to(x.device) class PixelSpaceModel(nn.Module): @@ -93,17 +93,17 @@ class PixelSpaceModel(nn.Module): def __init__(self, model): super().__init__() self.model = model - self.register_buffer("mean", CIFAR_MEAN.clone()) - self.register_buffer("std", CIFAR_STD.clone()) + self.register_buffer("mean", DATASET_MEAN.clone()) + self.register_buffer("std", DATASET_STD.clone()) def forward(self, x): return self.model((x - self.mean.to(x.device)) / self.std.to(x.device)) def make_torchattack(attack_cls, model, *args, **kwargs): - """Create a torchattacks instance configured for normalized CIFAR inputs.""" + """Create a torchattacks instance for normalized configured-dataset inputs.""" attack = attack_cls(model, *args, **kwargs) - attack.set_normalization_used(mean=CIFAR_MEAN_VALUES, std=CIFAR_STD_VALUES) + attack.set_normalization_used(mean=DATASET_MEAN_VALUES, std=DATASET_STD_VALUES) return attack @@ -414,8 +414,8 @@ def projected_pgd_attack(x, y, eps, alpha, steps, grad_fn): """PGD over normalized tensors with an externally supplied gradient function.""" clip_min = CLIP_MIN.to(device) clip_max = CLIP_MAX.to(device) - eps_normalized = eps / CIFAR_STD.to(x.device) - alpha_normalized = alpha / CIFAR_STD.to(x.device) + eps_normalized = eps / DATASET_STD.to(x.device) + alpha_normalized = alpha / DATASET_STD.to(x.device) x_adv = x.clone().detach() + torch.empty_like(x).uniform_(-1, 1) * eps_normalized x_adv = torch.max(torch.min(x_adv, clip_max), clip_min).detach() for _ in range(steps): @@ -710,7 +710,7 @@ def __init__(self, num_classes=SUBSTITUTE_NUM_CLASSES): super().__init__() self.features = nn.Sequential( nn.Conv2d( - 3, + DATASET_INPUT_CHANNELS, SUBSTITUTE_CONV1_CHANNELS, SUBSTITUTE_KERNEL_SIZE, padding=SUBSTITUTE_CONV_PADDING, @@ -995,7 +995,7 @@ def random_noise_attack( torch.manual_seed(seed) model.eval() clip_min, clip_max = CLIP_MIN.to(device), CLIP_MAX.to(device) - eps_normalized = eps / CIFAR_STD.to(device) + eps_normalized = eps / DATASET_STD.to(device) correct, total, vectors = 0, 0, [] with torch.no_grad(): for x, y in loader: @@ -1075,8 +1075,8 @@ def pgd_trajectory_diagnostics( """ model.eval() clip_min, clip_max = CLIP_MIN.to(device), CLIP_MAX.to(device) - eps_normalized = eps / CIFAR_STD.to(device) - alpha_normalized = alpha / CIFAR_STD.to(device) + eps_normalized = eps / DATASET_STD.to(device) + alpha_normalized = alpha / DATASET_STD.to(device) step_grad_norms = [0.0] * steps step_movement = [0.0] * steps n_batches = 0 @@ -1118,7 +1118,7 @@ def staircase_diagnostic( flat_norm = direction.flatten(1).norm(dim=1).view(-1, *([1] * (x.dim() - 1))) direction = direction / flat_norm clip_min, clip_max = CLIP_MIN.to(device), CLIP_MAX.to(device) - radius_normalized = radius / CIFAR_STD.to(device) + radius_normalized = radius / DATASET_STD.to(device) with torch.no_grad(): prev_logits = model(x) plateau_hits = 0.0 diff --git a/src/config.py b/src/config.py index 526dab8..2a02706 100644 --- a/src/config.py +++ b/src/config.py @@ -11,6 +11,8 @@ from pathlib import Path import torch +import torchvision.datasets as tv_datasets +import torchvision.transforms as tv_transforms logging.getLogger("torch.utils._pytree").setLevel(logging.ERROR) @@ -22,9 +24,48 @@ PROJECT_ROOT = CONFIG_DIR.parent if CONFIG_DIR.name in {"src", "src_old"} else CONFIG_DIR DATA_DIR = os.path.join(PROJECT_ROOT, "data") os.makedirs(DATA_DIR, exist_ok=True) -CIFAR10_DIR = os.path.join(PROJECT_ROOT, "cifar-10-batches-py") PYTORCHCV_MODEL_DIR = os.path.join(PROJECT_ROOT, "models", "pytorchcv") +DATASET_NAME = "CIFAR-100" +DATASET_CLASS = tv_datasets.CIFAR100 +DATASET_ROOT = PROJECT_ROOT +DATASET_DIR = os.path.join(DATASET_ROOT, "cifar-100-python") +DATASET_DOWNLOAD = False +DATASET_TRAIN_KWARGS = {"train": True} +DATASET_TEST_KWARGS = {"train": False} +DATASET_NUM_CLASSES = 100 +DATASET_INPUT_CHANNELS = 3 +DATASET_IMAGE_SIZE = 32 +DATASET_RANDOM_CROP_PADDING = 4 +DATASET_MEAN_VALUES = (0.5071, 0.4867, 0.4408) +DATASET_STD_VALUES = (0.2675, 0.2565, 0.2761) +DATASET_MEAN = torch.tensor(DATASET_MEAN_VALUES).view( + 1, DATASET_INPUT_CHANNELS, 1, 1 +) +DATASET_STD = torch.tensor(DATASET_STD_VALUES).view( + 1, DATASET_INPUT_CHANNELS, 1, 1 +) +DATASET_TRAIN_TRANSFORM = tv_transforms.Compose( + [ + tv_transforms.RandomCrop( + DATASET_IMAGE_SIZE, padding=DATASET_RANDOM_CROP_PADDING + ), + tv_transforms.RandomHorizontalFlip(), + tv_transforms.ToTensor(), + tv_transforms.Normalize(mean=DATASET_MEAN_VALUES, std=DATASET_STD_VALUES), + ] +) +DATASET_TEST_TRANSFORM = tv_transforms.Compose( + [ + tv_transforms.ToTensor(), + tv_transforms.Normalize(mean=DATASET_MEAN_VALUES, std=DATASET_STD_VALUES), + ] +) +CLIP_MIN = (0.0 - DATASET_MEAN) / DATASET_STD +CLIP_MAX = (1.0 - DATASET_MEAN) / DATASET_STD +CLIP_MIN_DEV = CLIP_MIN.to(device) +CLIP_MAX_DEV = CLIP_MAX.to(device) + RESULTS_CSV = os.path.join(DATA_DIR, "accuracyresult.csv") PER_EXAMPLE_DIR = os.path.join(DATA_DIR, "per_example") PERFORMANCE_CSV = os.path.join(DATA_DIR, "performance_metrics.csv") @@ -56,20 +97,17 @@ SEEDS = [0, 1, 2, 3, 4] PRETRAINED_NAMES = { - "ResNet56": "resnet56_cifar10", + "ResNet56_C100": "resnet56_cifar100", + "WRN28_10_C100": "wrn28_10_cifar100", + "DenseNet100_C100": "densenet100_k12_bc_cifar100", +} +PRETRAINED_MODEL_KWARGS = { + "pretrained": True, + "root": PYTORCHCV_MODEL_DIR, } QUANTIZATION_DEBUG_ONLY = False -CIFAR_MEAN_VALUES = (0.4914, 0.4822, 0.4465) -CIFAR_STD_VALUES = (0.2023, 0.1994, 0.2010) -CIFAR_MEAN = torch.tensor(CIFAR_MEAN_VALUES).view(1, 3, 1, 1) -CIFAR_STD = torch.tensor(CIFAR_STD_VALUES).view(1, 3, 1, 1) -CLIP_MIN = (0.0 - CIFAR_MEAN) / CIFAR_STD -CLIP_MAX = (1.0 - CIFAR_MEAN) / CIFAR_STD -CLIP_MIN_DEV = CLIP_MIN.to(device) -CLIP_MAX_DEV = CLIP_MAX.to(device) - DEFAULT_BATCH_SIZE = 512 DEFAULT_EVAL_N = 5000 DEFAULT_FINETUNE_N = 10000 @@ -78,9 +116,6 @@ PIN_MEMORY = True NON_BLOCKING_TRANSFER = PIN_MEMORY and device.type == "cuda" -CIFAR_IMAGE_SIZE = 32 -CIFAR_RANDOM_CROP_PADDING = 4 -CIFAR_DOWNLOAD = False TRAIN_SHUFFLE = True EVAL_SHUFFLE = False @@ -122,7 +157,7 @@ DEFAULT_EPS = 8 / 255 PGD_ALPHA = 2 / 255 -PGD_STEPS = 15 +PGD_STEPS = 20 PGD_RANDOM_START = True QAT_EPOCHS_DEFAULT = 3 @@ -194,7 +229,7 @@ CI_CONFIDENCE = 0.95 -SUBSTITUTE_NUM_CLASSES = 10 +SUBSTITUTE_NUM_CLASSES = DATASET_NUM_CLASSES SUBSTITUTE_ROUNDS = 6 SUBSTITUTE_EPOCHS_PER_ROUND = 10 SUBSTITUTE_LR = 1e-3 @@ -220,6 +255,7 @@ GRAD_DIAG_MAX_BATCHES = 3 GRAD_ZERO_THRESHOLD = 1e-8 PGD_ABLATION_STEPS = (0, 1, 2, 5, 10, 20, 50) +PGD_ABLATION_VARIANTS = ("FP32", "int8_PTQ", "int4_PTQ", "int8_QAT", "int4_QAT") TRAJECTORY_MAX_BATCHES = 3 LAYERWISE_MAX_BATCHES = 3 STAIRCASE_RADIUS = 1 / 255 diff --git a/src/data.py b/src/data.py index 9171a91..5bcfacf 100644 --- a/src/data.py +++ b/src/data.py @@ -67,7 +67,7 @@ def _read_csvs(paths: Iterable[Path], model_prefix: str | None = None) -> pd.Dat out = out.drop_duplicates( subset=[ c - for c in ["model", "epsilon", "steps", "layer", "config"] + for c in ["model", "epsilon", "steps", "attack", "layer", "config"] if c in out.columns ], keep="last", @@ -256,18 +256,26 @@ def plot_epsilon_sweep_curves( def plot_pgd_steps_ablation( df_ablation: pd.DataFrame, output: Path = ABLATION_PLOT_PNG ) -> None: - """Plot accuracy as PGD step count changes.""" + """Plot matched PGD and BPDA accuracy as attack step count changes.""" if ( df_ablation is None or df_ablation.empty - or not {"steps", "acc", "model"}.issubset(df_ablation.columns) + or not {"steps", "acc", "model", "attack"}.issubset(df_ablation.columns) ): return plt.figure(figsize=ABLATION_FIGSIZE) - sns.lineplot(data=df_ablation, x="steps", y="acc", hue="model", marker="o") - plt.title("PGD Accuracy vs Number of Steps (Gradient Masking Check)") - plt.xlabel("PGD steps") + sns.lineplot( + data=df_ablation, + x="steps", + y="acc", + hue="model", + style="attack", + markers=True, + dashes=True, + ) + plt.title("PGD vs BPDA Accuracy by Attack Steps (Gradient Masking Check)") + plt.xlabel("Attack steps") plt.ylabel("Accuracy") plt.ylim(0, PLOT_MAX_ACCURACY) plt.grid(linestyle="--", alpha=PLOT_GRID_ALPHA) @@ -933,7 +941,7 @@ def combine_csv_family(data_dir: Path, family: CsvFamily) -> pd.DataFrame: "ablation": CsvFamily( ABLATION_COMBINED_CSV, ("ablation_*.csv",), - ("model", "steps"), + ("model", "steps", "attack"), model_prefix="ablation_", ), "layerwise": CsvFamily( diff --git a/src/defense.py b/src/defense.py index 3ee7650..9a9dcba 100644 --- a/src/defense.py +++ b/src/defense.py @@ -32,19 +32,19 @@ def _device(): def normalize_pixels(x): - """Normalize pixel-space CIFAR tensors with the experiment constants.""" - return (x - CIFAR_MEAN.to(x.device)) / CIFAR_STD.to(x.device) + """Normalize pixel-space tensors with the configured dataset constants.""" + return (x - DATASET_MEAN.to(x.device)) / DATASET_STD.to(x.device) def denormalize_inputs(x): - """Map normalized CIFAR tensors back to pixel space.""" - return x * CIFAR_STD.to(x.device) + CIFAR_MEAN.to(x.device) + """Map normalized dataset tensors back to pixel space.""" + return x * DATASET_STD.to(x.device) + DATASET_MEAN.to(x.device) def make_attack(attack_cls, model, *args, **kwargs): - """Construct a torchattacks object for already-normalized CIFAR inputs.""" + """Construct a torchattacks object for normalized configured-dataset inputs.""" attack = attack_cls(model, *args, **kwargs) - attack.set_normalization_used(mean=CIFAR_MEAN_VALUES, std=CIFAR_STD_VALUES) + attack.set_normalization_used(mean=DATASET_MEAN_VALUES, std=DATASET_STD_VALUES) return attack @@ -248,7 +248,7 @@ def forward(self, x): ) pixels = F.interpolate( pixels, - size=(CIFAR_IMAGE_SIZE, CIFAR_IMAGE_SIZE), + size=(DATASET_IMAGE_SIZE, DATASET_IMAGE_SIZE), mode="bilinear", align_corners=False, ) @@ -304,7 +304,7 @@ def __init__(self): """Build the convolutional binary adversarial-example detector.""" super().__init__() self.net = nn.Sequential( - nn.Conv2d(3, 32, 3, padding=1), + nn.Conv2d(DATASET_INPUT_CHANNELS, 32, 3, padding=1), nn.ReLU(inplace=True), nn.MaxPool2d(2), nn.Conv2d(32, 64, 3, padding=1), From 45639df1a0e5b1cb87d01da8e31431632ab6dc32 Mon Sep 17 00:00:00 2001 From: Richard Date: Mon, 13 Jul 2026 13:49:12 +0800 Subject: [PATCH 3/4] Singular PGD | fmt - hopefully this is the last entry in PGD issues, merged the ablation and accuracy PGDs into the same utility - old ablation PGD took worst of 5 runs - new merged PGD uses static seeds - formatting --- src/QuantAdv.py | 98 +++++++-------- src/attack.py | 313 ++++++++++++++++++++++++++++++------------------ src/combine.py | 7 +- src/config.py | 12 +- src/data.py | 33 +++-- src/defense.py | 12 +- src/stats.py | 8 +- 7 files changed, 274 insertions(+), 209 deletions(-) diff --git a/src/QuantAdv.py b/src/QuantAdv.py index e18fb4e..ab2e28e 100644 --- a/src/QuantAdv.py +++ b/src/QuantAdv.py @@ -157,8 +157,7 @@ def _torchao_int_dtype(bits): def _make_torchao_fake_quantizer(bits, *, role): - """Build TorchAO fake quantization for a weight or activation tensor. - """ + """Build TorchAO fake quantization for a weight or activation tensor.""" if role == "activation": config = IntxFakeQuantizeConfig( dtype=_torchao_int_dtype(bits), @@ -234,7 +233,9 @@ class _QuantizedLayerMixin: def _init_quantizers(self, bits): self.bits = bits self.weight_fake_quantizer = _make_torchao_fake_quantizer(bits, role="weight") - self.activation_fake_quantizer = _make_torchao_fake_quantizer(bits, role="activation") + self.activation_fake_quantizer = _make_torchao_fake_quantizer( + bits, role="activation" + ) def _quant_params(self): return ( @@ -466,7 +467,9 @@ def quant_layer_chunks(layer_names, n_chunks): def count_quant_layers(model): return sum( - isinstance(mod, (QuantConv2d, QuantLinear, ChaoticQuantConv2d, ChaoticQuantLinear)) + isinstance( + mod, (QuantConv2d, QuantLinear, ChaoticQuantConv2d, ChaoticQuantLinear) + ) for mod in model.modules() ) @@ -494,7 +497,9 @@ def verify_quantization_layers( def set_quant_components(model, quant_weight, quant_act): for mod in model.modules(): - if isinstance(mod, (QuantConv2d, QuantLinear, ChaoticQuantConv2d, ChaoticQuantLinear)): + if isinstance( + mod, (QuantConv2d, QuantLinear, ChaoticQuantConv2d, ChaoticQuantLinear) + ): mod.quant_weight = quant_weight mod.quant_act = quant_act @@ -620,8 +625,7 @@ def add_paired_fp32_mcnemar_tests(df_results): def gradient_diagnostics( model, loader, fp32_ref=None, max_batches=GRAD_DIAG_MAX_BATCHES ): - """Ground-truth masking check: compare hard-round vs. STE input gradients. - """ + """Ground-truth masking check: compare hard-round vs. STE input gradients.""" frac_zero_hard, norm_hard = [], [] frac_zero_ste, norm_ste = [], [] cos_sims = [] @@ -674,6 +678,7 @@ def layerwise_grad_profile(model, loader, use_ste, max_batches=LAYERWISE_MAX_BAT def make_hook(name): """Create a backward hook that records gradient statistics for one layer.""" + def hook(module, grad_input, grad_output): """Record gradient statistics emitted by a backward hook.""" gi = grad_input[0] @@ -731,18 +736,9 @@ def run_quant_component_ablation(model, loader, name, eps=DEFAULT_EPS): set_quant_components(model, qw, qa) clean_acc = sanity_check_accuracy(model, loader) + pgd_hard = run_pgd(model, loader, eps=eps, seeds=SEEDS) + pgd_hard_acc = pgd_hard["PGD"] with ste_mode(model, False): - torch.manual_seed(0) - pgd_hard = make_torchattack( - torchattacks.PGD, - model, - eps=eps, - alpha=PGD_ALPHA, - steps=PGD_STEPS, - random_start=PGD_RANDOM_START, - ) - pgd_hard_acc = accuracy_under_attack(model, loader, pgd_hard) - x, y = next(iter(loader)) x, y = x.to(device), y.to(device) x_in = x.clone().requires_grad_(True) @@ -755,7 +751,7 @@ def run_quant_component_ablation(model, loader, name, eps=DEFAULT_EPS): loader, eps=eps, n_restarts=1, - seeds=SEEDS[:1], + seeds=SEEDS, ) rows.append( { @@ -766,6 +762,8 @@ def run_quant_component_ablation(model, loader, name, eps=DEFAULT_EPS): "clean_acc": clean_acc, "PGD_acc": pgd_hard_acc, "PGD_hard_acc": pgd_hard_acc, + "PGD_mean": pgd_hard["PGD_mean"], + "PGD_std": pgd_hard["PGD_std"], "BPDA_acc": bpda["BPDA_PGD"], "BPDA_mean": bpda["BPDA_PGD_mean"], "BPDA_std": bpda["BPDA_PGD_std"], @@ -817,15 +815,10 @@ def run_chunk_quantization_attacks( ) row["clean_acc"] = None try: - pgd = make_torchattack( - torchattacks.PGD, - chunk_model, - eps=eps, - alpha=PGD_ALPHA, - steps=PGD_STEPS, - random_start=PGD_RANDOM_START, - ) - row["PGD_acc"] = accuracy_under_attack(chunk_model, loader, pgd) + pgd = run_pgd(chunk_model, loader, eps=eps, seeds=SEEDS) + row["PGD_acc"] = pgd["PGD"] + row["PGD_mean"] = pgd["PGD_mean"] + row["PGD_std"] = pgd["PGD_std"] except Exception as e: print(f" [WARN] chunk PGD failed for {name} {row['chunk_label']}: {e}") row["PGD_acc"] = None @@ -938,9 +931,7 @@ def run_suite(model, loader, name, fp32_ref=None, eps=DEFAULT_EPS): if RUN_EXTRA_WHITEBOX_ATTACKS: safe_update( results, - lambda: run_extra_whitebox_attacks( - model, loader, eps=eps, use_ste=False - ), + lambda: run_extra_whitebox_attacks(model, loader, eps=eps, use_ste=False), "CW/DeepFool/JSMA failed", context=name, ) @@ -1061,9 +1052,7 @@ def save_pgd_ablation(): """Persist PGD and BPDA step-ablation diagnostics.""" rows = [] for attack_name, use_ste in (("PGD", False), ("BPDA_PGD", True)): - ablation = pgd_steps_ablation( - model, loader, eps=eps, use_ste=use_ste - ) + ablation = pgd_steps_ablation(model, loader, eps=eps, use_ste=use_ste) rows.extend( { "model": name, @@ -1239,6 +1228,20 @@ def run_epsilon_sweep_for_model_wrapped(model, loader, name, epsilons): ) +def completed_sweep_keys(df_sweep): + """Return only rows produced by the shared seeded sweep implementation.""" + required = {"model", "epsilon", "PGD_acc", "PGD_mean", "PGD_std"} + if df_sweep.empty or not required.issubset(df_sweep.columns): + return set() + complete = df_sweep[["PGD_acc", "PGD_mean", "PGD_std"]].notna().all(axis=1) + return set( + zip( + df_sweep.loc[complete, "model"].astype(str), + df_sweep.loc[complete, "epsilon"].round(6), + ) + ) + + def run_defense_suite(model_registry, finetune_loader, eval_loader): """Build defended models and evaluate them under adaptive attacks.""" summary_rows = [] @@ -1420,23 +1423,16 @@ def run_chaotic_dither_sweep(fp32_model, loader, arch_key, bits=8): .eval() ) clean = sanity_check_accuracy(model, loader) - attack = make_torchattack( - torchattacks.PGD, - model, - eps=DEFAULT_EPS, - alpha=PGD_ALPHA, - steps=PGD_STEPS, - random_start=PGD_RANDOM_START, - ) - torch.manual_seed(SEEDS[0]) - pgd = accuracy_under_attack(model, loader, attack) + pgd = run_pgd(model, loader, eps=DEFAULT_EPS, seeds=SEEDS) rows.append( { "model": arch_key, "bits": bits, "dither_amplitude": amplitude, "clean_acc": clean, - "PGD_acc": pgd, + "PGD_acc": pgd["PGD"], + "PGD_mean": pgd["PGD_mean"], + "PGD_std": pgd["PGD_std"], } ) del model @@ -1507,9 +1503,7 @@ def main(): traceback.print_exc() continue try: - int8_ptq = convert_to_quant( - fp32, bits=8, quant_weight=True, quant_act=True - ) + int8_ptq = convert_to_quant(fp32, bits=8, quant_weight=True, quant_act=True) verify_quantization_layers( arch_key, fp32, int8_ptq, "int8 PTQ", fp32_layer_names ) @@ -1697,9 +1691,7 @@ def main(): run_metrics = [] if RUN_EPSILON_SWEEP and os.path.exists(SWEEP_CSV): df_sweep = pd.read_csv(SWEEP_CSV) - sweep_done = set( - zip(df_sweep["model"].astype(str), df_sweep["epsilon"].round(6)) - ) + sweep_done = completed_sweep_keys(df_sweep) else: df_sweep = pd.DataFrame() sweep_done = set() @@ -1729,9 +1721,7 @@ def save_epsilon_sweep(): df_sweep = report_data.upsert_table( SWEEP_CSV, new_sweep, ["model", "epsilon"] ) - sweep_done = set( - zip(df_sweep["model"].astype(str), df_sweep["epsilon"].round(6)) - ) + sweep_done = completed_sweep_keys(df_sweep) safe_call( save_epsilon_sweep, diff --git a/src/attack.py b/src/attack.py index 6de9740..211edf7 100644 --- a/src/attack.py +++ b/src/attack.py @@ -10,6 +10,7 @@ - diagnostics that are attack-driven (trajectory, margins, staircase, PGD-steps ablation) """ +import random import warnings from contextlib import contextmanager @@ -31,8 +32,7 @@ def unwrap_model(model): def set_ste_mode(model, flag): - """Toggle straight-through gradients on fake-quantized modules. - """ + """Toggle straight-through gradients on fake-quantized modules.""" toggled = 0 for mod in model.modules(): if hasattr(mod, "use_ste") and hasattr(mod, "bits"): @@ -42,8 +42,7 @@ def set_ste_mode(model, flag): def get_ste_mode(model): - """Return the current ``use_ste`` state of a model's quantized modules. - """ + """Return the current ``use_ste`` state of a model's quantized modules.""" values = { mod.use_ste for mod in model.modules() @@ -127,7 +126,9 @@ def accuracy_from_adv_fn( break remaining = max_images - n_seen x, y = x[:remaining], y[:remaining] - x, y = x.to(device, non_blocking=NON_BLOCKING_TRANSFER), y.to(device, non_blocking=NON_BLOCKING_TRANSFER) + x, y = x.to(device, non_blocking=NON_BLOCKING_TRANSFER), y.to( + device, non_blocking=NON_BLOCKING_TRANSFER + ) x_adv = adv_fn(x, y) if adv_fn is not None else x with torch.no_grad(): if use_autocast: @@ -156,6 +157,7 @@ def accuracy_under_attack( model, loader, attack, target_model=None, max_images=None, return_vector=False ): """Measure accuracy under a torchattacks-compatible attack object.""" + def adv_fn(x, y): return attack(x, y) @@ -189,6 +191,94 @@ def seed_averaged_metrics(name, seeds, fn): } +def _seed_restart(seed): + """Seed every RNG used by an explicitly identified attack restart.""" + seed = int(seed) + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +def _aggregate_restarts(name, accuracies, vectors, return_vector=False): + """Aggregate seeded restarts by per-example worst-case correctness.""" + if not vectors: + raise ValueError(f"{name} requires at least one restart seed") + worst_vector = np.logical_and.reduce(vectors) + out = { + name: float(worst_vector.mean()), + f"{name}_mean": float(np.mean(accuracies)), + f"{name}_std": float(np.std(accuracies)), + } + if return_vector: + out["_vectors"] = {name: worst_vector} + return out + + +def _run_seeded_pgd( + model, + loader, + *, + eps, + alpha, + steps, + random_start, + seeds, + use_ste, + return_vector, +): + """Single implementation for seeded torchattacks PGD evaluation.""" + model.eval() + accuracies, vectors = [], [] + with ste_mode(model, use_ste): + for seed in seeds: + _seed_restart(seed) + attack = make_torchattack( + torchattacks.PGD, + model, + eps=eps, + alpha=alpha, + steps=steps, + random_start=random_start, + ) + accuracy, vector = accuracy_under_attack( + model, loader, attack, return_vector=True + ) + accuracies.append(accuracy) + vectors.append(vector) + return _aggregate_restarts("PGD", accuracies, vectors, return_vector) + + +def run_pgd( + model, + loader, + eps=DEFAULT_EPS, + alpha=PGD_ALPHA, + steps=PGD_STEPS, + random_start=PGD_RANDOM_START, + seeds=SEEDS, + return_vector=False, +): + """Evaluate hard-round PGD with the canonical seeded restart policy. + + Each seed is one restart. ``PGD`` is the intersection of per-example + correctness across restarts; ``PGD_mean`` and ``PGD_std`` summarize the + individual restart accuracies. + """ + return _run_seeded_pgd( + model, + loader, + eps=eps, + alpha=alpha, + steps=steps, + random_start=random_start, + seeds=seeds, + use_ste=False, + return_vector=return_vector, + ) + + def run_fgsm_pgd( model, loader, eps=DEFAULT_EPS, seeds=SEEDS, return_vectors=False, use_ste=False ): @@ -211,35 +301,31 @@ def adv_fn(x, y): fgsm_acc, fgsm_vector = attack_vector(fgsm) out["FGSM"] = fgsm_acc - pgd_vectors = [] - pgd_accs = [] - for seed in seeds: - torch.manual_seed(seed) - pgd = make_torchattack( - torchattacks.PGD, + pgd_out = ( + run_pgd( + model, loader, eps=eps, seeds=seeds, return_vector=return_vectors + ) + if not use_ste + else _run_seeded_pgd( model, + loader, eps=eps, alpha=PGD_ALPHA, steps=PGD_STEPS, random_start=PGD_RANDOM_START, + seeds=seeds, + use_ste=True, + return_vector=return_vectors, ) - acc, vector = attack_vector(pgd) - pgd_accs.append(acc) - pgd_vectors.append(vector) - # A multi-restart attack is the intersection of per-restart correctness, - # not the mean of restart accuracies. - pgd_vector = np.logical_and.reduce(pgd_vectors) - out["PGD"] = float(pgd_vector.mean()) - out["PGD_mean"] = float(np.mean(pgd_accs)) - out["PGD_std"] = float(np.std(pgd_accs)) + ) + pgd_vectors = pgd_out.pop("_vectors", {}) + out.update(pgd_out) if return_vectors: - out["_vectors"] = {"FGSM": fgsm_vector, "PGD": pgd_vector} + out["_vectors"] = {"FGSM": fgsm_vector, **pgd_vectors} return out -def run_autoattack( - model, loader, eps=DEFAULT_EPS, return_vector=False, use_ste=False -): +def run_autoattack(model, loader, eps=DEFAULT_EPS, return_vector=False, use_ste=False): """Run AutoAttack in pixel space while reporting normalized-model accuracy.""" model.eval() with ste_mode(model, use_ste): @@ -369,7 +455,9 @@ def build_uap( for x, y in loader: if n_seen >= max_images: break - x, y = x.to(device, non_blocking=NON_BLOCKING_TRANSFER), y.to(device, non_blocking=NON_BLOCKING_TRANSFER) + x, y = x.to(device, non_blocking=NON_BLOCKING_TRANSFER), y.to( + device, non_blocking=NON_BLOCKING_TRANSFER + ) x_pert = torch.max(torch.min(x + v, clip_max), clip_min) with torch.no_grad(): pred_orig = model(x).argmax(dim=1) @@ -446,13 +534,23 @@ def grad_fn(x_adv, labels): return projected_pgd_attack(x, y, eps, alpha, steps, grad_fn) -def _run_bpda_once(model, loader, eps, n_restarts, return_vector=False): +def _run_bpda_once( + model, + loader, + eps, + n_restarts, + alpha=PGD_ALPHA, + steps=PGD_STEPS, + return_vector=False, +): correct_masks = [] for x, y in loader: - x, y = x.to(device, non_blocking=NON_BLOCKING_TRANSFER), y.to(device, non_blocking=NON_BLOCKING_TRANSFER) + x, y = x.to(device, non_blocking=NON_BLOCKING_TRANSFER), y.to( + device, non_blocking=NON_BLOCKING_TRANSFER + ) worst_correct = torch.ones(y.size(0), dtype=torch.bool, device=device) for _ in range(n_restarts): - x_adv = bpda_pgd_attack(model, x, y, eps=eps) + x_adv = bpda_pgd_attack(model, x, y, eps=eps, alpha=alpha, steps=steps) with torch.no_grad(): pred = model(x_adv).argmax(dim=1) worst_correct &= pred == y @@ -473,26 +571,25 @@ def run_bpda( n_restarts=BPDA_RESTARTS_DEFAULT, seeds=SEEDS, return_vector=False, + alpha=PGD_ALPHA, + steps=PGD_STEPS, ): - """Evaluate BPDA-PGD with seed-level worst-case correctness aggregation. - """ + """Evaluate BPDA-PGD with seed-level worst-case correctness aggregation.""" vectors, accuracies = [], [] for seed in seeds: - torch.manual_seed(seed) + _seed_restart(seed) accuracy, vector = _run_bpda_once( - model, loader, eps, n_restarts, return_vector=True + model, + loader, + eps, + n_restarts, + alpha=alpha, + steps=steps, + return_vector=True, ) accuracies.append(accuracy) vectors.append(vector) - worst_vector = np.logical_and.reduce(vectors) - out = { - "BPDA_PGD": float(worst_vector.mean()), - "BPDA_PGD_mean": float(np.mean(accuracies)), - "BPDA_PGD_std": float(np.std(accuracies)), - } - if return_vector: - out["_vectors"] = {"BPDA_PGD": worst_vector} - return out + return _aggregate_restarts("BPDA_PGD", accuracies, vectors, return_vector) def adaptive_pgd_attack( @@ -912,7 +1009,9 @@ def run_boundary_attack( for x, y in loader: if total_seen >= max_images: break - x, y = x.to(device, non_blocking=NON_BLOCKING_TRANSFER), y.to(device, non_blocking=NON_BLOCKING_TRANSFER) + x, y = x.to(device, non_blocking=NON_BLOCKING_TRANSFER), y.to( + device, non_blocking=NON_BLOCKING_TRANSFER + ) with torch.no_grad(): pred = model(x).argmax(dim=1) for i in range(x.size(0)): @@ -999,7 +1098,9 @@ def random_noise_attack( correct, total, vectors = 0, 0, [] with torch.no_grad(): for x, y in loader: - x, y = x.to(device, non_blocking=NON_BLOCKING_TRANSFER), y.to(device, non_blocking=NON_BLOCKING_TRANSFER) + x, y = x.to(device, non_blocking=NON_BLOCKING_TRANSFER), y.to( + device, non_blocking=NON_BLOCKING_TRANSFER + ) worst_correct = torch.ones(y.size(0), dtype=torch.bool, device=device) for _ in range(n_restarts): noise = torch.empty_like(x).uniform_(-1, 1) * eps_normalized @@ -1021,44 +1122,35 @@ def run_random_noise_seeded( ): accuracies, vectors = [], [] for seed in seeds: + _seed_restart(seed) accuracy, vector = random_noise_attack( model, loader, eps=eps, seed=seed, return_vector=True ) accuracies.append(accuracy) vectors.append(vector) - worst_vector = np.logical_and.reduce(vectors) - out = { - "Random_Noise": float(worst_vector.mean()), - "Random_Noise_mean": float(np.mean(accuracies)), - "Random_Noise_std": float(np.std(accuracies)), - } - if return_vector: - out["_vectors"] = {"Random_Noise": worst_vector} - return out + return _aggregate_restarts("Random_Noise", accuracies, vectors, return_vector) def pgd_steps_ablation( model, loader, eps=DEFAULT_EPS, step_list=PGD_ABLATION_STEPS, use_ste=False ): - """Accuracy vs. PGD step count, at an explicit, chosen gradient regime. - """ + """Accuracy vs. PGD step count, at an explicit, chosen gradient regime.""" model.eval() out = {} - with ste_mode(model, use_ste): - for steps in step_list: - if steps == 0: - acc = random_noise_attack(model, loader, eps=eps, seed=0) - else: - pgd = make_torchattack( - torchattacks.PGD, - model, - eps=eps, - alpha=PGD_ALPHA, - steps=steps, - random_start=PGD_RANDOM_START, - ) - acc = accuracy_under_attack(model, loader, pgd) - out[steps] = acc + for steps in step_list: + if not use_ste: + result = run_pgd(model, loader, eps=eps, steps=steps, seeds=SEEDS) + out[steps] = result["PGD"] + else: + result = run_bpda( + model, + loader, + eps=eps, + n_restarts=1, + seeds=SEEDS, + steps=steps, + ) + out[steps] = result["BPDA_PGD"] return out @@ -1071,8 +1163,7 @@ def pgd_trajectory_diagnostics( max_batches=TRAJECTORY_MAX_BATCHES, use_ste=False, ): - """Per-step gradient norm and movement along a PGD trajectory. - """ + """Per-step gradient norm and movement along a PGD trajectory.""" model.eval() clip_min, clip_max = CLIP_MIN.to(device), CLIP_MAX.to(device) eps_normalized = eps / DATASET_STD.to(device) @@ -1084,7 +1175,9 @@ def pgd_trajectory_diagnostics( for bi, (x, y) in enumerate(loader): if bi >= max_batches: break - x, y = x.to(device, non_blocking=NON_BLOCKING_TRANSFER), y.to(device, non_blocking=NON_BLOCKING_TRANSFER) + x, y = x.to(device, non_blocking=NON_BLOCKING_TRANSFER), y.to( + device, non_blocking=NON_BLOCKING_TRANSFER + ) noise = torch.empty_like(x).uniform_(-1, 1) * eps_normalized x_start = torch.max(torch.min(x + noise, clip_max), clip_min).detach() x_adv = x_start.clone() @@ -1139,8 +1232,7 @@ def confidence_margin_diagnostic( max_batches=MARGIN_MAX_BATCHES, use_ste=False, ): - """Compare top-2 softmax margins clean vs. under PGD, at an explicit regime. - """ + """Compare top-2 softmax margins clean vs. under PGD, at an explicit regime.""" model.eval() with ste_mode(model, use_ste): pgd = make_torchattack( @@ -1155,7 +1247,9 @@ def confidence_margin_diagnostic( for bi, (x, y) in enumerate(loader): if bi >= max_batches: break - x, y = x.to(device, non_blocking=NON_BLOCKING_TRANSFER), y.to(device, non_blocking=NON_BLOCKING_TRANSFER) + x, y = x.to(device, non_blocking=NON_BLOCKING_TRANSFER), y.to( + device, non_blocking=NON_BLOCKING_TRANSFER + ) with torch.no_grad(): top2 = F.softmax(model(x), dim=1).topk(2, dim=1).values clean_margins.extend((top2[:, 0] - top2[:, 1]).cpu().tolist()) @@ -1169,21 +1263,10 @@ def confidence_margin_diagnostic( def run_epsilon_sweep_for_model( model, loader, name, epsilons, count_quant_layers_fn=None, safe_set=None ): - """ - `count_quant_layers_fn` and `safe_set` are injected by the caller so this - module does not need to import the quantization/report-plumbing code. - Falls back to local, dependency-free implementations if not provided. - - ``PGD_acc`` is explicitly computed under hard rounding (``use_ste=False``) - — this used to be *ambient*: whatever ``use_ste`` was left at by - ``run_suite``'s earlier attacks by the time the sweep ran, with no - guarantee it was False. For quantized models this is the number that, - if it tracks ``Random_Noise_acc`` instead of dropping as epsilon grows, - is showing you gradient masking, not robustness. - - ``BPDA_acc`` uses a restart budget (``n_restarts=1``) matched to - ``PGD_acc``. The sweep therefore reports only vanilla PGD and BPDA-PGD; - it does not execute or report a separate PGD-STE attack. + """Evaluate epsilon-sweep metrics with the main suite implementations. + + Shared metrics use the same seeded, per-example worst-case aggregation as + ``run_suite``, so matching model/epsilon rows are directly comparable. """ if safe_set is None: @@ -1200,41 +1283,43 @@ def safe_set(target, key, fn, warning, *, context=None, default=None): rows = [] for eps in epsilons: row = {"model": name, "epsilon": eps} + context = f"{name} eps={eps:.4f}" - def run_pgd_sweep(use_ste): - with ste_mode(model, use_ste): - pgd = make_torchattack( - torchattacks.PGD, - model, - eps=eps, - alpha=PGD_ALPHA, - steps=PGD_STEPS, - random_start=PGD_RANDOM_START, - ) - return accuracy_under_attack(model, loader, pgd) + def add_shared_metrics(result, source_name, output_name): + row[output_name] = result[source_name] if result is not None else None + row[f"{source_name}_mean"] = ( + result[f"{source_name}_mean"] if result is not None else None + ) + row[f"{source_name}_std"] = ( + result[f"{source_name}_std"] if result is not None else None + ) - context = f"{name} eps={eps:.4f}" - safe_set( - row, - "PGD_acc", - lambda: run_pgd_sweep(False), + pgd = safe_set( + {}, + "result", + lambda: run_pgd(model, loader, eps=eps, seeds=SEEDS), "PGD (hard-round) sweep failed", context=context, ) - safe_set( - row, - "Random_Noise_acc", - lambda: random_noise_attack(model, loader, eps=eps), + add_shared_metrics(pgd, "PGD", "PGD_acc") + + noise = safe_set( + {}, + "result", + lambda: run_random_noise_seeded(model, loader, eps=eps, seeds=SEEDS), "random_noise sweep failed", context=context, ) + add_shared_metrics(noise, "Random_Noise", "Random_Noise_acc") + if is_quant: - safe_set( - row, - "BPDA_acc", - lambda: _run_bpda_once(model, loader, eps=eps, n_restarts=1), + bpda = safe_set( + {}, + "result", + lambda: run_bpda(model, loader, eps=eps, n_restarts=1, seeds=SEEDS), "BPDA sweep failed", context=context, ) + add_shared_metrics(bpda, "BPDA_PGD", "BPDA_acc") rows.append(row) return rows diff --git a/src/combine.py b/src/combine.py index 5c55c76..7cd4d13 100644 --- a/src/combine.py +++ b/src/combine.py @@ -1,6 +1,5 @@ #!/usr/bin/env python3 -"""Recover partial QuantAdv outputs and regenerate report artifacts. -""" +"""Recover partial QuantAdv outputs and regenerate report artifacts.""" from __future__ import annotations @@ -17,7 +16,9 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument("--data-dir", type=Path, default=Path(report_data.DATA_DIR)) parser.add_argument("--no-plots", action="store_true") - parser.add_argument("--quiet", action="store_true", help="Do not print table sizes.") + parser.add_argument( + "--quiet", action="store_true", help="Do not print table sizes." + ) return parser.parse_args() diff --git a/src/config.py b/src/config.py index 2a02706..b340428 100644 --- a/src/config.py +++ b/src/config.py @@ -21,7 +21,9 @@ USE_AMP = torch.cuda.is_available() CONFIG_DIR = Path(__file__).resolve().parent -PROJECT_ROOT = CONFIG_DIR.parent if CONFIG_DIR.name in {"src", "src_old"} else CONFIG_DIR +PROJECT_ROOT = ( + CONFIG_DIR.parent if CONFIG_DIR.name in {"src", "src_old"} else CONFIG_DIR +) DATA_DIR = os.path.join(PROJECT_ROOT, "data") os.makedirs(DATA_DIR, exist_ok=True) PYTORCHCV_MODEL_DIR = os.path.join(PROJECT_ROOT, "models", "pytorchcv") @@ -39,12 +41,8 @@ DATASET_RANDOM_CROP_PADDING = 4 DATASET_MEAN_VALUES = (0.5071, 0.4867, 0.4408) DATASET_STD_VALUES = (0.2675, 0.2565, 0.2761) -DATASET_MEAN = torch.tensor(DATASET_MEAN_VALUES).view( - 1, DATASET_INPUT_CHANNELS, 1, 1 -) -DATASET_STD = torch.tensor(DATASET_STD_VALUES).view( - 1, DATASET_INPUT_CHANNELS, 1, 1 -) +DATASET_MEAN = torch.tensor(DATASET_MEAN_VALUES).view(1, DATASET_INPUT_CHANNELS, 1, 1) +DATASET_STD = torch.tensor(DATASET_STD_VALUES).view(1, DATASET_INPUT_CHANNELS, 1, 1) DATASET_TRAIN_TRANSFORM = tv_transforms.Compose( [ tv_transforms.RandomCrop( diff --git a/src/data.py b/src/data.py index 5bcfacf..7116394 100644 --- a/src/data.py +++ b/src/data.py @@ -197,14 +197,11 @@ def plot_summary_results(df_results: pd.DataFrame, output: Path = PLOT_PNG) -> N def plot_epsilon_sweep_curves( df_sweep: pd.DataFrame, output: Path = SWEEP_PLOT_PNG ) -> None: - """Plot robustness metrics as a function of attack epsilon. - """ + """Plot robustness metrics as a function of attack epsilon.""" if df_sweep is None or df_sweep.empty: return value_cols = [ - c - for c in ["PGD_acc", "Random_Noise_acc", "BPDA_acc"] - if c in df_sweep.columns + c for c in ["PGD_acc", "Random_Noise_acc", "BPDA_acc"] if c in df_sweep.columns ] if not value_cols: return @@ -390,9 +387,7 @@ def plot_component_ablation( return value_vars = [ - c - for c in ["clean_acc", "PGD_acc", "BPDA_acc"] - if c in df_component.columns + c for c in ["clean_acc", "PGD_acc", "BPDA_acc"] if c in df_component.columns ] df_long = df_component.melt( id_vars=["model", "config"], @@ -425,8 +420,7 @@ def plot_component_ablation( def plot_gradient_masking_summary( df_results: pd.DataFrame, output: Path = MASKING_SUMMARY_PLOT_PNG ) -> None: - """Plot direct evidence of gradient masking under hard-round quantization. - """ + """Plot direct evidence of gradient masking under hard-round quantization.""" if ( df_results is None or df_results.empty @@ -592,8 +586,7 @@ def plot_results_heatmap( def add_sweep_masking_metrics(df_sweep: pd.DataFrame) -> pd.DataFrame: - """Add an explicit masking-gap column to the epsilon-sweep table. - """ + """Add an explicit masking-gap column to the epsilon-sweep table.""" if df_sweep is None or df_sweep.empty: return df_sweep if {"PGD_acc", "BPDA_acc"}.issubset(df_sweep.columns): @@ -645,9 +638,7 @@ def add_derived_metrics(df: pd.DataFrame) -> pd.DataFrame: architecture = ( df["model"] .astype(str) - .str.replace( - r"_(FP32|int8_PTQ|int4_PTQ|int8_QAT|int4_QAT).*", "", regex=True - ) + .str.replace(r"_(FP32|int8_PTQ|int4_PTQ|int8_QAT|int4_QAT).*", "", regex=True) ) df["Architecture"] = architecture if "Worst_Robust_Acc" in df: @@ -666,9 +657,7 @@ def add_derived_metrics(df: pd.DataFrame) -> pd.DataFrame: def add_paired_tests(df: pd.DataFrame) -> pd.DataFrame: """Add paired McNemar tests between each model variant and FP32 baseline.""" - return qstats.add_paired_mcnemar_tests( - df, baseline_name=qstats.fp32_baseline_name - ) + return qstats.add_paired_mcnemar_tests(df, baseline_name=qstats.fp32_baseline_name) def _read_if_present(path: Path) -> pd.DataFrame: @@ -923,6 +912,7 @@ def combine_csv_family(data_dir: Path, family: CsvFamily) -> pd.DataFrame: _write(frame, output) return frame + CSV_FAMILIES = { "results": CsvFamily( RESULTS_CSV, @@ -969,6 +959,7 @@ def combine_csv_family(data_dir: Path, family: CsvFamily) -> pd.DataFrame: ), } + def combine_all(data_dir: Path) -> dict[str, pd.DataFrame]: """Combine all available experiment result families.""" data_dir = Path(data_dir) @@ -995,6 +986,7 @@ class PlotSpec: function: Callable[[pd.DataFrame, Path], None] output: str + PLOT_SPECS = ( PlotSpec("results", plot_summary_results, PLOT_PNG), PlotSpec("sweep", plot_epsilon_sweep_curves, SWEEP_PLOT_PNG), @@ -1013,11 +1005,14 @@ class PlotSpec: PlotSpec("defense", plot_defense_comparison, DEFENSE_PLOT_PNG), ) + def plot_all(dfs: dict[str, pd.DataFrame], output_dir: Path = DATA_DIR) -> None: """Generate all summary plots from combined result tables.""" output_dir = Path(output_dir) for spec in PLOT_SPECS: - spec.function(dfs.get(spec.table, pd.DataFrame()), output_dir / _name(spec.output)) + spec.function( + dfs.get(spec.table, pd.DataFrame()), output_dir / _name(spec.output) + ) overview_dir = output_dir / "visualizations" for name, frame in dfs.items(): diff --git a/src/defense.py b/src/defense.py index 9a9dcba..d7b1c85 100644 --- a/src/defense.py +++ b/src/defense.py @@ -88,11 +88,7 @@ def forward(self, x): else self.weight ) out = self._conv_forward(x, w, self.bias) - return ( - _quantize_tensor(out, self.bits, self.use_ste) - if self.quant_act - else out - ) + return _quantize_tensor(out, self.bits, self.use_ste) if self.quant_act else out class _QuantLinear(nn.Linear): @@ -106,11 +102,7 @@ def forward(self, x): else self.weight ) out = F.linear(x, w, self.bias) - return ( - _quantize_tensor(out, self.bits, self.use_ste) - if self.quant_act - else out - ) + return _quantize_tensor(out, self.bits, self.use_ste) if self.quant_act else out def _to_quant_module(module, bits): diff --git a/src/stats.py b/src/stats.py index ad61c0c..59c250e 100644 --- a/src/stats.py +++ b/src/stats.py @@ -69,7 +69,9 @@ def mcnemar_exact(vector_a, vector_b) -> dict[str, float]: a = _as_bool_vector(vector_a) b = _as_bool_vector(vector_b) if a.shape != b.shape: - raise ValueError(f"McNemar vectors must have equal shape, got {a.shape} and {b.shape}") + raise ValueError( + f"McNemar vectors must have equal shape, got {a.shape} and {b.shape}" + ) a_only = int(np.sum(a & ~b)) b_only = int(np.sum(~a & b)) @@ -128,7 +130,9 @@ def add_paired_mcnemar_tests( with np.load(baseline_path) as baseline_vectors, np.load( variant_path ) as variant_vectors: - for metric in sorted(set(baseline_vectors.files) & set(variant_vectors.files)): + for metric in sorted( + set(baseline_vectors.files) & set(variant_vectors.files) + ): a = np.asarray(baseline_vectors[metric], dtype=bool) b = np.asarray(variant_vectors[metric], dtype=bool) if a.shape != b.shape: From da4a79ecc608154128bdef3923b7fd36583eae2f Mon Sep 17 00:00:00 2001 From: Richard Date: Tue, 14 Jul 2026 14:12:30 +0800 Subject: [PATCH 4/4] TorchAO Critical STE (PGD) Fix | Refactor - TorchAO was incorrectly integrated earlier such that (A) layers were detached and (B) quantization differed between hard and BPDA (STE-smoothed) - quantization has now been factored out into its own file - upgrades to graphing utilities; moved into own package - committing "formaldata.py", used for condensed result plotting for paper - implemeted unified framework for caching, increased reuse of results - enabled a lot more attacks --- README.md | 7 +- src/QuantAdv.py | 656 +++++++----------------------------- src/attack.py | 376 +++++++++++++-------- src/attack_cache.py | 119 +++++++ src/config.py | 14 +- src/graphs/__init__.py | 5 + src/{ => graphs}/combine.py | 2 +- src/{ => graphs}/data.py | 4 +- src/graphs/formaldata.py | 383 +++++++++++++++++++++ src/quantization.py | 272 +++++++++++++++ 10 files changed, 1159 insertions(+), 679 deletions(-) create mode 100644 src/attack_cache.py create mode 100644 src/graphs/__init__.py rename src/{ => graphs}/combine.py (96%) rename src/{ => graphs}/data.py (99%) create mode 100644 src/graphs/formaldata.py create mode 100644 src/quantization.py diff --git a/README.md b/README.md index 728ce74..587b6f9 100644 --- a/README.md +++ b/README.md @@ -24,8 +24,11 @@ or set Download=true To run `python src/QuantAdv.py` -To graph incomplete results -`python src/combine.py` +To graph incomplete results +`python -m src.graphs.combine` + +To generate the consolidated formal figures +`python -m src.graphs.formaldata` Obsolete run `python src_old/launcher.py` diff --git a/src/QuantAdv.py b/src/QuantAdv.py index ab2e28e..67444f9 100644 --- a/src/QuantAdv.py +++ b/src/QuantAdv.py @@ -15,28 +15,26 @@ import time import threading import psutil -from torch.amp import autocast, GradScaler +from pathlib import Path + +PROJECT_ROOT = Path(__file__).resolve().parent.parent +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) -from torchao.quantization.qat import IntxFakeQuantizeConfig, IntxFakeQuantizer from pytorchcv.model_provider import get_model as ptcv_get_model import torchattacks from autoattack import AutoAttack import defense as dfn -import data as report_data +from src.graphs import data as report_data import stats as qstats -import sys -from pathlib import Path - -PROJECT_ROOT = Path(__file__).resolve().parent.parent -if str(PROJECT_ROOT) not in sys.path: - sys.path.insert(0, str(PROJECT_ROOT)) - from config import * +from quantization import * from ResourceMonitor import ResourceMonitor from attack import * +from attack_cache import AttackResultCache """QuantAdv experiment runner for quantization and adversarial robustness. @@ -143,467 +141,6 @@ def sanity_check_accuracy(model, loader): return accuracy_from_adv_fn(model, loader, use_autocast=True) -def _torchao_int_dtype(bits): - """Return TorchAO's integer dtype corresponding to the requested bit width.""" - if bits == 8: - return torch.int8 - dtype = getattr(torch, f"int{bits}", None) - if dtype is None: - raise RuntimeError( - f"This PyTorch build does not expose torch.int{bits}; upgrade PyTorch/TorchAO " - f"to use {bits}-bit TorchAO fake quantization." - ) - return dtype - - -def _make_torchao_fake_quantizer(bits, *, role): - """Build TorchAO fake quantization for a weight or activation tensor.""" - if role == "activation": - config = IntxFakeQuantizeConfig( - dtype=_torchao_int_dtype(bits), - granularity="per_token", - is_symmetric=False, - is_dynamic=True, - eps=QUANT_SCALE_MIN, - ) - elif role == "weight": - config = IntxFakeQuantizeConfig( - dtype=_torchao_int_dtype(bits), - granularity="per_channel", - is_symmetric=True, - is_dynamic=True, - eps=QUANT_SCALE_MIN, - ) - else: - raise ValueError(f"Unknown fake-quantizer role: {role!r}") - return IntxFakeQuantizer(config) - - -def _hard_or_ste(fake_quantizer, tensor, use_ste): - """Use TorchAO numerics while selecting STE or true hard-round gradients.""" - quantized = fake_quantizer(tensor) - return quantized if use_ste else quantized.detach() - - -def chaotic_sequence_like(t, seed=CHAOTIC_QUANT_SEED, map_name=CHAOTIC_QUANT_MAP): - """Create a deterministic chaotic sequence with the same shape as ``t``.""" - n = t.numel() - if n == 0: - return torch.empty_like(t) - dtype = torch.float32 if t.dtype in (torch.float16, torch.bfloat16) else t.dtype - idx = torch.arange(n, device=t.device, dtype=dtype) - z = torch.frac(seed + (idx + 1.0) * 0.6180339887498949) - z = torch.clamp(z, 1e-6, 1.0 - 1e-6) - if map_name == "tent": - for _ in range(CHAOTIC_QUANT_WARMUP): - z = torch.where( - z < 0.5, CHAOTIC_QUANT_MU * z * 0.5, CHAOTIC_QUANT_MU * (1.0 - z) * 0.5 - ) - else: - for _ in range(CHAOTIC_QUANT_WARMUP): - z = CHAOTIC_QUANT_R * z * (1.0 - z) - return z.view_as(t).to(dtype=t.dtype) - - -def chaotic_quantize_tensor( - t, - bits, - use_ste, - fake_quantizer, - quantize=True, - dither_amplitude=CHAOTIC_QUANT_DITHER, -): - """Apply subtractive chaotic dither around a TorchAO fake quantizer.""" - if bits is None or not quantize: - return t - chaos = chaotic_sequence_like(t) - # Dither is expressed in units of the tensor's dynamic range. TorchAO - # selects the actual scale/zero-point according to the configured scheme. - amplitude = t.detach().abs().amax().clamp_min(QUANT_SCALE_MIN) - dither = (chaos - 0.5) * dither_amplitude * amplitude - quantized = fake_quantizer(t + dither) - dither - return quantized if use_ste else quantized.detach() - - -class _QuantizedLayerMixin: - """Shared TorchAO fake-quantization controls for Conv2d/Linear wrappers.""" - - chaotic = False - - def _init_quantizers(self, bits): - self.bits = bits - self.weight_fake_quantizer = _make_torchao_fake_quantizer(bits, role="weight") - self.activation_fake_quantizer = _make_torchao_fake_quantizer( - bits, role="activation" - ) - - def _quant_params(self): - return ( - getattr(self, "bits", None), - getattr(self, "use_ste", QUANT_DEFAULT_USE_STE), - getattr(self, "quant_weight", QUANT_DEFAULT_WEIGHT), - getattr(self, "quant_act", QUANT_DEFAULT_ACT), - ) - - def _quantize_weight(self, tensor, enabled): - if not enabled: - return tensor - original_shape = tensor.shape - flattened = tensor.reshape(tensor.shape[0], -1) - if self.chaotic: - quantized = chaotic_quantize_tensor( - flattened, - self.bits, - self.use_ste, - self.weight_fake_quantizer, - True, - getattr(self, "dither_amplitude", CHAOTIC_QUANT_DITHER), - ) - else: - quantized = _hard_or_ste( - self.weight_fake_quantizer, flattened, self.use_ste - ) - return quantized.reshape(original_shape) - - def _quantize_activation(self, tensor, enabled): - if not enabled: - return tensor - if self.chaotic: - return chaotic_quantize_tensor( - tensor, - self.bits, - self.use_ste, - self.activation_fake_quantizer, - True, - getattr(self, "dither_amplitude", CHAOTIC_QUANT_DITHER), - ) - return _hard_or_ste(self.activation_fake_quantizer, tensor, self.use_ste) - - -class QuantConv2d(_QuantizedLayerMixin, nn.Conv2d): - """Conv2d using TorchAO fake quantization for weights and output activations.""" - - def forward(self, x): - _, _, quant_weight, quant_act = self._quant_params() - weight = self._quantize_weight(self.weight, quant_weight) - out = self._conv_forward(x, weight, self.bias) - return self._quantize_activation(out, quant_act) - - -class ChaoticQuantConv2d(QuantConv2d): - chaotic = True - - -class QuantLinear(_QuantizedLayerMixin, nn.Linear): - """Linear using TorchAO fake quantization for weights and output activations.""" - - def forward(self, x): - _, _, quant_weight, quant_act = self._quant_params() - weight = self._quantize_weight(self.weight, quant_weight) - out = F.linear(x, weight, self.bias) - return self._quantize_activation(out, quant_act) - - -class ChaoticQuantLinear(QuantLinear): - chaotic = True - - -def _to_quant_module( - mod, - bits, - quant_weight=True, - quant_act=True, - chaotic=False, - dither_amplitude=CHAOTIC_QUANT_DITHER, -): - """Convert Conv2d/Linear to TorchAO-backed fake-quantized wrappers.""" - if isinstance(mod, nn.Conv2d): - cls = ChaoticQuantConv2d if chaotic else QuantConv2d - new = cls( - mod.in_channels, - mod.out_channels, - mod.kernel_size, - mod.stride, - mod.padding, - mod.dilation, - mod.groups, - mod.bias is not None, - mod.padding_mode, - device=mod.weight.device, - dtype=mod.weight.dtype, - ) - elif isinstance(mod, nn.Linear): - cls = ChaoticQuantLinear if chaotic else QuantLinear - new = cls( - mod.in_features, - mod.out_features, - bias=mod.bias is not None, - device=mod.weight.device, - dtype=mod.weight.dtype, - ) - else: - return None - - new.weight = mod.weight - if mod.bias is not None: - new.bias = mod.bias - new._init_quantizers(bits) - new.use_ste = QUANT_DEFAULT_USE_STE - new.quant_weight = quant_weight - new.quant_act = quant_act - new.dither_amplitude = dither_amplitude - return new - - -def _replace_recursive( - module, - bits, - quant_weight=True, - quant_act=True, - chaotic=False, - dither_amplitude=CHAOTIC_QUANT_DITHER, -): - for name, child in list(module.named_children()): - replacement = _to_quant_module( - child, - bits, - quant_weight, - quant_act, - chaotic=chaotic, - dither_amplitude=dither_amplitude, - ) - if replacement is not None: - setattr(module, name, replacement) - else: - _replace_recursive( - child, - bits, - quant_weight, - quant_act, - chaotic=chaotic, - dither_amplitude=dither_amplitude, - ) - - -def convert_to_quant( - model, - bits, - quant_weight=True, - quant_act=True, - chaotic=False, - dither_amplitude=CHAOTIC_QUANT_DITHER, -): - """Return a deep-copied model backed by TorchAO fake quantizers.""" - m = copy.deepcopy(model) - _replace_recursive( - m, - bits, - quant_weight, - quant_act, - chaotic=chaotic, - dither_amplitude=dither_amplitude, - ) - return m - - -def convert_to_chaotic_quant( - model, - bits, - quant_weight=True, - quant_act=True, - dither_amplitude=CHAOTIC_QUANT_DITHER, -): - return convert_to_quant( - model, - bits, - quant_weight=quant_weight, - quant_act=quant_act, - chaotic=True, - dither_amplitude=dither_amplitude, - ) - - -def quantizable_layer_names(model): - quant_types = (QuantConv2d, QuantLinear, ChaoticQuantConv2d, ChaoticQuantLinear) - return [ - name - for name, mod in model.named_modules() - if isinstance(mod, (nn.Conv2d, nn.Linear)) and not isinstance(mod, quant_types) - ] - - -def set_child_module(root, module_name, new_module): - parts = module_name.split(".") - parent = root - for part in parts[:-1]: - parent = parent._modules[part] - parent._modules[parts[-1]] = new_module - - -def convert_layer_chunk_to_quant( - model, layer_names, bits, quant_weight=True, quant_act=True, chaotic=False -): - m = copy.deepcopy(model) - targets = set(layer_names) - for name, mod in list(m.named_modules()): - if name not in targets: - continue - new_mod = _to_quant_module(mod, bits, quant_weight, quant_act, chaotic=chaotic) - if new_mod is not None: - set_child_module(m, name, new_mod) - return m.to(device).eval() - - -def quant_layer_chunks(layer_names, n_chunks): - if not layer_names: - return [] - n_chunks = max(1, min(n_chunks, len(layer_names))) - return [ - list(chunk) - for chunk in np.array_split(np.array(layer_names, dtype=object), n_chunks) - if len(chunk) > 0 - ] - - -def count_quant_layers(model): - return sum( - isinstance( - mod, (QuantConv2d, QuantLinear, ChaoticQuantConv2d, ChaoticQuantLinear) - ) - for mod in model.modules() - ) - - -def verify_quantization_layers( - arch_key, fp32_model, quant_model, label, fp32_layer_names=None -): - fp32_layer_names = ( - quantizable_layer_names(fp32_model) - if fp32_layer_names is None - else fp32_layer_names - ) - if not fp32_layer_names: - raise RuntimeError(f"{arch_key} exposes zero FP32 nn.Conv2d/nn.Linear layers.") - quant_count = count_quant_layers(quant_model) - threshold = int(np.ceil(0.8 * len(fp32_layer_names))) - print(f" {label} quantized layers: {quant_count}", flush=True) - if quant_count < threshold: - raise RuntimeError( - f"{arch_key} {label} replaced {quant_count}/{len(fp32_layer_names)} " - f"quantizable layers; expected at least {threshold}." - ) - return quant_count - - -def set_quant_components(model, quant_weight, quant_act): - for mod in model.modules(): - if isinstance( - mod, (QuantConv2d, QuantLinear, ChaoticQuantConv2d, ChaoticQuantLinear) - ): - mod.quant_weight = quant_weight - mod.quant_act = quant_act - - -def prepare_qat( - fp32_model, - bits, - finetune_loader, - epochs=QAT_EPOCHS_DEFAULT, - lr=QAT_LR, - chaotic=False, -): - """Fine-tune a TorchAO fake-quantized copy with STE enabled.""" - m = ( - convert_to_chaotic_quant(fp32_model, bits, quant_weight=True, quant_act=True) - if chaotic - else convert_to_quant(fp32_model, bits, quant_weight=True, quant_act=True) - ) - if torch.cuda.device_count() > 1: - m = nn.DataParallel(m) - set_ste_mode(m, True) - m.train() - opt = torch.optim.SGD( - m.parameters(), - lr=lr, - momentum=QAT_MOMENTUM, - weight_decay=QAT_WEIGHT_DECAY, - ) - scaler = GradScaler(device=device.type, enabled=device.type == "cuda") - for epoch in range(epochs): - running = 0.0 - for batch_idx, (x, y) in enumerate(finetune_loader): - x = x.to(device, non_blocking=NON_BLOCKING_TRANSFER) - y = y.to(device, non_blocking=NON_BLOCKING_TRANSFER) - opt.zero_grad(set_to_none=True) - with autocast(device_type=device.type, enabled=device.type == "cuda"): - loss = F.cross_entropy(m(x), y) - scaler.scale(loss).backward() - scaler.step(opt) - scaler.update() - running += loss.item() - if batch_idx == 0 or (batch_idx + 1) % QAT_LOG_EVERY_BATCHES == 0: - print( - f" QAT epoch {epoch + 1}/{epochs} batch " - f"{batch_idx + 1}/{len(finetune_loader)} loss {loss.item():.4f}", - flush=True, - ) - print( - f" QAT epoch {epoch + 1}/{epochs} avg loss " - f"{running / len(finetune_loader):.4f}", - flush=True, - ) - set_ste_mode(m, False) - return m.eval() - - -class ImageCompressionSTE(torch.autograd.Function): - """Quantize pixel values with identity gradients for compression defenses.""" - - @staticmethod - def forward(ctx, x, bits): - """Apply image compression in the forward pass for input preprocessing.""" - levels = float(2**bits - 1) - return torch.round(x * levels).div(levels) - - @staticmethod - def backward(ctx, grad_output): - """Pass gradients through image compression unchanged.""" - return grad_output, None - - -class CompressedInputModel(nn.Module): - """Wrap a classifier with resize and pixel-bit-depth input compression.""" - - def __init__( - self, - model, - size=COMPRESS_IMAGE_SIZE, - bits=COMPRESS_IMAGE_BITS, - mode=COMPRESS_IMAGE_MODE, - ): - """Wrap a model with differentiable image-compression preprocessing.""" - super().__init__() - self.model = model - self.size = size - self.bits = bits - self.mode = mode - - def forward(self, x): - """Compress inputs before forwarding them to the wrapped model.""" - pixels = denormalize_inputs(x).clamp(0.0, 1.0) - if self.size and self.size < pixels.shape[-1]: - kwargs = {"mode": self.mode} - if self.mode in ("linear", "bilinear", "bicubic", "trilinear"): - kwargs["align_corners"] = COMPRESS_IMAGE_ALIGN_CORNERS - pixels = F.interpolate(pixels, size=(self.size, self.size), **kwargs) - pixels = F.interpolate( - pixels, size=(DATASET_IMAGE_SIZE, DATASET_IMAGE_SIZE), **kwargs - ) - if self.bits is not None: - pixels = ImageCompressionSTE.apply(pixels, self.bits) - return self.model(normalize_pixels(pixels.clamp(0.0, 1.0))) - - def with_image_compression( model, size=COMPRESS_IMAGE_SIZE, bits=COMPRESS_IMAGE_BITS, mode=COMPRESS_IMAGE_MODE ): @@ -709,11 +246,60 @@ def hook(module, grad_input, grad_output): } +def component_ablation_row( + name, label, quant_weight, quant_act, clean_acc, pgd, bpda, frac_zero_grad_hard +): + """Build one consistently shaped quantization-component ablation row.""" + pgd_acc = pgd.get("PGD") + bpda_acc = bpda.get("BPDA_PGD") + return { + "model": name, + "config": label, + "quant_weight": quant_weight, + "quant_act": quant_act, + "clean_acc": clean_acc, + "PGD_acc": pgd_acc, + "PGD_hard_acc": pgd_acc, + "PGD_mean": pgd.get("PGD_mean"), + "PGD_std": pgd.get("PGD_std"), + "BPDA_acc": bpda_acc, + "BPDA_mean": bpda.get("BPDA_PGD_mean"), + "BPDA_std": bpda.get("BPDA_PGD_std"), + "PGD_minus_BPDA": ( + pgd_acc - bpda_acc + if pgd_acc is not None and bpda_acc is not None + else None + ), + "frac_zero_grad_hard": frac_zero_grad_hard, + } + + +def main_both_component_ablation_row(name, results): + """Reuse main-suite metrics for the already evaluated fully quantized model.""" + return component_ablation_row( + name, + "both", + True, + True, + results.get("clean_acc"), + { + key: results.get(key) + for key in ("PGD", "PGD_mean", "PGD_std") + }, + { + key: results.get(key) + for key in ("BPDA_PGD", "BPDA_PGD_mean", "BPDA_PGD_std") + }, + results.get("frac_zero_grad_hard"), + ) + + def run_quant_component_ablation(model, loader, name, eps=DEFAULT_EPS): """Evaluate which quantized components are responsible for observed effects. - The rows distinguish three interpretations of "testing quantized": - weight-only quantization, activation-only quantization, and both together. + The final artifact distinguishes weight-only, activation-only, and fully + quantized models. This function computes only the first two; the ``both`` + row is populated from the main suite rather than attacked a second time. For each interpretation the experiment reports ordinary hard-round PGD and a budget-matched BPDA-PGD. A large PGD vs BPDA-PGD gap is evidence of gradient masking, not evidence that quantization itself improves @@ -726,57 +312,39 @@ def run_quant_component_ablation(model, loader, name, eps=DEFAULT_EPS): rate and a small PGD-vs-BPDA gap. Only ``act_only``/``both`` quantize a tensor that ``x`` actually flows through, so only those can mask. """ - configs = [ - ("weight_only", True, False), - ("act_only", False, True), - ("both", True, True), - ] + configs = [("weight_only", True, False), ("act_only", False, True)] rows = [] - for label, qw, qa in configs: - set_quant_components(model, qw, qa) - clean_acc = sanity_check_accuracy(model, loader) + try: + for label, qw, qa in configs: + set_quant_components(model, qw, qa) + clean_acc = sanity_check_accuracy(model, loader) - pgd_hard = run_pgd(model, loader, eps=eps, seeds=SEEDS) - pgd_hard_acc = pgd_hard["PGD"] - with ste_mode(model, False): - x, y = next(iter(loader)) - x, y = x.to(device), y.to(device) - x_in = x.clone().requires_grad_(True) - loss = F.cross_entropy(model(x_in), y) - g_hard = torch.autograd.grad(loss, x_in)[0].flatten() - frac_zero = (g_hard.abs() < GRAD_ZERO_THRESHOLD).float().mean().item() - - bpda = run_bpda( - model, - loader, - eps=eps, - n_restarts=1, - seeds=SEEDS, - ) - rows.append( - { - "model": name, - "config": label, - "quant_weight": qw, - "quant_act": qa, - "clean_acc": clean_acc, - "PGD_acc": pgd_hard_acc, - "PGD_hard_acc": pgd_hard_acc, - "PGD_mean": pgd_hard["PGD_mean"], - "PGD_std": pgd_hard["PGD_std"], - "BPDA_acc": bpda["BPDA_PGD"], - "BPDA_mean": bpda["BPDA_PGD_mean"], - "BPDA_std": bpda["BPDA_PGD_std"], - "PGD_minus_BPDA": ( - pgd_hard_acc - bpda["BPDA_PGD"] - if pgd_hard_acc is not None and bpda["BPDA_PGD"] is not None - else None - ), - "frac_zero_grad_hard": frac_zero, - } - ) - # restore original (both quantized) state - set_quant_components(model, True, True) + pgd_hard = run_pgd(model, loader, eps=eps, seeds=SEEDS) + with ste_mode(model, False): + x, y = next(iter(loader)) + x, y = x.to(device), y.to(device) + x_in = x.clone().requires_grad_(True) + loss = F.cross_entropy(model(x_in), y) + g_hard = torch.autograd.grad(loss, x_in)[0].flatten() + frac_zero = ( + (g_hard.abs() < GRAD_ZERO_THRESHOLD).float().mean().item() + ) + + bpda = run_bpda( + model, + loader, + eps=eps, + n_restarts=1, + seeds=SEEDS, + ) + rows.append( + component_ablation_row( + name, label, qw, qa, clean_acc, pgd_hard, bpda, frac_zero + ) + ) + finally: + # Never leave the shared registry model in an ablated state after failure. + set_quant_components(model, True, True) return rows @@ -893,7 +461,9 @@ def save_json(path, data, *, indent=None): json.dump(data, handle, indent=indent) -def run_suite(model, loader, name, fp32_ref=None, eps=DEFAULT_EPS): +def run_suite( + model, loader, name, fp32_ref=None, eps=DEFAULT_EPS, attack_cache=None +): """Run the full attack and diagnostic suite for one model.""" model.eval() results = {"model": name} @@ -912,7 +482,8 @@ def run_suite(model, loader, name, fp32_ref=None, eps=DEFAULT_EPS): results, vectors, lambda: run_fgsm_pgd( - model, loader, eps=eps, return_vectors=True, use_ste=False + model, loader, eps=eps, return_vectors=True, use_ste=False, + cache=attack_cache, ), "FGSM/PGD failed", context=name, @@ -963,6 +534,7 @@ def run_suite(model, loader, name, fp32_ref=None, eps=DEFAULT_EPS): eps=eps, return_vector=True, use_ste=False, + cache=attack_cache, ), "transfer_attack failed", context=name, @@ -978,6 +550,7 @@ def run_suite(model, loader, name, fp32_ref=None, eps=DEFAULT_EPS): eps=eps, return_vector=True, use_ste=False, + cache=attack_cache, ), "MIM transfer_attack failed", context=name, @@ -986,7 +559,9 @@ def run_suite(model, loader, name, fp32_ref=None, eps=DEFAULT_EPS): safe_set( results, "UAP_Transfer", - lambda: transfer_uap_attack(fp32_ref, model, loader, eps=eps), + lambda: transfer_uap_attack( + fp32_ref, model, loader, eps=eps, cache=attack_cache + ), "UAP transfer_attack failed", context=name, ) @@ -1019,7 +594,9 @@ def run_suite(model, loader, name, fp32_ref=None, eps=DEFAULT_EPS): safe_update_vectors( results, vectors, - lambda: run_random_noise_seeded(model, loader, eps=eps, return_vector=True), + lambda: run_random_noise_seeded( + model, loader, eps=eps, return_vector=True, cache=attack_cache + ), "random_noise_attack failed", context=name, defaults={"Random_Noise": None}, @@ -1076,6 +653,7 @@ def save_pgd_ablation(): eps=eps, n_restarts=1, return_vector=True, + cache=attack_cache, ), "BPDA failed", context=name, @@ -1177,6 +755,7 @@ def save_layerwise_profile(): def save_component_ablation(): """Persist quantization component-ablation diagnostics for the current model.""" rows = run_quant_component_ablation(model, loader, name, eps=eps) + rows.append(main_both_component_ablation_row(name, results)) pd.DataFrame(rows).to_csv( csv_path(name, "component_ablation"), index=False ) @@ -1216,7 +795,9 @@ def save_confidence_margins(): return results -def run_epsilon_sweep_for_model_wrapped(model, loader, name, epsilons): +def run_epsilon_sweep_for_model_wrapped( + model, loader, name, epsilons, attack_cache=None +): """Run and annotate an epsilon sweep for one named model.""" return run_epsilon_sweep_for_model( model, @@ -1225,6 +806,7 @@ def run_epsilon_sweep_for_model_wrapped(model, loader, name, epsilons): epsilons, count_quant_layers_fn=count_quant_layers, safe_set=safe_set, + cache=attack_cache, ) @@ -1696,6 +1278,8 @@ def main(): df_sweep = pd.DataFrame() sweep_done = set() + attack_cache = AttackResultCache() + def run_pending_epsilon_sweep(name, model): """Run an epsilon sweep for a model that still needs one.""" nonlocal df_sweep, sweep_done @@ -1714,7 +1298,7 @@ def save_epsilon_sweep(): """Persist epsilon-sweep diagnostics for pending models.""" nonlocal df_sweep, sweep_done rows = run_epsilon_sweep_for_model_wrapped( - model, eval_loader, name, pending_eps + model, eval_loader, name, pending_eps, attack_cache=attack_cache ) if rows: new_sweep = pd.DataFrame(rows) @@ -1740,9 +1324,21 @@ def save_epsilon_sweep(): try: if RECORD_RUN_METRICS: with monitor: - res = run_suite(model, eval_loader, name, fp32_ref=ref) + res = run_suite( + model, + eval_loader, + name, + fp32_ref=ref, + attack_cache=attack_cache, + ) else: - res = run_suite(model, eval_loader, name, fp32_ref=ref) + res = run_suite( + model, + eval_loader, + name, + fp32_ref=ref, + attack_cache=attack_cache, + ) except Exception as e: print(f" [FAIL] run_suite failed for {name}: {e}") traceback.print_exc() diff --git a/src/attack.py b/src/attack.py index 211edf7..e479e85 100644 --- a/src/attack.py +++ b/src/attack.py @@ -12,7 +12,6 @@ import random import warnings -from contextlib import contextmanager import numpy as np import torch @@ -24,6 +23,8 @@ import defense as dfn from torch.amp import autocast from config import * +from quantization import get_ste_mode, set_ste_mode, ste_mode +from attack_cache import AttackKey, AttackResultCache def unwrap_model(model): @@ -31,51 +32,6 @@ def unwrap_model(model): return model.module if isinstance(model, nn.DataParallel) else model -def set_ste_mode(model, flag): - """Toggle straight-through gradients on fake-quantized modules.""" - toggled = 0 - for mod in model.modules(): - if hasattr(mod, "use_ste") and hasattr(mod, "bits"): - mod.use_ste = flag - toggled += 1 - return toggled - - -def get_ste_mode(model): - """Return the current ``use_ste`` state of a model's quantized modules.""" - values = { - mod.use_ste - for mod in model.modules() - if hasattr(mod, "use_ste") and hasattr(mod, "bits") - } - if not values: - return None - if len(values) > 1: - return "mixed" - return next(iter(values)) - - -@contextmanager -def ste_mode(model, flag): - """Run a block with an explicit, restored ``use_ste`` state. - - Every attack or diagnostic that cares whether gradients are computed - through the true (masking) rounding op or through the straight-through - bypass should wrap its gradient computation in this, e.g.:: - with ste_mode(model, False): # real hard-round gradient - with ste_mode(model, True): # STE / BPDA-style bypass - """ - previous = get_ste_mode(model) - set_ste_mode(model, flag) - try: - yield - finally: - # Restore exactly what was there before (True/False), or leave - # everything at False if the model had no quantized layers / no - # prior state to speak of. - set_ste_mode(model, previous if isinstance(previous, bool) else False) - - def normalize_pixels(x): """Normalize pixel-space tensors with the configured dataset constants.""" return (x - DATASET_MEAN.to(x.device)) / DATASET_STD.to(x.device) @@ -171,6 +127,49 @@ def adv_fn(x, y): ) +def _cache_adversarial_batches(loader, adv_fn): + """Generate adversarial examples once and retain device-independent batches.""" + batches = [] + for x, y in loader: + x = x.to(device, non_blocking=NON_BLOCKING_TRANSFER) + y = y.to(device, non_blocking=NON_BLOCKING_TRANSFER) + x_adv = adv_fn(x, y) + batches.append((x_adv.detach().cpu(), y.detach().cpu())) + return batches + + +def _evaluate_cached_batches(target_model, batches, return_vector=False): + """Evaluate a target model on previously generated adversarial batches.""" + correct, total = 0, 0 + correct_vectors = [] + target_model.eval() + for x_adv, y in batches: + x_adv = x_adv.to(device, non_blocking=NON_BLOCKING_TRANSFER) + y = y.to(device, non_blocking=NON_BLOCKING_TRANSFER) + with torch.no_grad(): + batch_correct = target_model(x_adv).argmax(dim=1) == y + correct += batch_correct.sum().item() + total += y.size(0) + if return_vector: + correct_vectors.append(batch_correct.detach().cpu()) + accuracy = correct / total if total else None + if not return_vector: + return accuracy + vector = ( + torch.cat(correct_vectors).numpy().astype(bool) + if correct_vectors + else np.empty(0, dtype=bool) + ) + return accuracy, vector + + +def _cached_transfer_batches(cache, key, loader, adv_fn): + """Return cached examples for a transfer attack, generating them on a miss.""" + if cache is None: + return _cache_adversarial_batches(loader, adv_fn) + return cache.get_or_compute(key, lambda: _cache_adversarial_batches(loader, adv_fn)) + + def evaluate_normalized_attack(model, loader, attack_fn): """Evaluate an attack function that consumes and returns normalized tensors.""" model.eval() @@ -216,6 +215,16 @@ def _aggregate_restarts(name, accuracies, vectors, return_vector=False): return out +def _copy_cached_result(result, return_vector): + """Return a caller-owned result dictionary without mutating cached state.""" + copied = dict(result) + if "_vectors" in copied: + copied["_vectors"] = dict(copied["_vectors"]) + if not return_vector: + copied.pop("_vectors", None) + return copied + + def _run_seeded_pgd( model, loader, @@ -227,27 +236,33 @@ def _run_seeded_pgd( seeds, use_ste, return_vector, + cache=None, ): """Single implementation for seeded torchattacks PGD evaluation.""" - model.eval() - accuracies, vectors = [], [] - with ste_mode(model, use_ste): - for seed in seeds: - _seed_restart(seed) - attack = make_torchattack( - torchattacks.PGD, - model, - eps=eps, - alpha=alpha, - steps=steps, - random_start=random_start, - ) - accuracy, vector = accuracy_under_attack( - model, loader, attack, return_vector=True - ) - accuracies.append(accuracy) - vectors.append(vector) - return _aggregate_restarts("PGD", accuracies, vectors, return_vector) + key = AttackKey.create( + model_id=id(unwrap_model(model)), attack_name="PGD_RESULT", loader=loader, + epsilon=eps, alpha=alpha, steps=steps, seeds=seeds, + restarts=len(tuple(seeds)), use_ste=use_ste, + attack_parameters={"random_start": random_start}, + ) + def compute(): + model.eval() + accuracies, vectors = [], [] + with ste_mode(model, use_ste): + for seed in seeds: + _seed_restart(seed) + attack = make_torchattack( + torchattacks.PGD, model, eps=eps, alpha=alpha, steps=steps, + random_start=random_start, + ) + accuracy, vector = accuracy_under_attack( + model, loader, attack, return_vector=True + ) + accuracies.append(accuracy) + vectors.append(vector) + return _aggregate_restarts("PGD", accuracies, vectors, True) + result = compute() if cache is None else cache.get_or_compute(key, compute) + return _copy_cached_result(result, return_vector) def run_pgd( @@ -259,6 +274,7 @@ def run_pgd( random_start=PGD_RANDOM_START, seeds=SEEDS, return_vector=False, + cache=None, ): """Evaluate hard-round PGD with the canonical seeded restart policy. @@ -276,11 +292,13 @@ def run_pgd( seeds=seeds, use_ste=False, return_vector=return_vector, + cache=cache, ) def run_fgsm_pgd( - model, loader, eps=DEFAULT_EPS, seeds=SEEDS, return_vectors=False, use_ste=False + model, loader, eps=DEFAULT_EPS, seeds=SEEDS, return_vectors=False, + use_ste=False, cache=None ): """Run FGSM and multi-restart PGD, aggregating PGD by worst-case correctness. @@ -303,7 +321,8 @@ def adv_fn(x, y): out["FGSM"] = fgsm_acc pgd_out = ( run_pgd( - model, loader, eps=eps, seeds=seeds, return_vector=return_vectors + model, loader, eps=eps, seeds=seeds, return_vector=return_vectors, + cache=cache, ) if not use_ste else _run_seeded_pgd( @@ -316,6 +335,7 @@ def adv_fn(x, y): seeds=seeds, use_ste=True, return_vector=return_vectors, + cache=cache, ) ) pgd_vectors = pgd_out.pop("_vectors", {}) @@ -386,24 +406,40 @@ def transfer_attack( eps=DEFAULT_EPS, return_vector=False, use_ste=False, + cache=None, ): - """Generate PGD examples on ``source_model`` and evaluate ``target_model``.""" - with ste_mode(source_model, use_ste): - pgd = make_torchattack( - torchattacks.PGD, - source_model, - eps=eps, - alpha=PGD_ALPHA, - steps=PGD_STEPS, - random_start=PGD_RANDOM_START, - ) - return accuracy_under_attack( - source_model, - loader, - pgd, - target_model=target_model, - return_vector=return_vector, - ) + """Generate/cache FP32-source PGD examples and evaluate ``target_model``.""" + key = AttackKey.create( + model_id=id(unwrap_model(source_model)), + attack_name="PGD_TRANSFER_ARTIFACTS", + loader=loader, + epsilon=eps, + alpha=PGD_ALPHA, + steps=PGD_STEPS, + use_ste=use_ste, + attack_parameters={"random_start": PGD_RANDOM_START}, + ) + if cache is None: + with ste_mode(source_model, use_ste): + pgd = make_torchattack( + torchattacks.PGD, + source_model, + eps=eps, + alpha=PGD_ALPHA, + steps=PGD_STEPS, + random_start=PGD_RANDOM_START, + ) + batches = _cached_transfer_batches(cache, key, loader, pgd) + else: + def generate(): + with ste_mode(source_model, use_ste): + pgd = make_torchattack( + torchattacks.PGD, source_model, eps=eps, alpha=PGD_ALPHA, + steps=PGD_STEPS, random_start=PGD_RANDOM_START, + ) + return _cache_adversarial_batches(loader, pgd) + batches = cache.get_or_compute(key, generate) + return _evaluate_cached_batches(target_model, batches, return_vector) def transfer_attack_mim( @@ -413,23 +449,39 @@ def transfer_attack_mim( eps=DEFAULT_EPS, return_vector=False, use_ste=False, + cache=None, ): - with ste_mode(source_model, use_ste): - mim = make_torchattack( - torchattacks.MIFGSM, - source_model, - eps=eps, - alpha=PGD_ALPHA, - steps=PGD_STEPS, - decay=MIFGSM_DECAY, - ) - return accuracy_under_attack( - source_model, - loader, - mim, - target_model=target_model, - return_vector=return_vector, - ) + key = AttackKey.create( + model_id=id(unwrap_model(source_model)), + attack_name="MIM_TRANSFER_ARTIFACTS", + loader=loader, + epsilon=eps, + alpha=PGD_ALPHA, + steps=PGD_STEPS, + use_ste=use_ste, + attack_parameters={"decay": MIFGSM_DECAY}, + ) + if cache is None: + with ste_mode(source_model, use_ste): + mim = make_torchattack( + torchattacks.MIFGSM, + source_model, + eps=eps, + alpha=PGD_ALPHA, + steps=PGD_STEPS, + decay=MIFGSM_DECAY, + ) + batches = _cached_transfer_batches(cache, key, loader, mim) + else: + def generate(): + with ste_mode(source_model, use_ste): + mim = make_torchattack( + torchattacks.MIFGSM, source_model, eps=eps, alpha=PGD_ALPHA, + steps=PGD_STEPS, decay=MIFGSM_DECAY, + ) + return _cache_adversarial_batches(loader, mim) + batches = cache.get_or_compute(key, generate) + return _evaluate_cached_batches(target_model, batches, return_vector) def build_uap( @@ -487,15 +539,48 @@ def adv_fn(x, y): def run_uap_attack(model, loader, eps=DEFAULT_EPS, max_images=UAP_MAX_IMAGES): - return _run_uap_attack(model, model, loader, eps=eps, max_images=max_images) + with ste_mode(model, False): + return _run_uap_attack(model, model, loader, eps=eps, max_images=max_images) def transfer_uap_attack( - source_model, target_model, loader, eps=DEFAULT_EPS, max_images=UAP_MAX_IMAGES + source_model, + target_model, + loader, + eps=DEFAULT_EPS, + max_images=UAP_MAX_IMAGES, + cache=None, ): - return _run_uap_attack( - source_model, target_model, loader, eps=eps, max_images=max_images - ) + with ste_mode(source_model, False): + key = AttackKey.create( + model_id=id(unwrap_model(source_model)), + attack_name="UAP_TRANSFER_ARTIFACTS", + loader=loader, + epsilon=eps, + attack_parameters={ + "delta": UAP_DELTA, "max_iter": UAP_MAX_ITER, + "deepfool_steps": UAP_DEEPFOOL_STEPS, + "overshoot": UAP_OVERSHOOT, "max_images": max_images, + }, + ) + if cache is None: + v = build_uap(source_model, loader, eps=eps, max_images=max_images) + clip_min, clip_max = CLIP_MIN.to(device), CLIP_MAX.to(device) + + def apply_uap(x, _y): + return torch.max(torch.min(x + v, clip_max), clip_min) + + batches = _cached_transfer_batches(cache, key, loader, apply_uap) + else: + def generate(): + v = build_uap(source_model, loader, eps=eps, max_images=max_images) + clip_min, clip_max = CLIP_MIN.to(device), CLIP_MAX.to(device) + return _cache_adversarial_batches( + loader, + lambda x, _y: torch.max(torch.min(x + v, clip_max), clip_min), + ) + batches = cache.get_or_compute(key, generate) + return _evaluate_cached_batches(target_model, batches) def projected_pgd_attack(x, y, eps, alpha, steps, grad_fn): @@ -573,23 +658,27 @@ def run_bpda( return_vector=False, alpha=PGD_ALPHA, steps=PGD_STEPS, + cache=None, ): """Evaluate BPDA-PGD with seed-level worst-case correctness aggregation.""" - vectors, accuracies = [], [] - for seed in seeds: - _seed_restart(seed) - accuracy, vector = _run_bpda_once( - model, - loader, - eps, - n_restarts, - alpha=alpha, - steps=steps, - return_vector=True, - ) - accuracies.append(accuracy) - vectors.append(vector) - return _aggregate_restarts("BPDA_PGD", accuracies, vectors, return_vector) + key = AttackKey.create( + model_id=id(unwrap_model(model)), attack_name="BPDA_PGD_RESULT", loader=loader, + epsilon=eps, alpha=alpha, steps=steps, seeds=seeds, + restarts=n_restarts, use_ste=True, + ) + def compute(): + vectors, accuracies = [], [] + for seed in seeds: + _seed_restart(seed) + accuracy, vector = _run_bpda_once( + model, loader, eps, n_restarts, alpha=alpha, steps=steps, + return_vector=True, + ) + accuracies.append(accuracy) + vectors.append(vector) + return _aggregate_restarts("BPDA_PGD", accuracies, vectors, True) + result = compute() if cache is None else cache.get_or_compute(key, compute) + return _copy_cached_result(result, return_vector) def adaptive_pgd_attack( @@ -644,7 +733,7 @@ def loss_fn(x_adv, labels): loss = loss + F.cross_entropy(base_model(noisy), labels) return loss / eot_samples - return adaptive_pgd_attack(model, x, y, loss_fn, eps=eps) + return adaptive_pgd_attack(model, x, y, loss_fn, eps=eps, use_ste=True) return {"EOT_PGD": evaluate_normalized_attack(model, loader, attack)} @@ -664,7 +753,7 @@ def loss_fn(x_adv, labels): ).mean() return F.cross_entropy(logits, labels) - lam * guardrail_penalty - return adaptive_pgd_attack(model, x, y, loss_fn, eps=eps) + return adaptive_pgd_attack(model, x, y, loss_fn, eps=eps, use_ste=True) return {"Adaptive_Guardrail": evaluate_normalized_attack(model, loader, attack)} @@ -681,7 +770,7 @@ def loss_fn(x_adv, labels): detector_penalty = F.cross_entropy(defense_model.detector(x_adv), benign) return F.cross_entropy(logits, labels) - lam * detector_penalty - return adaptive_pgd_attack(model, x, y, loss_fn, eps=eps) + return adaptive_pgd_attack(model, x, y, loss_fn, eps=eps, use_ste=True) return {"Adaptive_DetectGuard": evaluate_normalized_attack(model, loader, attack)} @@ -1118,17 +1207,25 @@ def random_noise_attack( def run_random_noise_seeded( - model, loader, eps=DEFAULT_EPS, seeds=SEEDS, return_vector=False + model, loader, eps=DEFAULT_EPS, seeds=SEEDS, return_vector=False, cache=None ): - accuracies, vectors = [], [] - for seed in seeds: - _seed_restart(seed) - accuracy, vector = random_noise_attack( - model, loader, eps=eps, seed=seed, return_vector=True - ) - accuracies.append(accuracy) - vectors.append(vector) - return _aggregate_restarts("Random_Noise", accuracies, vectors, return_vector) + key = AttackKey.create( + model_id=id(unwrap_model(model)), attack_name="RANDOM_NOISE_RESULT", + loader=loader, epsilon=eps, seeds=seeds, restarts=len(tuple(seeds)), + attack_parameters={"distribution": "uniform_linf"}, + ) + def compute(): + accuracies, vectors = [], [] + for seed in seeds: + _seed_restart(seed) + accuracy, vector = random_noise_attack( + model, loader, eps=eps, seed=seed, return_vector=True + ) + accuracies.append(accuracy) + vectors.append(vector) + return _aggregate_restarts("Random_Noise", accuracies, vectors, True) + result = compute() if cache is None else cache.get_or_compute(key, compute) + return _copy_cached_result(result, return_vector) def pgd_steps_ablation( @@ -1261,7 +1358,8 @@ def confidence_margin_diagnostic( def run_epsilon_sweep_for_model( - model, loader, name, epsilons, count_quant_layers_fn=None, safe_set=None + model, loader, name, epsilons, count_quant_layers_fn=None, safe_set=None, + cache=None, ): """Evaluate epsilon-sweep metrics with the main suite implementations. @@ -1297,7 +1395,7 @@ def add_shared_metrics(result, source_name, output_name): pgd = safe_set( {}, "result", - lambda: run_pgd(model, loader, eps=eps, seeds=SEEDS), + lambda: run_pgd(model, loader, eps=eps, seeds=SEEDS, cache=cache), "PGD (hard-round) sweep failed", context=context, ) @@ -1306,7 +1404,9 @@ def add_shared_metrics(result, source_name, output_name): noise = safe_set( {}, "result", - lambda: run_random_noise_seeded(model, loader, eps=eps, seeds=SEEDS), + lambda: run_random_noise_seeded( + model, loader, eps=eps, seeds=SEEDS, cache=cache + ), "random_noise sweep failed", context=context, ) @@ -1316,7 +1416,9 @@ def add_shared_metrics(result, source_name, output_name): bpda = safe_set( {}, "result", - lambda: run_bpda(model, loader, eps=eps, n_restarts=1, seeds=SEEDS), + lambda: run_bpda( + model, loader, eps=eps, n_restarts=1, seeds=SEEDS, cache=cache + ), "BPDA sweep failed", context=context, ) diff --git a/src/attack_cache.py b/src/attack_cache.py new file mode 100644 index 0000000..d39eb70 --- /dev/null +++ b/src/attack_cache.py @@ -0,0 +1,119 @@ +"""Typed, configuration-complete cache keys and storage for attack artifacts.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Hashable, Mapping, Sequence + +from torch.utils.data import Subset + + +def _freeze(value: Any) -> Hashable: + """Convert nested attack parameters to a deterministic hashable value.""" + if isinstance(value, Mapping): + return tuple(sorted((str(key), _freeze(item)) for key, item in value.items())) + if isinstance(value, (list, tuple)): + return tuple(_freeze(item) for item in value) + if isinstance(value, set): + return tuple(sorted(_freeze(item) for item in value)) + if value is None or isinstance(value, (str, int, float, bool, bytes)): + return value + return repr(value) + + +def loader_dataset_indices(loader) -> tuple[int, ...]: + """Return the underlying dataset indices represented by a data loader.""" + dataset = loader.dataset + indices: Sequence[int] = range(len(dataset)) + while isinstance(dataset, Subset): + parent_indices = dataset.indices + indices = tuple(int(parent_indices[index]) for index in indices) + dataset = dataset.dataset + return tuple(int(index) for index in indices) + + +@dataclass(frozen=True) +class AttackKey: + """Identity of an attack computation whose result can be reused safely.""" + + model_id: int | str + attack_name: str + epsilon: float | None = None + alpha: float | None = None + steps: int | None = None + seeds: tuple[int, ...] = () + restarts: int = 1 + use_ste: bool = False + dataset_indices: tuple[int, ...] = () + attack_parameters: tuple[tuple[str, Hashable], ...] = () + + @classmethod + def create( + cls, + *, + model_id: int | str, + attack_name: str, + loader, + epsilon: float | None = None, + alpha: float | None = None, + steps: int | None = None, + seeds: Sequence[int] = (), + restarts: int = 1, + use_ste: bool = False, + attack_parameters: Mapping[str, Any] | None = None, + ) -> "AttackKey": + parameters = { + "batch_size": getattr(loader, "batch_size", None), + "drop_last": getattr(loader, "drop_last", False), + "dataset_type": type(loader.dataset).__qualname__, + **(attack_parameters or {}), + } + return cls( + model_id=model_id, + attack_name=str(attack_name), + epsilon=None if epsilon is None else float(epsilon), + alpha=None if alpha is None else float(alpha), + steps=None if steps is None else int(steps), + seeds=tuple(int(seed) for seed in seeds), + restarts=int(restarts), + use_ste=bool(use_ste), + dataset_indices=loader_dataset_indices(loader), + attack_parameters=tuple( + sorted((str(name), _freeze(value)) for name, value in parameters.items()) + ), + ) + + +@dataclass +class AttackResultCache: + """In-memory cache for attack results or generated adversarial artifacts.""" + + _values: dict[AttackKey, Any] = field(default_factory=dict) + hits: int = 0 + misses: int = 0 + + def get_or_compute(self, key: AttackKey, compute: Callable[[], Any]) -> Any: + if key in self._values: + self.hits += 1 + return self._values[key] + self.misses += 1 + value = compute() + self._values[key] = value + return value + + def get(self, key: AttackKey, default=None): + if key in self._values: + self.hits += 1 + return self._values[key] + self.misses += 1 + return default + + def store(self, key: AttackKey, value: Any) -> Any: + self._values[key] = value + return value + + def clear(self) -> None: + self._values.clear() + + def __len__(self) -> int: + return len(self._values) diff --git a/src/config.py b/src/config.py index b340428..971db30 100644 --- a/src/config.py +++ b/src/config.py @@ -209,17 +209,17 @@ # Paper-run allocation. Expensive, weakly informative analyses remain available # but are off by default so the core attacks can use more images/steps/restarts. RUN_CHAOTIC_COMPRESS = False -RUN_EXTRA_WHITEBOX_ATTACKS = False +RUN_EXTRA_WHITEBOX_ATTACKS = True RUN_UAP_ATTACKS = False -RUN_SURROGATE_ATTACK = False -RUN_REVERSE_TRANSFERS = False +RUN_SURROGATE_ATTACK = True +RUN_REVERSE_TRANSFERS = True RUN_BOUNDARY_ATTACK = False RUN_NES_ATTACK = False RUN_DEFENSE_SUITE = False -RUN_CHUNK_QUANTIZATION = False -RUN_PGD_TRAJECTORY = False -RUN_LAYERWISE_PROFILE = False -RUN_CONFIDENCE_MARGIN = False +RUN_CHUNK_QUANTIZATION = True +RUN_PGD_TRAJECTORY = True +RUN_LAYERWISE_PROFILE = True +RUN_CONFIDENCE_MARGIN = True RUN_COMPONENT_ABLATION = True RUN_EPSILON_SWEEP = True RUN_CHAOTIC_DITHER_SWEEP = False diff --git a/src/graphs/__init__.py b/src/graphs/__init__.py new file mode 100644 index 0000000..f7a5b9b --- /dev/null +++ b/src/graphs/__init__.py @@ -0,0 +1,5 @@ +"""Result collation and visualization tools for QuantAdv. + +Run command-line utilities from the repository root with ``python -m``; for +example, ``python -m src.graphs.combine``. +""" diff --git a/src/combine.py b/src/graphs/combine.py similarity index 96% rename from src/combine.py rename to src/graphs/combine.py index 7cd4d13..b343429 100644 --- a/src/combine.py +++ b/src/graphs/combine.py @@ -6,7 +6,7 @@ import argparse from pathlib import Path -import data as report_data +from . import data as report_data def parse_args() -> argparse.Namespace: diff --git a/src/data.py b/src/graphs/data.py similarity index 99% rename from src/data.py rename to src/graphs/data.py index 7116394..6442930 100644 --- a/src/data.py +++ b/src/graphs/data.py @@ -25,8 +25,8 @@ import matplotlib.pyplot as plt import seaborn as sns -from config import * -import stats as qstats +from ..config import * +from .. import stats as qstats def _model_from_prefixed_path(path: Path, prefix: str, suffix: str) -> str: diff --git a/src/graphs/formaldata.py b/src/graphs/formaldata.py new file mode 100644 index 0000000..902a88e --- /dev/null +++ b/src/graphs/formaldata.py @@ -0,0 +1,383 @@ +#!/usr/bin/env python3 +"""Build consolidated formal figures from every QuantAdv result table. + +Combination and derived-stat logic stays in :mod:`data`. This module groups +related metrics and creates one figure per group, always comparing every model +and quantization for which that group is available. +""" + +from __future__ import annotations + +import math +import re +from dataclasses import dataclass +from pathlib import Path +from typing import Callable, Iterable + +import matplotlib + +matplotlib.use("Agg") +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns + +from . import data as report_data +from ..config import DATA_DIR, PLOT_BBOX_INCHES, PLOT_DPI + + +MetricPredicate = Callable[[str, pd.DataFrame], bool] + + +@dataclass(frozen=True) +class MetricGroup: + """A semantic group rendered as one consolidated figure.""" + + name: str + title: str + predicate: MetricPredicate + + +@dataclass(frozen=True) +class TableSpec: + """Plot structure and grouping rules for one combined artifact table.""" + + groups: tuple[MetricGroup, ...] + x: str | None = None + dimensions: tuple[str, ...] = () + excluded: tuple[str, ...] = () + distribution: bool = False + + +def _contains(*tokens: str) -> MetricPredicate: + return lambda column, _frame: any(token in column.lower() for token in tokens) + + +def _starts(*prefixes: str) -> MetricPredicate: + return lambda column, _frame: column.lower().startswith(prefixes) + + +def _ends(*suffixes: str) -> MetricPredicate: + return lambda column, _frame: column.lower().endswith(suffixes) + + +def _bounded(column: str, frame: pd.DataFrame) -> bool: + values = pd.to_numeric(frame[column], errors="coerce").dropna() + return bool(len(values) and values.between(0, 1).all()) + + +def _accuracy_metric(column: str, frame: pd.DataFrame) -> bool: + """Recognize attack/clean accuracies without a brittle attack-name list.""" + lower = column.lower() + if lower.endswith(("_std", "_pm", "_low", "_high", "_n", "_correct")): + return False + if any( + token in lower + for token in ("mcnemar", "grad", "plateau", "masking", "retention") + ): + return False + bases_with_vectors = { + name[: -len("_n")].lower() + for name in frame.columns + if name.lower().endswith("_n") + } + base = re.sub(r"_(mean|acc)$", "", lower) + return ( + "acc" in lower + or base in bases_with_vectors + or ("boundary" not in lower and _bounded(column, frame)) + ) + + +def _uncertainty_metric(column: str, _frame: pd.DataFrame) -> bool: + return bool( + re.search(r"_(n|correct|std|wilson_(low|high|pm)|cp_(low|high|pm))$", column) + ) + + +RESULT_GROUPS = ( + MetricGroup("accuracy", "Clean and Adversarial Accuracy", _accuracy_metric), + MetricGroup("uncertainty", "Accuracy Uncertainty and Sample Counts", _uncertainty_metric), + MetricGroup("robustness", "Derived Robustness and Masking Scores", _contains("robustness", "robust_acc", "masking_score", "masking_gap")), + MetricGroup("gradients", "Gradient and Staircase Diagnostics", _contains("grad", "plateau", "staircase")), + MetricGroup("boundary", "Decision-Boundary Diagnostics", _starts("boundary_")), + MetricGroup("mcnemar", "Paired McNemar Comparisons", _starts("mcnemar_")), +) + + +TABLE_SPECS: dict[str, TableSpec] = { + "results": TableSpec(RESULT_GROUPS), + "sweep": TableSpec( + ( + MetricGroup("accuracy", "Accuracy Across Perturbation Budgets", _contains("acc", "mean")), + MetricGroup("variability", "Sweep Variability", _ends("_std")), + MetricGroup("masking_gap", "PGD–BPDA Gap Across Budgets", _contains("minus", "gap")), + ), + x="epsilon", + ), + "ablation": TableSpec( + (MetricGroup("accuracy", "Accuracy Across Attack Steps", _accuracy_metric),), + x="steps", + dimensions=("attack",), + ), + "trajectory": TableSpec( + (MetricGroup("trajectory", "PGD Trajectory Diagnostics", lambda _c, _f: True),), + x="step", + ), + "layerwise": TableSpec( + (MetricGroup("gradients", "Layerwise Gradient Propagation", lambda _c, _f: True),), + x="layer", + ), + "component_ablation": TableSpec( + (MetricGroup("accuracy", "Quantization Component Ablation", _accuracy_metric),), + dimensions=("config",), + ), + "performance": TableSpec( + ( + MetricGroup("time", "Runtime and CPU Cost", _contains("second", "cpu_core")), + MetricGroup("memory", "CPU and CUDA Memory", _contains("mib")), + MetricGroup("model_size", "Model Size and Precision", _contains("parameter", "weight_bits")), + ) + ), + "dither": TableSpec( + (MetricGroup("accuracy", "Chaotic-Dither Accuracy", _accuracy_metric),), + x="dither_amplitude", + dimensions=("bits",), + ), + "chunk_quant": TableSpec( + (MetricGroup("accuracy", "Chunk-Quantization Accuracy", _accuracy_metric),), + x="chunk_id", + dimensions=("layers",), + ), + "margins": TableSpec( + (MetricGroup("margins", "Clean and Adversarial Confidence Margins", _contains("margin")),), + dimensions=("kind",), + excluded=("index",), + distribution=True, + ), + "defense": TableSpec( + (MetricGroup("accuracy_rates", "Defense Accuracy and Detection Rates", _accuracy_metric),), + dimensions=("defense",), + ), +} + + +def _numeric_metrics(frame: pd.DataFrame, spec: TableSpec) -> list[str]: + excluded = {spec.x, *spec.excluded} + return [ + column + for column in frame.columns + if column not in excluded + and pd.to_numeric(frame[column], errors="coerce").notna().any() + and not pd.api.types.is_bool_dtype(frame[column]) + ] + + +def group_metrics(frame: pd.DataFrame, spec: TableSpec) -> dict[str, list[str]]: + """Assign every numeric metric to its first semantic group or ``other``.""" + available = _numeric_metrics(frame, spec) + remaining = set(available) + grouped: dict[str, list[str]] = {} + for group in spec.groups: + matches = [c for c in available if c in remaining and group.predicate(c, frame)] + if matches: + grouped[group.name] = matches + remaining.difference_update(matches) + if remaining: + grouped["other"] = [c for c in available if c in remaining] + return grouped + + +def _labels(frame: pd.DataFrame, dimensions: Iterable[str]) -> pd.Series: + columns = [c for c in ("model", *dimensions) if c in frame.columns] + if not columns: + return pd.Series("All", index=frame.index) + return frame[columns].fillna("NA").astype(str).agg(" | ".join, axis=1) + + +def _long(frame: pd.DataFrame, metrics: list[str], spec: TableSpec) -> pd.DataFrame: + prepared = frame.copy() + prepared["Series"] = _labels(prepared, spec.dimensions) + ids = ["Series", *([spec.x] if spec.x and spec.x in prepared else [])] + long = prepared.melt(ids, metrics, "Metric", "Value") + long["Value"] = pd.to_numeric(long["Value"], errors="coerce") + return long.dropna(subset=["Value"]) + + +def _style_axes(ax: plt.Axes) -> None: + ax.grid(axis="y", linestyle="--", alpha=0.3) + ax.tick_params(axis="x", rotation=45) + for label in ax.get_xticklabels(): + label.set_horizontalalignment("right") + + +def _plot_grouped_bars(ax: plt.Axes, long: pd.DataFrame) -> None: + sns.barplot(data=long, x="Series", y="Value", hue="Metric", errorbar=None, ax=ax) + _style_axes(ax) + + +def _plot_lines(ax: plt.Axes, long: pd.DataFrame, x: str) -> None: + long = long.copy() + long["Series / Metric"] = long["Series"] + " | " + long["Metric"] + sns.lineplot( + data=long, + x=x, + y="Value", + hue="Series / Metric", + marker="o", + estimator=None, + ax=ax, + ) + ax.grid(linestyle="--", alpha=0.3) + + +def _plot_distribution(ax: plt.Axes, long: pd.DataFrame) -> None: + sns.boxplot(data=long, x="Series", y="Value", hue="Metric", ax=ax) + _style_axes(ax) + + +def _plot_heatmap(long: pd.DataFrame, title: str) -> plt.Figure: + """Render a dense metric group without losing models or metric columns.""" + matrix = long.pivot_table( + index="Series", columns="Metric", values="Value", aggfunc="mean", sort=False + ) + width = max(12, min(30, 0.34 * len(matrix.columns) + 6)) + height = max(5, min(18, 0.45 * len(matrix.index) + 3)) + fig, ax = plt.subplots(figsize=(width, height)) + sns.heatmap(matrix, cmap="viridis", mask=matrix.isna(), ax=ax) + ax.set_title(title) + ax.set_xlabel("Statistic") + ax.set_ylabel("") + ax.tick_params(axis="x", rotation=60) + return fig + + +def _needs_subplots(long: pd.DataFrame) -> bool: + """Use small multiples when metric magnitudes would make bars unreadable.""" + # Accuracy, rates, proportions, and signed gains share a meaningful scale + # even when some attacks drive accuracy very close to zero. + if long["Value"].between(-1, 1).all(): + return False + scales = long.groupby("Metric")["Value"].apply(lambda s: s.abs().max()).replace(0, np.nan) + return len(scales) > 1 and scales.max() / scales.min() > 100 + + +def _plot_small_multiples(long: pd.DataFrame, title: str) -> plt.Figure: + metrics = list(long["Metric"].drop_duplicates()) + columns = min(3, len(metrics)) + rows = math.ceil(len(metrics) / columns) + fig, axes = plt.subplots(rows, columns, figsize=(6 * columns, 4.5 * rows), squeeze=False) + for ax, metric in zip(axes.flat, metrics): + subset = long[long["Metric"] == metric] + sns.barplot(data=subset, x="Series", y="Value", errorbar=None, ax=ax, color="#4C72B0") + ax.set_title(metric.replace("_", " ")) + _style_axes(ax) + for ax in axes.flat[len(metrics) :]: + ax.axis("off") + fig.suptitle(title) + return fig + + +def plot_metric_group( + frame: pd.DataFrame, + metrics: list[str], + spec: TableSpec, + title: str, + output: Path, +) -> None: + """Render all metrics and all models for one semantic group in one figure.""" + long = _long(frame, metrics, spec) + if long.empty: + return + if not spec.x and not spec.distribution and len(metrics) > 20: + fig = _plot_heatmap(long, title) + elif not spec.x and not spec.distribution and _needs_subplots(long): + fig = _plot_small_multiples(long, title) + else: + width = max(10, min(24, 5 + 0.45 * long["Series"].nunique())) + fig, ax = plt.subplots(figsize=(width, 6.5)) + if spec.distribution: + _plot_distribution(ax, long) + elif spec.x and spec.x in long: + _plot_lines(ax, long, spec.x) + else: + _plot_grouped_bars(ax, long) + ax.set_title(title) + ax.set_ylabel("Value") + ax.legend(fontsize=8, loc="best") + fig.tight_layout() + output.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(output, dpi=PLOT_DPI, bbox_inches=PLOT_BBOX_INCHES) + plt.close(fig) + + +def plot_table_groups(frame: pd.DataFrame, table: str, root: Path) -> list[Path]: + if frame is None or frame.empty: + return [] + spec = TABLE_SPECS.get( + table, + TableSpec((MetricGroup("all", table.replace("_", " ").title(), lambda _c, _f: True),)), + ) + groups = group_metrics(frame, spec) + titles = {group.name: group.title for group in spec.groups} + outputs = [] + for name, metrics in groups.items(): + output = root / table / f"{name}.png" + plot_metric_group( + frame, + metrics, + spec, + titles.get(name, f"{table.replace('_', ' ').title()}: Other Diagnostics"), + output, + ) + outputs.append(output) + return outputs + + +def plot_all(tables: dict[str, pd.DataFrame], output_dir: Path = DATA_DIR) -> list[Path]: + """Create one consolidated figure per semantic group.""" + root = Path(output_dir) / "formal_figures" + outputs = [ + output + for table, frame in tables.items() + for output in plot_table_groups(frame, table, root) + ] + # Remove obsolete per-stat figures created by older versions of this module. + expected = {path.resolve() for path in outputs} + for path in root.rglob("*.png") if root.exists() else (): + if path.resolve() not in expected: + path.unlink() + return outputs + + +def print_report(tables: dict[str, pd.DataFrame], plots: Iterable[Path] = ()) -> None: + rows = [] + for table, frame in tables.items(): + spec = TABLE_SPECS.get(table) + rows.append( + { + "table": table, + "rows": len(frame), + "figure_groups": len(group_metrics(frame, spec)) if spec and not frame.empty else 0, + } + ) + print(pd.DataFrame(rows).to_string(index=False)) + print(f"generated {len(list(plots))} consolidated formal figures") + + +def combine_all(data_dir: Path = DATA_DIR) -> dict[str, pd.DataFrame]: + return report_data.combine_all(Path(data_dir)) + + +def generate_reports( + data_dir: Path = DATA_DIR, *, plots: bool = True, summary: bool = True +) -> dict[str, pd.DataFrame]: + tables = combine_all(Path(data_dir)) + outputs = plot_all(tables, Path(data_dir)) if plots else [] + if summary: + print_report(tables, outputs) + return tables + + +if __name__ == "__main__": + generate_reports() diff --git a/src/quantization.py b/src/quantization.py new file mode 100644 index 0000000..cf8af03 --- /dev/null +++ b/src/quantization.py @@ -0,0 +1,272 @@ +"""Fake-quantization layers, conversion helpers, and hard/STE mode controls.""" + +import copy +from contextlib import contextmanager + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.amp import GradScaler, autocast +from torchao.quantization.qat import IntxFakeQuantizeConfig, IntxFakeQuantizer + +from config import * + +WEIGHT = "weight" +ACTIVATION = "activation" + + +def _dtype(bits): + dtype = torch.int8 if bits == 8 else getattr(torch, f"int{bits}", None) + if dtype is None: + raise RuntimeError(f"This PyTorch build does not expose torch.int{bits}") + return dtype + + +def make_fake_quantizer(bits, role): + if role not in (WEIGHT, ACTIVATION): + raise ValueError(f"Unknown quantizer role: {role!r}") + return IntxFakeQuantizer( + IntxFakeQuantizeConfig( + dtype=_dtype(bits), + granularity="per_channel" if role == WEIGHT else "per_token", + is_symmetric=role == WEIGHT, + is_dynamic=True, + eps=QUANT_SCALE_MIN, + ) + ) + + +def hard_fake_quantize(tensor, bits, role): + """Match TorchAO forward numerics using native hard ``torch.round``.""" + qmin, qmax = -(2 ** (bits - 1)), 2 ** (bits - 1) - 1 + with torch.no_grad(): + zero = tensor.new_zeros(()) + lo = tensor.amin(dim=-1, keepdim=True).minimum(zero) + hi = tensor.amax(dim=-1, keepdim=True).maximum(zero) + if role == WEIGHT: + scale = (2 * torch.maximum(-lo, hi) / (qmax - qmin)).clamp_min( + QUANT_SCALE_MIN + ) + zero_point = torch.zeros_like(scale, dtype=torch.int32) + elif role == ACTIVATION: + scale = ((hi - lo) / (qmax - qmin)).clamp_min(QUANT_SCALE_MIN) + zero_point = (qmin - torch.round(lo / scale)).clamp(qmin, qmax).int() + else: + raise ValueError(f"Unknown quantizer role: {role!r}") + integer = torch.clamp( + torch.round(tensor * (1.0 / scale)) + zero_point, qmin, qmax + ) + return (integer - zero_point) * scale + + +def chaotic_sequence_like(tensor, seed=CHAOTIC_QUANT_SEED, map_name=CHAOTIC_QUANT_MAP): + if tensor.numel() == 0: + return torch.empty_like(tensor) + dtype = torch.float32 if tensor.dtype in (torch.float16, torch.bfloat16) else tensor.dtype + z = torch.frac(seed + (torch.arange(tensor.numel(), device=tensor.device, dtype=dtype) + 1) * 0.6180339887498949) + z = z.clamp(1e-6, 1 - 1e-6) + for _ in range(CHAOTIC_QUANT_WARMUP): + z = (torch.where(z < .5, CHAOTIC_QUANT_MU * z * .5, CHAOTIC_QUANT_MU * (1 - z) * .5) + if map_name == "tent" else CHAOTIC_QUANT_R * z * (1 - z)) + return z.view_as(tensor).to(tensor.dtype) + + +def fake_quantize(tensor, bits, role, quantizer, use_ste, dither_amplitude=None): + """Uniform hard/STE entry point, optionally with subtractive chaotic dither.""" + dither = 0 + if dither_amplitude is not None: + with torch.no_grad(): + amplitude = tensor.abs().amax().clamp_min(QUANT_SCALE_MIN) + dither = (chaotic_sequence_like(tensor) - .5) * dither_amplitude * amplitude + value = tensor + dither + quantized = quantizer(value) if use_ste else hard_fake_quantize(value, bits, role) + return quantized - dither + + +class _QuantizedMixin: + chaotic = False + + def _init_quantization(self, bits, quant_weight, quant_act, dither_amplitude): + self.bits, self.use_ste = bits, QUANT_DEFAULT_USE_STE + self.quant_weight, self.quant_act = quant_weight, quant_act + self.dither_amplitude = dither_amplitude + self.weight_fake_quantizer = make_fake_quantizer(bits, WEIGHT) + self.activation_fake_quantizer = make_fake_quantizer(bits, ACTIVATION) + + def _quantize(self, tensor, role, enabled): + if not enabled: + return tensor + shape = tensor.shape + value = tensor.reshape(tensor.shape[0], -1) if role == WEIGHT else tensor + quantizer = self.weight_fake_quantizer if role == WEIGHT else self.activation_fake_quantizer + result = fake_quantize( + value, self.bits, role, quantizer, self.use_ste, + self.dither_amplitude if self.chaotic else None, + ) + return result.reshape(shape) + + +class QuantConv2d(_QuantizedMixin, nn.Conv2d): + def forward(self, x): + weight = self._quantize(self.weight, WEIGHT, self.quant_weight) + return self._quantize(self._conv_forward(x, weight, self.bias), ACTIVATION, self.quant_act) + + +class ChaoticQuantConv2d(QuantConv2d): + chaotic = True + + +class QuantLinear(_QuantizedMixin, nn.Linear): + def forward(self, x): + weight = self._quantize(self.weight, WEIGHT, self.quant_weight) + return self._quantize(F.linear(x, weight, self.bias), ACTIVATION, self.quant_act) + + +class ChaoticQuantLinear(QuantLinear): + chaotic = True + + +QUANT_TYPES = (QuantConv2d, QuantLinear, ChaoticQuantConv2d, ChaoticQuantLinear) + + +def _convert_module(module, bits, quant_weight, quant_act, chaotic, dither_amplitude): + if isinstance(module, nn.Conv2d): + cls = ChaoticQuantConv2d if chaotic else QuantConv2d + new = cls(module.in_channels, module.out_channels, module.kernel_size, + module.stride, module.padding, module.dilation, module.groups, + module.bias is not None, module.padding_mode, + device=module.weight.device, dtype=module.weight.dtype) + elif isinstance(module, nn.Linear): + cls = ChaoticQuantLinear if chaotic else QuantLinear + new = cls(module.in_features, module.out_features, module.bias is not None, + device=module.weight.device, dtype=module.weight.dtype) + else: + return None + new.weight = module.weight + if module.bias is not None: + new.bias = module.bias + new._init_quantization(bits, quant_weight, quant_act, dither_amplitude) + return new + + +def _replace(module, bits, quant_weight, quant_act, chaotic, dither_amplitude): + for name, child in list(module.named_children()): + replacement = _convert_module(child, bits, quant_weight, quant_act, chaotic, dither_amplitude) + if replacement is None: + _replace(child, bits, quant_weight, quant_act, chaotic, dither_amplitude) + else: + setattr(module, name, replacement) + + +def convert_to_quant(model, bits, quant_weight=True, quant_act=True, chaotic=False, + dither_amplitude=CHAOTIC_QUANT_DITHER): + result = copy.deepcopy(model) + _replace(result, bits, quant_weight, quant_act, chaotic, dither_amplitude) + return result + + +def convert_to_chaotic_quant(model, bits, quant_weight=True, quant_act=True, + dither_amplitude=CHAOTIC_QUANT_DITHER): + return convert_to_quant(model, bits, quant_weight, quant_act, True, dither_amplitude) + + +def quantizable_layer_names(model): + return [name for name, mod in model.named_modules() + if isinstance(mod, (nn.Conv2d, nn.Linear)) and not isinstance(mod, QUANT_TYPES)] + + +def convert_layer_chunk_to_quant(model, layer_names, bits, quant_weight=True, + quant_act=True, chaotic=False): + result, targets = copy.deepcopy(model), set(layer_names) + for name, module in list(result.named_modules()): + if name not in targets: + continue + replacement = _convert_module(module, bits, quant_weight, quant_act, chaotic, CHAOTIC_QUANT_DITHER) + if replacement is not None: + parent_name, _, child_name = name.rpartition(".") + parent = result.get_submodule(parent_name) if parent_name else result + setattr(parent, child_name, replacement) + return result.to(device).eval() + + +def quant_layer_chunks(layer_names, n_chunks): + if not layer_names: + return [] + return [list(x) for x in np.array_split(layer_names, min(max(1, n_chunks), len(layer_names))) if len(x)] + + +def count_quant_layers(model): + return sum(isinstance(module, QUANT_TYPES) for module in model.modules()) + + +def verify_quantization_layers(arch_key, fp32_model, quant_model, label, fp32_layer_names=None): + names = quantizable_layer_names(fp32_model) if fp32_layer_names is None else fp32_layer_names + if not names: + raise RuntimeError(f"{arch_key} exposes zero quantizable layers") + count, threshold = count_quant_layers(quant_model), int(np.ceil(.8 * len(names))) + print(f" {label} quantized layers: {count}", flush=True) + if count < threshold: + raise RuntimeError(f"{arch_key} {label} replaced {count}/{len(names)} layers; expected at least {threshold}") + return count + + +def set_quant_components(model, quant_weight, quant_act): + for module in model.modules(): + if isinstance(module, QUANT_TYPES): + module.quant_weight, module.quant_act = quant_weight, quant_act + + +def prepare_qat(fp32_model, bits, finetune_loader, epochs=QAT_EPOCHS_DEFAULT, + lr=QAT_LR, chaotic=False): + """Fine-tune a quantized copy through TorchAO's STE path.""" + model = convert_to_quant(fp32_model, bits, chaotic=chaotic).to(device) + if torch.cuda.device_count() > 1: + model = nn.DataParallel(model) + set_ste_mode(model, True) + model.train() + optimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=QAT_MOMENTUM, + weight_decay=QAT_WEIGHT_DECAY) + scaler = GradScaler(device=device.type, enabled=device.type == "cuda") + for epoch in range(epochs): + running = 0.0 + for batch_idx, (x, y) in enumerate(finetune_loader): + x = x.to(device, non_blocking=NON_BLOCKING_TRANSFER) + y = y.to(device, non_blocking=NON_BLOCKING_TRANSFER) + optimizer.zero_grad(set_to_none=True) + with autocast(device_type=device.type, enabled=device.type == "cuda"): + loss = F.cross_entropy(model(x), y) + scaler.scale(loss).backward() + scaler.step(optimizer) + scaler.update() + running += loss.item() + if batch_idx == 0 or (batch_idx + 1) % QAT_LOG_EVERY_BATCHES == 0: + print(f" QAT epoch {epoch + 1}/{epochs} batch {batch_idx + 1}/{len(finetune_loader)} loss {loss.item():.4f}", flush=True) + print(f" QAT epoch {epoch + 1}/{epochs} avg loss {running / len(finetune_loader):.4f}", flush=True) + set_ste_mode(model, False) + return model.eval() + + +def set_ste_mode(model, flag): + # Keep the mode protocol structural so defense wrappers and any external + # quantized modules exposing the same controls participate as well. + modules = [m for m in model.modules() if hasattr(m, "bits") and hasattr(m, "use_ste")] + for module in modules: + module.use_ste = flag + return len(modules) + + +def get_ste_mode(model): + values = {m.use_ste for m in model.modules() + if hasattr(m, "bits") and hasattr(m, "use_ste")} + return None if not values else next(iter(values)) if len(values) == 1 else "mixed" + + +@contextmanager +def ste_mode(model, flag): + previous = get_ste_mode(model) + set_ste_mode(model, flag) + try: + yield + finally: + set_ste_mode(model, previous if isinstance(previous, bool) else False)