diff --git a/src/QuantAdv.py b/src/QuantAdv.py index 36fd025..d124784 100644 --- a/src/QuantAdv.py +++ b/src/QuantAdv.py @@ -14,7 +14,6 @@ import traceback import sys import warnings -import math import time import threading import psutil @@ -27,6 +26,7 @@ import defense as dfn import data as report_data +import stats as qstats import sys from pathlib import Path @@ -39,24 +39,33 @@ from ResourceMonitor import ResourceMonitor from attack import * -""" -Actively developed analysis of metrics of attacks per model and eplison +"""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, +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. """ def csv_path(model_name, type): + """Return the per-model CSV path for an experiment artifact family.""" return os.path.join(DATA_DIR, f"{type}_{model_name}.csv") def json_path(model_name, type): + """Return the per-model JSON path for an experiment artifact family.""" return os.path.join(DATA_DIR, f"{type}_{model_name}.json") def defense_summary_csv_path(): + """Return the aggregate defense-summary CSV path.""" return os.path.join(DATA_DIR, "defense_summary.csv") def check_environment(): + """Validate runtime packages and local CIFAR-10 data before a full run.""" missing = [ pkg for pkg in ("torchattacks", "autoattack", "pytorchcv") @@ -75,6 +84,7 @@ def check_environment(): 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), @@ -121,6 +131,7 @@ 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'.") if ptcv_get_model is None: @@ -133,21 +144,27 @@ def load_pretrained(arch_key): def sanity_check_accuracy(model, loader): + """Compute clean accuracy with the same evaluation path as attack metrics.""" model.eval() return accuracy_from_adv_fn(model, loader, use_autocast=True) class FakeQuantSTE(torch.autograd.Function): + """Round in forward and pass identity gradients backward.""" + @staticmethod def forward(ctx, x): + """Round values during the forward pass of the STE quantizer.""" return torch.round(x) @staticmethod def backward(ctx, grad_output): + """Pass gradients through the STE quantizer unchanged.""" return grad_output def quantize_tensor(t, bits, use_ste): + """Apply signed per-tensor fake quantization to ``t``.""" if bits is None: return t qmax = 2 ** (bits - 1) - 1 @@ -159,6 +176,7 @@ def quantize_tensor(t, bits, use_ste): 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) @@ -180,6 +198,7 @@ def chaotic_sequence_like(t, seed=CHAOTIC_QUANT_SEED, map_name=CHAOTIC_QUANT_MAP def chaotic_quantize_tensor( t, bits, use_ste, quantize=True, dither_amplitude=CHAOTIC_QUANT_DITHER ): + """Fake-quantize with deterministic chaotic subtractive dithering.""" if bits is None or not quantize: return t qmax = 2 ** (bits - 1) - 1 @@ -193,9 +212,12 @@ def chaotic_quantize_tensor( class _QuantizedLayerMixin: + """Shared fake-quantization controls for custom Conv2d/Linear layers.""" + chaotic = False def _quant_params(self): + """Return this layer's effective quantization bit widths.""" bits = getattr(self, "bits", None) use_ste = getattr(self, "use_ste", QUANT_DEFAULT_USE_STE) quant_weight = getattr(self, "quant_weight", QUANT_DEFAULT_WEIGHT) @@ -203,6 +225,7 @@ def _quant_params(self): return bits, use_ste, quant_weight, quant_act def _quantize(self, t, bits, use_ste, enabled=True): + """Fake-quantize a tensor when the corresponding quantization path is enabled.""" if self.chaotic: return chaotic_quantize_tensor( t, @@ -215,7 +238,10 @@ def _quantize(self, t, bits, use_ste, enabled=True): class QuantConv2d(_QuantizedLayerMixin, nn.Conv2d): + """Conv2d that can fake-quantize weights and output activations.""" + def forward(self, x): + """Run convolution with this layer's quantization settings applied.""" bits, use_ste, quant_weight, quant_act = self._quant_params() w = self._quantize(self.weight, bits, use_ste, quant_weight) out = self._conv_forward(x, w, self.bias) @@ -225,11 +251,16 @@ def forward(self, x): class ChaoticQuantConv2d(QuantConv2d): + """Quantized Conv2d variant that uses chaotic dither before rounding.""" + chaotic = True class QuantLinear(_QuantizedLayerMixin, nn.Linear): + """Linear layer that can fake-quantize weights and output activations.""" + def forward(self, x): + """Run a linear projection with this layer's quantization settings applied.""" bits, use_ste, quant_weight, quant_act = self._quant_params() w = self._quantize(self.weight, bits, use_ste, quant_weight) out = F.linear(x, w, self.bias) @@ -239,6 +270,8 @@ def forward(self, x): class ChaoticQuantLinear(QuantLinear): + """Quantized Linear variant that uses chaotic dither before rounding.""" + chaotic = True @@ -250,6 +283,7 @@ def _to_quant_module( chaotic=False, dither_amplitude=CHAOTIC_QUANT_DITHER, ): + """Convert a supported leaf module to the matching fake-quantized module.""" if isinstance(mod, nn.Conv2d): cls = ChaoticQuantConv2d if chaotic else QuantConv2d new = cls( @@ -299,6 +333,7 @@ def _replace_recursive( chaotic=False, dither_amplitude=CHAOTIC_QUANT_DITHER, ): + """Recursively replace supported children with fake-quantized layers.""" for name, child in list(module.named_children()): nc = _to_quant_module( child, @@ -329,6 +364,7 @@ def convert_to_quant( chaotic=False, dither_amplitude=CHAOTIC_QUANT_DITHER, ): + """Return a deep-copied model with Conv2d/Linear layers fake-quantized.""" m = copy.deepcopy(model) _replace_recursive( m, @@ -348,6 +384,7 @@ def convert_to_chaotic_quant( quant_act=True, dither_amplitude=CHAOTIC_QUANT_DITHER, ): + """Return a fake-quantized copy that uses chaotic dither in quantized layers.""" return convert_to_quant( model, bits, @@ -359,6 +396,7 @@ def convert_to_chaotic_quant( def quantizable_layer_names(model): + """List FP32 Conv2d/Linear module names that are eligible for conversion.""" quant_types = (QuantConv2d, QuantLinear, ChaoticQuantConv2d, ChaoticQuantLinear) return [ n @@ -368,6 +406,7 @@ def quantizable_layer_names(model): def set_child_module(root, module_name, new_module): + """Replace a dotted child module path inside ``root``.""" parts = module_name.split(".") parent = root for part in parts[:-1]: @@ -378,6 +417,7 @@ def set_child_module(root, module_name, new_module): def convert_layer_chunk_to_quant( model, layer_names, bits, quant_weight=True, quant_act=True, chaotic=False ): + """Quantize only the named layer subset for localization experiments.""" m = copy.deepcopy(model) targets = set(layer_names) for name, mod in list(m.named_modules()): @@ -390,6 +430,7 @@ def convert_layer_chunk_to_quant( def quant_layer_chunks(layer_names, n_chunks): + """Split quantizable layer names into non-empty ordered chunks.""" if not layer_names: return [] n_chunks = max(1, min(n_chunks, len(layer_names))) @@ -401,6 +442,7 @@ def quant_layer_chunks(layer_names, n_chunks): def count_quant_layers(model): + """Count custom fake-quantized Conv2d/Linear layers in ``model``.""" return sum( 1 for m in model.modules() @@ -413,6 +455,7 @@ def count_quant_layers(model): def verify_quantization_layers( arch_key, fp32_model, quant_model, label, fp32_layer_names=None ): + """Fail fast when a converted model did not replace most eligible layers.""" fp32_layer_names = ( quantizable_layer_names(fp32_model) if fp32_layer_names is None @@ -431,6 +474,7 @@ def verify_quantization_layers( def set_quant_components(model, quant_weight, quant_act): + """Toggle weight and activation fake quantization on converted layers.""" for mod in model.modules(): if isinstance( mod, (QuantConv2d, QuantLinear, ChaoticQuantConv2d, ChaoticQuantLinear) @@ -447,6 +491,7 @@ def prepare_qat( lr=QAT_LR, chaotic=False, ): + """Fine-tune a fake-quantized copy with STE-enabled quantization-aware training.""" m = ( convert_to_chaotic_quant(fp32_model, bits, quant_weight=True, quant_act=True) if chaotic @@ -483,17 +528,23 @@ def prepare_qat( 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, @@ -501,6 +552,7 @@ def __init__( bits=COMPRESS_IMAGE_BITS, mode=COMPRESS_IMAGE_MODE, ): + """Wrap a model with differentiable image-compression preprocessing.""" super().__init__() self.model = model self.size = size @@ -508,6 +560,7 @@ def __init__( 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} @@ -525,6 +578,7 @@ def forward(self, x): def with_image_compression( model, size=COMPRESS_IMAGE_SIZE, bits=COMPRESS_IMAGE_BITS, mode=COMPRESS_IMAGE_MODE ): + """Return an eval-mode copy of ``model`` behind input compression.""" return ( CompressedInputModel(copy.deepcopy(model), size=size, bits=bits, mode=mode) .to(device) @@ -532,135 +586,18 @@ def with_image_compression( ) -def wilson_interval(correct, total, confidence=CI_CONFIDENCE): - """Dependency-free Wilson score interval for a binomial proportion.""" - if total <= 0: - return None, None - # 1.95996 is the two-sided 95% normal quantile; keep the configured - # confidence explicit and use statistics.NormalDist for other levels. - from statistics import NormalDist - - z = NormalDist().inv_cdf(0.5 + confidence / 2) - p = correct / total - denom = 1 + z * z / total - centre = (p + z * z / (2 * total)) / denom - radius = z * math.sqrt(p * (1 - p) / total + z * z / (4 * total * total)) / denom - return max(0.0, centre - radius), min(1.0, centre + radius) - - -def clopper_pearson_interval(correct, total, confidence=CI_CONFIDENCE): - """Exact binomial interval when scipy is present; Wilson is a safe fallback.""" - if total <= 0: - return None, None - try: - from scipy.stats import beta - - alpha = 1 - confidence - low = ( - 0.0 - if correct == 0 - else float(beta.ppf(alpha / 2, correct, total - correct + 1)) - ) - high = ( - 1.0 - if correct == total - else float(beta.ppf(1 - alpha / 2, correct + 1, total - correct)) - ) - return low, high - except ImportError: - return wilson_interval(correct, total, confidence) - - -def add_binomial_statistics(results, metric, correct_vector): - vector = np.asarray(correct_vector, dtype=bool) - n = int(vector.size) - k = int(vector.sum()) - wlo, whi = wilson_interval(k, n) - clo, chi = clopper_pearson_interval(k, n) - results.update( - { - f"{metric}_n": n, - f"{metric}_correct": k, - f"{metric}_wilson_low": wlo, - f"{metric}_wilson_high": whi, - f"{metric}_wilson_pm": (whi - wlo) / 2 if wlo is not None else None, - f"{metric}_cp_low": clo, - f"{metric}_cp_high": chi, - f"{metric}_cp_pm": (chi - clo) / 2 if clo is not None else None, - } - ) - - -def mcnemar_exact(vector_a, vector_b): - """Two-sided exact McNemar test over paired correctness outcomes.""" - a = np.asarray(vector_a, dtype=bool) - b = np.asarray(vector_b, dtype=bool) - if a.shape != 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)) - discordant = a_only + b_only - if discordant == 0: - p_value = 1.0 - else: - tail = sum(math.comb(discordant, i) for i in range(min(a_only, b_only) + 1)) / ( - 2**discordant - ) - p_value = min(1.0, 2.0 * tail) - return { - "a_only": a_only, - "b_only": b_only, - "discordant": discordant, - "p_value": p_value, - } - - -def save_correctness_vectors(model_name, vectors): - os.makedirs(PER_EXAMPLE_DIR, exist_ok=True) - safe_name = "".join(c if c.isalnum() or c in "._-" else "_" for c in model_name) - path = os.path.join(PER_EXAMPLE_DIR, f"{safe_name}.npz") - np.savez_compressed( - path, **{key: np.asarray(value, dtype=bool) for key, value in vectors.items()} - ) - return path - - def add_paired_fp32_mcnemar_tests(df_results): """Post-hoc paired tests between each variant and its architecture's FP32 model.""" - if "correctness_vectors_path" not in df_results: - return df_results - for idx, row in df_results.iterrows(): - model_name = str(row["model"]) - architecture = model_name.split("_", 1)[0] - baseline_rows = df_results[ - df_results["model"].astype(str) == f"{architecture}_FP32" - ] - if baseline_rows.empty or not isinstance( - row.get("correctness_vectors_path"), str - ): - continue - baseline_path = baseline_rows.iloc[0].get("correctness_vectors_path") - if not isinstance(baseline_path, str) or not os.path.exists(baseline_path): - continue - if not os.path.exists(row["correctness_vectors_path"]): - continue - with np.load(baseline_path) as baseline, np.load( - row["correctness_vectors_path"] - ) as variant: - for metric in sorted(set(baseline.files) & set(variant.files)): - test = mcnemar_exact(baseline[metric], variant[metric]) - prefix = f"McNemar_vs_FP32_{metric}" - for key, value in test.items(): - df_results.loc[idx, f"{prefix}_{key}"] = value - return df_results + return qstats.add_paired_mcnemar_tests( + df_results, baseline_name=qstats.fp32_baseline_name + ) def gradient_diagnostics( model, loader, fp32_ref=None, max_batches=GRAD_DIAG_MAX_BATCHES ): - set_ste_mode(model, False) + """Ground-truth masking check: compare hard-round vs. STE input gradients. + """ frac_zero_hard, norm_hard = [], [] frac_zero_ste, norm_ste = [], [] cos_sims = [] @@ -668,19 +605,18 @@ def gradient_diagnostics( if bi >= max_batches: break x, y = x.to(device), y.to(device) - set_ste_mode(model, False) - 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() + with ste_mode(model, False): + 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_hard.append( (g_hard.abs() < GRAD_ZERO_THRESHOLD).float().mean().item() ) norm_hard.append(g_hard.norm().item()) - set_ste_mode(model, True) - x_in2 = x.clone().requires_grad_(True) - loss2 = F.cross_entropy(model(x_in2), y) - g_ste = torch.autograd.grad(loss2, x_in2)[0].flatten() - set_ste_mode(model, False) + with ste_mode(model, True): + x_in2 = x.clone().requires_grad_(True) + loss2 = F.cross_entropy(model(x_in2), y) + g_ste = torch.autograd.grad(loss2, x_in2)[0].flatten() frac_zero_ste.append((g_ste.abs() < GRAD_ZERO_THRESHOLD).float().mean().item()) norm_ste.append(g_ste.norm().item()) if fp32_ref is not None: @@ -703,6 +639,7 @@ def gradient_diagnostics( def layerwise_grad_profile(model, loader, use_ste, max_batches=LAYERWISE_MAX_BATCHES): + """Collect per-layer gradient statistics for a model and loader.""" quant_layers = [ (n, m) for n, m in model.named_modules() @@ -712,7 +649,9 @@ def layerwise_grad_profile(model, loader, use_ste, max_batches=LAYERWISE_MAX_BAT handles = [] 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] if gi is not None: norms[name].append(gi.flatten(1).norm(dim=1).mean().item()) @@ -722,20 +661,19 @@ def hook(module, grad_input, grad_output): try: for n, m in quant_layers: handles.append(m.register_full_backward_hook(make_hook(n))) - set_ste_mode(model, use_ste) model.eval() - for bi, (x, y) in enumerate(loader): - if bi >= max_batches: - break - x, y = x.to(device), y.to(device) - x = x.clone().requires_grad_(True) - loss = F.cross_entropy(model(x), y) - model.zero_grad(set_to_none=True) - loss.backward() + with ste_mode(model, use_ste): + for bi, (x, y) in enumerate(loader): + if bi >= max_batches: + break + x, y = x.to(device), y.to(device) + x = x.clone().requires_grad_(True) + loss = F.cross_entropy(model(x), y) + model.zero_grad(set_to_none=True) + loss.backward() finally: for h in handles: h.remove() - set_ste_mode(model, False) ordered_names = [n for n, _ in quant_layers] return { n: (float(np.mean(norms[n])) if len(norms[n]) else None) for n in ordered_names @@ -743,6 +681,22 @@ def hook(module, grad_input, grad_output): 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. + 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 + of gradient masking, not evidence that quantization itself improves + robustness. + + Weight quantization alone cannot produce this gap: rounding a weight + 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 + tensor that ``x`` actually flows through, so only those can mask. + """ configs = [ ("weight_only", True, False), ("act_only", False, True), @@ -752,22 +706,45 @@ def run_quant_component_ablation(model, loader, name, eps=DEFAULT_EPS): for label, qw, qa in configs: set_quant_components(model, qw, qa) clean_acc = sanity_check_accuracy(model, loader) - torch.manual_seed(0) - pgd = make_torchattack( - torchattacks.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) + 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() + + 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, eps=eps, - alpha=PGD_ALPHA, - steps=PGD_STEPS, - random_start=PGD_RANDOM_START, + n_restarts=1, + seeds=SEEDS[:1], ) - pgd_acc = accuracy_under_attack(model, loader, pgd) - 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() rows.append( { "model": name, @@ -775,7 +752,22 @@ def run_quant_component_ablation(model, loader, name, eps=DEFAULT_EPS): "quant_weight": qw, "quant_act": qa, "clean_acc": clean_acc, - "PGD_acc": pgd_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 + else None + ), "frac_zero_grad_hard": frac_zero, } ) @@ -792,6 +784,7 @@ def run_chunk_quantization_attacks( n_chunks=CHUNK_QUANT_NUM_CHUNKS, eps=DEFAULT_EPS, ): + """Evaluate attacks as contiguous chunks of layers are quantized.""" layer_names = quantizable_layer_names(fp32_model) chunks = quant_layer_chunks(layer_names, n_chunks) rows = [] @@ -852,11 +845,13 @@ def safe_call( def safe_set(target, key, fn, warning, *, context=None, default=None): + """Run a metric function and store its value with warning-based fallback.""" target[key] = safe_call(fn, warning, context=context, default=default) return target[key] def safe_update(target, fn, warning, *, context=None, defaults=None): + """Merge a metric dictionary into a target dictionary with warning fallback.""" update = safe_call(fn, warning, context=context, default=None) if update is not None: target.update(update) @@ -872,6 +867,7 @@ def safe_update(target, fn, warning, *, context=None, defaults=None): def safe_set_vector(results, vectors, metric, fn, warning, *, context=None): + """Run a vector-producing metric and store values plus correctness vectors.""" pair = safe_call(fn, warning, context=context, default=None) if pair is None: results[metric] = None @@ -881,6 +877,7 @@ def safe_set_vector(results, vectors, metric, fn, warning, *, context=None): def safe_update_vectors(results, vectors, fn, warning, *, context=None, defaults=None): + """Merge metric and vector outputs into their destination dictionaries.""" update = safe_call(fn, warning, context=context, default=None) if update is None: if defaults: @@ -892,11 +889,13 @@ def safe_update_vectors(results, vectors, fn, warning, *, context=None, defaults def save_json(path, data, *, indent=None): + """Serialize a Python object as pretty-printed JSON.""" with open(path, "w") as handle: json.dump(data, handle, indent=indent) def run_suite(model, loader, name, fp32_ref=None, eps=DEFAULT_EPS): + """Run the full attack and diagnostic suite for one model.""" model.eval() results = {"model": name} vectors = {} @@ -1026,6 +1025,13 @@ def run_suite(model, loader, name, fp32_ref=None, eps=DEFAULT_EPS): defaults=adaptive_defaults, ) if count_quant_layers(model) > 0: + # n_restarts=1 here (with the default 5-element SEEDS) gives 5 total + # restarts, matching run_fgsm_pgd's PGD exactly (one restart per + # seed, no inner restart loop). BPDA_RESTARTS_SUITE (5) would nest + # inside run_bpda's outer seed loop for ~25 effective restarts, + # making any PGD-vs-BPDA_PGD gap partly a bigger-attack-budget + # artifact rather than purely a gradient-regime effect — which is + # what Gradient_Masking_Gap downstream is supposed to isolate. safe_update_vectors( results, vectors, @@ -1033,7 +1039,7 @@ def run_suite(model, loader, name, fp32_ref=None, eps=DEFAULT_EPS): model, loader, eps=eps, - n_restarts=BPDA_RESTARTS_SUITE, + n_restarts=1, return_vector=True, ), "BPDA failed", @@ -1099,6 +1105,7 @@ def run_suite(model, loader, name, fp32_ref=None, eps=DEFAULT_EPS): ) 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()] @@ -1108,6 +1115,7 @@ def save_pgd_ablation(): if RUN_PGD_TRAJECTORY: def save_trajectory(): + """Persist PGD trajectory diagnostics for the current model.""" traj = pgd_trajectory_diagnostics( model, loader, eps=eps, max_batches=TRAJECTORY_MAX_BATCHES ) @@ -1119,6 +1127,7 @@ def save_trajectory(): if RUN_LAYERWISE_PROFILE: def save_layerwise_profile(): + """Persist layerwise gradient diagnostics for the current model.""" prof_hard = layerwise_grad_profile(model, loader, use_ste=False) prof_ste = layerwise_grad_profile(model, loader, use_ste=True) rows = [ @@ -1139,6 +1148,7 @@ def save_layerwise_profile(): if RUN_COMPONENT_ABLATION: def save_component_ablation(): + """Persist quantization component-ablation diagnostics for the current model.""" rows = run_quant_component_ablation(model, loader, name, eps=eps) pd.DataFrame(rows).to_csv( csv_path(name, "component_ablation"), index=False @@ -1152,6 +1162,7 @@ def save_component_ablation(): if RUN_CONFIDENCE_MARGIN: def save_confidence_margins(): + """Persist confidence-margin diagnostics for the current model.""" margins = confidence_margin_diagnostic( model, loader, eps=eps, max_batches=MARGIN_MAX_BATCHES ) @@ -1163,18 +1174,23 @@ def save_confidence_margins(): context=name, ) for metric, vector in vectors.items(): - add_binomial_statistics(results, metric, vector) + qstats.add_binomial_statistics( + results, metric, vector, confidence=CI_CONFIDENCE + ) if "PGD" in vectors and "BPDA_PGD" in vectors: - test = mcnemar_exact(vectors["PGD"], vectors["BPDA_PGD"]) + test = qstats.mcnemar_exact(vectors["PGD"], vectors["BPDA_PGD"]) results.update( {f"McNemar_PGD_vs_BPDA_{key}": value for key, value in test.items()} ) if vectors: - results["correctness_vectors_path"] = save_correctness_vectors(name, vectors) + results["correctness_vectors_path"] = qstats.save_correctness_vectors( + name, vectors, PER_EXAMPLE_DIR + ) return results def run_epsilon_sweep_for_model_wrapped(model, loader, name, epsilons): + """Run and annotate an epsilon sweep for one named model.""" return run_epsilon_sweep_for_model( model, loader, @@ -1186,6 +1202,7 @@ def run_epsilon_sweep_for_model_wrapped(model, loader, name, epsilons): def run_defense_suite(model_registry, finetune_loader, eval_loader): + """Build defended models and evaluate them under adaptive attacks.""" summary_rows = [] arch_keys = sorted( {name.split("_FP32")[0] for name in model_registry if name.endswith("_FP32")} @@ -1198,6 +1215,7 @@ def run_defense_suite(model_registry, finetune_loader, eval_loader): fp32_model = fp32_entry[0] def add_fp32_at(): + """Add the adversarially trained FP32 defense to the registry when available.""" fp32_at = dfn.prepare_adversarial_training( fp32_model, finetune_loader, bits=None ) @@ -1212,6 +1230,7 @@ def add_fp32_at(): ) def add_int8_at(): + """Add the adversarially trained INT8 defense to the registry when available.""" int8_at = dfn.prepare_adversarial_training( fp32_model, finetune_loader, bits=QAT_BITS ) @@ -1238,6 +1257,7 @@ def add_int8_at(): entry_name = f"{arch_key}_{tag}" def add_sanitized(): + """Add the input-sanitization defense to the registry.""" sanitized = dfn.SanitizedModel(base_model).to(device).eval() model_registry[f"{entry_name}_Sanitized"] = (sanitized, fp32_model) @@ -1246,6 +1266,7 @@ def add_sanitized(): ) def add_smoothed(): + """Add the randomized-smoothing defense to the registry.""" smoothed = dfn.SmoothedModel(base_model).to(device).eval() model_registry[f"{entry_name}_Smoothed"] = (smoothed, fp32_model) cert_stats = dfn.run_certified_accuracy(smoothed, eval_loader) @@ -1266,6 +1287,7 @@ def add_smoothed(): ) def add_guardrail(): + """Add the confidence guardrail defense to the registry.""" guardrail = dfn.GuardrailModel(base_model).to(device).eval() model_registry[f"{entry_name}_Guardrail"] = (guardrail, fp32_model) pgd_for_flagging = make_torchattack( @@ -1293,6 +1315,7 @@ def add_guardrail(): if detector is not None: def add_detect_guard(): + """Add the detector-based guardrail defense to the registry.""" detect_guard = ( dfn.DetectGuardModel(base_model, detector).to(device).eval() ) @@ -1329,16 +1352,19 @@ def add_detect_guard(): def _palette_for(values): + """Choose a plotting palette sized to the provided values.""" return {v: ATTACK_PALETTE[v] for v in values if v in ATTACK_PALETTE} def parallelize(model): + """Wrap a model in DataParallel when multiple CUDA devices are available.""" if torch.cuda.device_count() > 1 and not isinstance(model, nn.DataParallel): return nn.DataParallel(model) return model def run_chaotic_dither_sweep(fp32_model, loader, arch_key, bits=QAT_BITS): + """Evaluate chaotic quantization over dither-amplitude settings.""" rows = [] for amplitude in CHAOTIC_DITHER_AMPLITUDES: model = ( @@ -1415,6 +1441,7 @@ def add_general_scores(df_results): def main(): + """Run the command-line entry point for this module.""" check_environment() finetune_loader, eval_loader = get_dataloaders() model_registry = {} @@ -1622,6 +1649,7 @@ def main(): sweep_done = set() def run_pending_epsilon_sweep(name, model): + """Run an epsilon sweep for a model that still needs one.""" nonlocal df_sweep, sweep_done if not RUN_EPSILON_SWEEP: return @@ -1635,6 +1663,7 @@ def run_pending_epsilon_sweep(name, model): print(f"\nSweeping {name} ...") 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 @@ -1737,6 +1766,7 @@ def save_epsilon_sweep(): print("\nEpsilon sweep completed. Results saved to", SWEEP_CSV) 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) diff --git a/src/ResourceMonitor.py b/src/ResourceMonitor.py index 68e2d4e..d3554ae 100644 --- a/src/ResourceMonitor.py +++ b/src/ResourceMonitor.py @@ -11,6 +11,7 @@ class ResourceMonitor: """Record resource use around an existing run without executing extra model work.""" def __init__(self, model, name, sample_interval_seconds=0.1): + """Initialize process-resource tracking for a model run.""" self.model = model self.name = name self.sample_interval_seconds = sample_interval_seconds @@ -41,11 +42,13 @@ def _process_tree_usage(process): return rss_bytes, user_seconds, system_seconds def _sample_memory(self): + """Sample current process-tree memory usage into the monitor state.""" while not self.stop_event.wait(self.sample_interval_seconds): rss_bytes, _, _ = self._process_tree_usage(self.process) self.peak_rss_bytes = max(self.peak_rss_bytes, rss_bytes) def __enter__(self): + """Start resource monitoring and return this context manager.""" self.start_rss_bytes, self.start_user_seconds, self.start_system_seconds = ( self._process_tree_usage(self.process) ) @@ -64,6 +67,7 @@ def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback_value): + """Stop monitoring and finalize peak memory and elapsed-time metrics.""" if torch.cuda.is_available(): torch.cuda.synchronize() elapsed_seconds = time.perf_counter() - self.start_time diff --git a/src/attack.py b/src/attack.py index df23889..0d03fd6 100644 --- a/src/attack.py +++ b/src/attack.py @@ -11,6 +11,8 @@ """ import warnings +from contextlib import contextmanager + import numpy as np import torch import torch.nn as nn @@ -24,10 +26,13 @@ def unwrap_model(model): + """Return the underlying module when ``model`` is wrapped in DataParallel.""" 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"): @@ -36,15 +41,61 @@ def set_ste_mode(model, flag): return toggled +def get_ste_mode(model): + """Return the current ``use_ste`` state of a model's quantized modules. + + Returns ``None`` if the model has no quantized modules, ``True``/``False`` + if all quantized modules agree, or the string ``"mixed"`` if they don't + (which should never happen in normal use, and indicates something set + ``use_ste`` on a subset of layers directly rather than through + ``set_ste_mode``/``ste_mode``). + """ + 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 CIFAR tensors with the experiment constants.""" return (x - CIFAR_MEAN.to(x.device)) / CIFAR_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) class PixelSpaceModel(nn.Module): + """Adapter exposing a normalized-input model as a pixel-space model.""" + def __init__(self, model): super().__init__() self.model = model @@ -56,6 +107,7 @@ def forward(self, x): def make_torchattack(attack_cls, model, *args, **kwargs): + """Create a torchattacks instance configured for normalized CIFAR inputs.""" attack = attack_cls(model, *args, **kwargs) attack.set_normalization_used(mean=CIFAR_MEAN_VALUES, std=CIFAR_STD_VALUES) return attack @@ -70,6 +122,7 @@ def accuracy_from_adv_fn( use_autocast=False, return_vector=False, ): + """Measure accuracy after an optional adversarial-example function.""" target = target_model if target_model is not None else model correct, total, n_seen = 0, 0, 0 @@ -108,6 +161,7 @@ def accuracy_from_adv_fn( 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) @@ -122,11 +176,13 @@ def adv_fn(x, y): def evaluate_normalized_attack(model, loader, attack_fn): + """Evaluate an attack function that consumes and returns normalized tensors.""" model.eval() return accuracy_from_adv_fn(model, loader, attack_fn) def seed_averaged_metrics(name, seeds, fn): + """Run a metric for multiple seeds and return mean/std summaries.""" accs = [] for seed in seeds: torch.manual_seed(seed) @@ -140,9 +196,15 @@ def seed_averaged_metrics(name, seeds, fn): def run_fgsm_pgd(model, loader, eps=DEFAULT_EPS, seeds=SEEDS, return_vectors=False): + """Run FGSM and multi-restart PGD, aggregating PGD by worst-case correctness. + + Uses STE (straight-through) gradients explicitly: this is the + "attacker who knows to bypass fake-quantization rounding" number, not + the naive hard-round number. See ``pgd_steps_ablation`` / + ``run_epsilon_sweep_for_model`` for the explicit hard-round counterpart. + """ model.eval() - set_ste_mode(model, True) - try: + with ste_mode(model, True): fgsm = make_torchattack(torchattacks.FGSM, model, eps=eps) out = {} @@ -178,14 +240,12 @@ def adv_fn(x, y): if return_vectors: out["_vectors"] = {"FGSM": fgsm_vector, "PGD": pgd_vector} return out - finally: - set_ste_mode(model, False) def run_autoattack(model, loader, eps=DEFAULT_EPS, return_vector=False): + """Run AutoAttack in pixel space while reporting normalized-model accuracy.""" model.eval() - set_ste_mode(model, True) - try: + with ste_mode(model, True): pixel_model = PixelSpaceModel(model).to(device).eval() adversary = AutoAttack( pixel_model, @@ -204,16 +264,13 @@ def adv_fn(x, y): ) return accuracy_from_adv_fn(model, loader, adv_fn, return_vector=return_vector) - finally: - set_ste_mode(model, False) def run_extra_whitebox_attacks( model, loader, eps=DEFAULT_EPS, jsma_max_images=JSMA_MAX_IMAGES ): model.eval() - set_ste_mode(model, True) - try: + with ste_mode(model, True): out = {} cw = make_torchattack( torchattacks.CW, model, c=CW_C, kappa=CW_KAPPA, steps=CW_STEPS, lr=CW_LR @@ -233,15 +290,13 @@ def run_extra_whitebox_attacks( model, loader, jsma, max_images=jsma_max_images ) return out - finally: - set_ste_mode(model, False) def transfer_attack( source_model, target_model, loader, eps=DEFAULT_EPS, return_vector=False ): - set_ste_mode(source_model, True) - try: + """Generate PGD examples on ``source_model`` and evaluate ``target_model``.""" + with ste_mode(source_model, True): pgd = make_torchattack( torchattacks.PGD, source_model, @@ -257,8 +312,6 @@ def transfer_attack( target_model=target_model, return_vector=return_vector, ) - finally: - set_ste_mode(source_model, False) def transfer_attack_mim( @@ -346,6 +399,7 @@ def transfer_uap_attack( 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) @@ -364,23 +418,20 @@ def projected_pgd_attack(x, y, eps, alpha, steps, grad_fn): def bpda_pgd_attack( model, x, y, eps=DEFAULT_EPS, alpha=PGD_ALPHA, steps=PGD_STEPS, backward_model=None ): - set_ste_mode(model, True) + """Run PGD with STE/BPDA gradients through non-differentiable defenses.""" backward_model = backward_model if backward_model is not None else model - if backward_model is not model: - set_ste_mode(backward_model, True) - try: - def grad_fn(x_adv, labels): - x_adv.requires_grad_(True) - loss = F.cross_entropy(backward_model(x_adv), labels) - grad = torch.autograd.grad(loss, x_adv, allow_unused=True)[0] - return grad + def grad_fn(x_adv, labels): + x_adv.requires_grad_(True) + loss = F.cross_entropy(backward_model(x_adv), labels) + grad = torch.autograd.grad(loss, x_adv, allow_unused=True)[0] + return grad - return projected_pgd_attack(x, y, eps, alpha, steps, grad_fn) - finally: - set_ste_mode(model, False) + with ste_mode(model, True): if backward_model is not model: - set_ste_mode(backward_model, False) + with ste_mode(backward_model, True): + return projected_pgd_attack(x, y, eps, alpha, steps, grad_fn) + return projected_pgd_attack(x, y, eps, alpha, steps, grad_fn) def _run_bpda_once(model, loader, eps, n_restarts, return_vector=False): @@ -411,6 +462,18 @@ def run_bpda( seeds=SEEDS, return_vector=False, ): + """Evaluate BPDA-PGD with seed-level worst-case correctness aggregation. + + IMPORTANT for masking-gap comparisons: the true attack budget here is + ``len(seeds) * n_restarts`` random restarts (an outer seed loop wrapping + an inner ``n_restarts`` loop in ``_run_bpda_once``), not just + ``len(seeds)``. ``run_fgsm_pgd``'s ``PGD`` metric, by contrast, uses + exactly ``len(seeds)`` restarts (one per seed, no inner loop). If you + call this with ``n_restarts > 1`` and then compare the result to + ``PGD`` as a "masking gap", part of any gap you see is simply a bigger + attack budget for BPDA, not evidence of masking. Pass ``n_restarts=1`` + to keep the comparison budget-matched to ``PGD``. + """ vectors, accuracies = [], [] for seed in seeds: torch.manual_seed(seed) @@ -433,18 +496,14 @@ def run_bpda( def adaptive_pgd_attack( model, x, y, loss_fn, eps=DEFAULT_EPS, alpha=PGD_ALPHA, steps=PGD_STEPS ): - set_ste_mode(model, True) - try: - - def grad_fn(x_adv, labels): - x_adv.requires_grad_(True) - loss = loss_fn(x_adv, labels) - grad = torch.autograd.grad(loss, x_adv, allow_unused=True)[0] - return grad + def grad_fn(x_adv, labels): + x_adv.requires_grad_(True) + loss = loss_fn(x_adv, labels) + grad = torch.autograd.grad(loss, x_adv, allow_unused=True)[0] + return grad + with ste_mode(model, True): return projected_pgd_attack(x, y, eps, alpha, steps, grad_fn) - finally: - set_ste_mode(model, False) def run_sanitized_bpda(model, loader, eps=DEFAULT_EPS): @@ -969,23 +1028,39 @@ def run_random_noise_seeded( return out -def pgd_steps_ablation(model, loader, eps=DEFAULT_EPS, step_list=PGD_ABLATION_STEPS): +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. + + ``use_ste=False`` (the default) is the "naive attacker against the real + fake-quantized model" regime: if accuracy stays flat as ``steps`` + increases, that's gradient masking, not robustness — a real attack gets + *more* effective with more steps unless the gradient carries no signal. + Pass ``use_ste=True`` to see the same sweep with the masking bypassed, + for a direct side-by-side comparison. + + Previously this inherited whatever ``use_ste`` state a prior, unrelated + attack elsewhere in the run happened to leave the model in, rather than + choosing one explicitly. + """ model.eval() out = {} - 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 + 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 return out @@ -996,7 +1071,18 @@ def pgd_trajectory_diagnostics( alpha=PGD_ALPHA, steps=PGD_STEPS, max_batches=TRAJECTORY_MAX_BATCHES, + use_ste=False, ): + """Per-step gradient norm and movement along a PGD trajectory. + + Computes gradients manually via ``torch.autograd.grad`` rather than + through a torchattacks object, so it never went through + ``set_ste_mode`` at all previously — it just used whatever ``use_ste`` + happened to be set to by the last thing that ran before it. Now + explicit; defaults to ``False`` (hard round) since a near-zero + ``grad_norm_per_step`` here is exactly the masking signature this + diagnostic exists to catch. + """ model.eval() clip_min, clip_max = CLIP_MIN.to(device), CLIP_MAX.to(device) eps_normalized = eps / CIFAR_STD.to(device) @@ -1004,25 +1090,28 @@ def pgd_trajectory_diagnostics( step_grad_norms = [0.0] * steps step_movement = [0.0] * steps n_batches = 0 - for bi, (x, y) in enumerate(loader): - if bi >= max_batches: - break - x, y = x.to(device), y.to(device) - 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() - for s in range(steps): - x_adv.requires_grad_(True) - loss = F.cross_entropy(model(x_adv), y) - grad = torch.autograd.grad(loss, x_adv)[0] - step_grad_norms[s] += grad.flatten(1).norm(dim=1).mean().item() - x_adv = x_adv.detach() + alpha_normalized * grad.sign() - x_adv = torch.min(torch.max(x_adv, x - eps_normalized), x + eps_normalized) - x_adv = torch.max(torch.min(x_adv, clip_max), clip_min).detach() - step_movement[s] += ( - (x_adv - x_start).flatten(1).abs().max(dim=1).values.mean().item() - ) - n_batches += 1 + with ste_mode(model, use_ste): + for bi, (x, y) in enumerate(loader): + if bi >= max_batches: + break + x, y = x.to(device), y.to(device) + 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() + for s in range(steps): + x_adv.requires_grad_(True) + loss = F.cross_entropy(model(x_adv), y) + grad = torch.autograd.grad(loss, x_adv)[0] + step_grad_norms[s] += grad.flatten(1).norm(dim=1).mean().item() + x_adv = x_adv.detach() + alpha_normalized * grad.sign() + x_adv = torch.min( + torch.max(x_adv, x - eps_normalized), x + eps_normalized + ) + x_adv = torch.max(torch.min(x_adv, clip_max), clip_min).detach() + step_movement[s] += ( + (x_adv - x_start).flatten(1).abs().max(dim=1).values.mean().item() + ) + n_batches += 1 return { "grad_norm_per_step": [g / n_batches for g in step_grad_norms], "movement_from_random_start_per_step": [m / n_batches for m in step_movement], @@ -1053,29 +1142,42 @@ def staircase_diagnostic( def confidence_margin_diagnostic( - model, loader, eps=DEFAULT_EPS, steps=MARGIN_STEPS, max_batches=MARGIN_MAX_BATCHES + model, + loader, + eps=DEFAULT_EPS, + steps=MARGIN_STEPS, + max_batches=MARGIN_MAX_BATCHES, + use_ste=False, ): + """Compare top-2 softmax margins clean vs. under PGD, at an explicit regime. + + Previously never called ``set_ste_mode`` itself, so it ran under + whatever state was leftover from earlier in the run. Defaults to + ``False`` (hard round) to match that historical behavior, but now makes + the choice explicit rather than accidental. + """ model.eval() - pgd = make_torchattack( - torchattacks.PGD, - model, - eps=eps, - alpha=PGD_ALPHA, - steps=steps, - random_start=PGD_RANDOM_START, - ) - clean_margins, adv_margins = [], [] - for bi, (x, y) in enumerate(loader): - if bi >= max_batches: - break - x, y = x.to(device), y.to(device) - 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()) - x_adv = pgd(x, y) - with torch.no_grad(): - top2_adv = F.softmax(model(x_adv), dim=1).topk(2, dim=1).values - adv_margins.extend((top2_adv[:, 0] - top2_adv[:, 1]).cpu().tolist()) + with ste_mode(model, use_ste): + pgd = make_torchattack( + torchattacks.PGD, + model, + eps=eps, + alpha=PGD_ALPHA, + steps=steps, + random_start=PGD_RANDOM_START, + ) + clean_margins, adv_margins = [], [] + for bi, (x, y) in enumerate(loader): + if bi >= max_batches: + break + x, y = x.to(device), y.to(device) + 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()) + x_adv = pgd(x, y) + with torch.no_grad(): + top2_adv = F.softmax(model(x_adv), dim=1).topk(2, dim=1).values + adv_margins.extend((top2_adv[:, 0] - top2_adv[:, 1]).cpu().tolist()) return {"clean_margins": clean_margins, "adv_margins": adv_margins} @@ -1086,6 +1188,21 @@ def run_epsilon_sweep_for_model( `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. + + ``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. """ if safe_set is None: @@ -1103,19 +1220,34 @@ def safe_set(target, key, fn, warning, *, context=None, default=None): for eps in epsilons: row = {"model": name, "epsilon": eps} - def run_pgd_sweep(): - 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 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) context = f"{name} eps={eps:.4f}" - safe_set(row, "PGD_acc", run_pgd_sweep, "PGD sweep failed", context=context) + safe_set( + row, + "PGD_acc", + lambda: run_pgd_sweep(False), + "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", @@ -1127,9 +1259,7 @@ def run_pgd_sweep(): safe_set( row, "BPDA_acc", - lambda: _run_bpda_once( - model, loader, eps=eps, n_restarts=BPDA_RESTARTS_SWEEP - ), + lambda: _run_bpda_once(model, loader, eps=eps, n_restarts=1), "BPDA sweep failed", context=context, ) diff --git a/src/combine.py b/src/combine.py index f586b24..ec541af 100644 --- a/src/combine.py +++ b/src/combine.py @@ -122,6 +122,7 @@ 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) :] @@ -131,6 +132,7 @@ def _model_from_prefixed_path(path: Path, prefix: str, suffix: str) -> str: 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: @@ -161,6 +163,7 @@ def _read_csvs(paths: Iterable[Path], model_prefix: str | None = None) -> pd.Dat 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) @@ -169,6 +172,7 @@ def _write(df: pd.DataFrame, path: Path) -> pd.DataFrame: def _result_files(data_dir: Path) -> list[Path]: + """Collect result-file paths by experiment family.""" blocked = { RESULTS_CSV.name, SWEEP_CSV.name, @@ -190,12 +194,14 @@ def _result_files(data_dir: Path) -> list[Path]: 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): @@ -206,6 +212,7 @@ def combine_sweeps(data_dir: Path, output: Path = SWEEP_CSV) -> pd.DataFrame: 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): @@ -216,6 +223,7 @@ def combine_ablation( 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) @@ -224,6 +232,7 @@ def combine_layerwise( 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 ] @@ -234,6 +243,7 @@ def combine_component_ablation( 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") @@ -263,6 +273,7 @@ def combine_trajectories( 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") @@ -284,6 +295,7 @@ def combine_margins(data_dir: Path, output: Path = MARGIN_COMBINED_CSV) -> pd.Da 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: @@ -292,6 +304,7 @@ def _model_names(*dfs: pd.DataFrame) -> list[str]: 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 @@ -342,6 +355,7 @@ 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.""" if df_sweep is None or df_sweep.empty: return value_cols = [ @@ -397,6 +411,7 @@ 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.""" if ( df_ablation is None or df_ablation.empty @@ -420,6 +435,7 @@ def plot_pgd_steps_ablation( 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", @@ -464,6 +480,7 @@ def plot_pgd_trajectory( 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 @@ -509,6 +526,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, and BPDA component-ablation accuracies.""" required = {"model", "config", "clean_acc", "PGD_acc"} if ( df_component is None @@ -517,9 +535,12 @@ def plot_component_ablation( ): 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=["clean_acc", "PGD_acc"], + value_vars=value_vars, var_name="Metric", value_name="Accuracy", ).dropna(subset=["Accuracy"]) @@ -548,6 +569,7 @@ def plot_component_ablation( 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 @@ -594,6 +616,7 @@ def plot_gradient_masking_summary( 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 @@ -651,6 +674,7 @@ def plot_confidence_margin_diagnostic( 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 @@ -703,6 +727,7 @@ def plot_results_heatmap( 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) @@ -730,6 +755,7 @@ def combine_all(data_dir: Path) -> dict[str, pd.DataFrame]: 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) @@ -742,6 +768,7 @@ def plot_all(dfs: dict[str, pd.DataFrame]) -> None: def main() -> None: + """Run the command-line entry point for this module.""" parser = argparse.ArgumentParser( description="Merge incomplete QuantAdv outputs and regenerate reports." ) diff --git a/src/config.py b/src/config.py index 8b4aed9..d03222b 100644 --- a/src/config.py +++ b/src/config.py @@ -1,5 +1,9 @@ -""" -Central configuration and constants for QuantAdvOld.py. +"""Central configuration for active QuantAdv experiments. + +Constants in this module define dataset paths, model names, perturbation +budgets, quantization defaults, attack settings, optional experiment switches, +and report output locations. Runtime modules import these values directly, so +changes here affect both the main runner and standalone analysis scripts. """ import logging diff --git a/src/data.py b/src/data.py index b9fbffb..b460551 100644 --- a/src/data.py +++ b/src/data.py @@ -23,12 +23,13 @@ matplotlib.use("Agg") import matplotlib.pyplot as plt import seaborn as sns -from scipy.stats import binomtest from config import * +import stats as qstats 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) :] @@ -38,10 +39,12 @@ def _model_from_prefixed_path(path: Path, prefix: str, suffix: str) -> str: def _name(path) -> str: + """Return a display-safe stem for a path-like value.""" return Path(path).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: @@ -72,6 +75,7 @@ def _read_csvs(paths: Iterable[Path], model_prefix: str | None = None) -> pd.Dat 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) @@ -80,6 +84,7 @@ def _write(df: pd.DataFrame, path: Path) -> pd.DataFrame: def _result_files(data_dir: Path) -> list[Path]: + """Collect result-file paths by experiment family.""" blocked = { _name(RESULTS_CSV), _name(SWEEP_CSV), @@ -101,12 +106,14 @@ def _result_files(data_dir: Path) -> list[Path]: 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): @@ -117,6 +124,7 @@ def combine_sweeps(data_dir: Path, output: Path = SWEEP_CSV) -> pd.DataFrame: 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): @@ -127,6 +135,7 @@ def combine_ablation( 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) @@ -135,6 +144,7 @@ def combine_layerwise( 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 ] @@ -145,6 +155,7 @@ def combine_component_ablation( 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") @@ -174,6 +185,7 @@ def combine_trajectories( 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") @@ -195,6 +207,7 @@ def combine_margins(data_dir: Path, output: Path = MARGIN_COMBINED_CSV) -> pd.Da 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: @@ -203,6 +216,7 @@ def _model_names(*dfs: pd.DataFrame) -> list[str]: 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 @@ -253,6 +267,7 @@ 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.""" if df_sweep is None or df_sweep.empty: return value_cols = [ @@ -308,6 +323,7 @@ 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.""" if ( df_ablation is None or df_ablation.empty @@ -331,6 +347,7 @@ def plot_pgd_steps_ablation( 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", @@ -376,6 +393,7 @@ def plot_pgd_trajectory( 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 @@ -421,6 +439,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, and BPDA component-ablation accuracies.""" required = {"model", "config", "clean_acc", "PGD_acc"} if ( df_component is None @@ -429,9 +448,12 @@ def plot_component_ablation( ): 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=["clean_acc", "PGD_acc"], + value_vars=value_vars, var_name="Metric", value_name="Accuracy", ).dropna(subset=["Accuracy"]) @@ -460,6 +482,7 @@ def plot_component_ablation( 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 @@ -504,6 +527,7 @@ def plot_gradient_masking_summary( 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 @@ -561,6 +585,7 @@ def plot_confidence_margin_diagnostic( 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 @@ -613,6 +638,7 @@ def plot_results_heatmap( def add_derived_metrics(df: pd.DataFrame) -> pd.DataFrame: + """Add derived robustness and masking metrics to a results table.""" if df is None or df.empty or "model" not in df: return df df = df.copy() @@ -670,44 +696,14 @@ def add_derived_metrics(df: pd.DataFrame) -> pd.DataFrame: def add_paired_tests(df: pd.DataFrame) -> pd.DataFrame: - if df is None or df.empty or "correctness_vectors_path" not in df: - return df - df = df.copy() - rows = df.set_index("model", drop=False) - for index, row in df.iterrows(): - model = str(row["model"]) - architecture = model.split("_", 1)[0] - baseline_name = f"{architecture}_FP32" - if baseline_name not in rows.index: - continue - paths = ( - rows.loc[baseline_name, "correctness_vectors_path"], - row["correctness_vectors_path"], - ) - if not all(isinstance(path, str) and Path(path).exists() for path in paths): - continue - with np.load(paths[0]) as baseline, np.load(paths[1]) as variant: - for metric in sorted(set(baseline.files) & set(variant.files)): - a = np.asarray(baseline[metric], dtype=bool) - b = np.asarray(variant[metric], dtype=bool) - if a.shape != b.shape: - continue - baseline_only = int(np.sum(a & ~b)) - variant_only = int(np.sum(~a & b)) - discordant = baseline_only + variant_only - prefix = f"McNemar_vs_FP32_{metric}" - df.loc[index, f"{prefix}_a_only"] = baseline_only - df.loc[index, f"{prefix}_b_only"] = variant_only - df.loc[index, f"{prefix}_discordant"] = discordant - df.loc[index, f"{prefix}_p_value"] = ( - binomtest(baseline_only, discordant, 0.5).pvalue - if discordant - else 1.0 - ) - return df + """Add paired McNemar tests between each model variant and FP32 baseline.""" + return qstats.add_paired_mcnemar_tests( + df, baseline_name=qstats.fp32_baseline_name + ) def _read_if_present(path: Path) -> pd.DataFrame: + """Read a CSV file when it exists, otherwise return an empty table.""" try: return pd.read_csv(path) if path.exists() else pd.DataFrame() except (OSError, pd.errors.ParserError, pd.errors.EmptyDataError) as exc: @@ -716,16 +712,19 @@ def _read_if_present(path: Path) -> pd.DataFrame: def read_table(path) -> pd.DataFrame: + """Read an existing CSV table or return an empty table.""" return _read_if_present(Path(path)) def upsert_table(path, rows: pd.DataFrame, keys: Iterable[str]) -> pd.DataFrame: + """Merge rows into a keyed CSV table and persist the result.""" path = Path(path) combined = _merge_frames([_read_if_present(path), rows], list(keys)) return _write(combined, path) def _merge_frames(frames: Iterable[pd.DataFrame], keys: list[str]) -> pd.DataFrame: + """Merge multiple dataframes on shared key columns.""" frames = [frame for frame in frames if frame is not None and not frame.empty] if not frames: return pd.DataFrame() @@ -741,6 +740,7 @@ 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] ) -> pd.DataFrame: + """Combine a family of CSV files using filename patterns and keys.""" paths = { path.resolve() for pattern in patterns @@ -757,6 +757,7 @@ def _combine_csv_family( def _plot_metric_grid( df: pd.DataFrame, title: str, output: Path, id_columns=("model",), page_size=12 ) -> None: + """Render a paginated grid of metric plots for a dataframe.""" if df is None or df.empty: return numeric = [ @@ -799,6 +800,7 @@ def _plot_metric_grid( def _plot_model_subgraphs(df: pd.DataFrame, table_name: str, output_dir: Path) -> None: + """Plot model-level subgraphs for a combined result table.""" if df is None or df.empty or "model" not in df or df["model"].nunique() < 2: return subgraph_dir = output_dir / "subgraphs" / table_name @@ -808,6 +810,7 @@ def _plot_model_subgraphs(df: pd.DataFrame, table_name: str, output_dir: Path) - def plot_performance(df: pd.DataFrame, output: Path = PERFORMANCE_PLOT_PNG) -> None: + """Plot performance and resource metrics for evaluated models.""" columns = [ c for c in ( @@ -838,6 +841,7 @@ def plot_performance(df: pd.DataFrame, output: Path = PERFORMANCE_PLOT_PNG) -> N def plot_dither_sweep(df: pd.DataFrame, output: Path = DITHER_PLOT_PNG) -> None: + """Plot chaotic-dither sweep results.""" values = [c for c in ("clean_acc", "PGD_acc") if c in df] if df is None or df.empty or "dither_amplitude" not in df or not values: return @@ -865,6 +869,7 @@ def plot_dither_sweep(df: pd.DataFrame, output: Path = DITHER_PLOT_PNG) -> None: def plot_chunk_quantization_attacks( df: pd.DataFrame, output: Path = CHUNK_QUANT_PLOT_PNG ) -> None: + """Plot chunk-wise quantization attack results.""" if df is None or df.empty: return x = next( @@ -894,6 +899,7 @@ def plot_chunk_quantization_attacks( def plot_defense_comparison(df: pd.DataFrame, output: Path = DEFENSE_PLOT_PNG) -> None: + """Plot defense accuracy and attack-resistance comparisons.""" if df is None or df.empty: return model_col = next((c for c in ("model", "defense", "name") if c in df), None) @@ -922,6 +928,7 @@ def plot_defense_comparison(df: pd.DataFrame, output: Path = DEFENSE_PLOT_PNG) - 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) @@ -990,6 +997,7 @@ def combine_all(data_dir: Path) -> dict[str, pd.DataFrame]: 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)) @@ -1024,6 +1032,7 @@ def plot_all(dfs: dict[str, pd.DataFrame], output_dir: Path = DATA_DIR) -> None: def print_report(dfs: dict[str, pd.DataFrame]) -> None: + """Print a concise console report from combined result tables.""" print( pd.DataFrame( [ diff --git a/src/defense.py b/src/defense.py index 7dc501b..3ee7650 100644 --- a/src/defense.py +++ b/src/defense.py @@ -1,3 +1,13 @@ +"""Defense wrappers and defense-training utilities for QuantAdv experiments. + +The main runner uses this module for input transformations, randomized +smoothing, rejection-style guardrails, adversarial-example detectors, and +optional adversarial training. The lightweight quantized copy used for +adversarial training intentionally mirrors the public attributes of the main +QuantAdv fake-quantized layers so attacks can switch between hard rounding and +straight-through gradients. +""" + import copy import math @@ -12,53 +22,99 @@ def _cfg(name, default): + """Return a config value while keeping this module importable in isolation.""" return globals().get(name, default) def _device(): + """Return the configured torch device as a ``torch.device`` instance.""" return device if isinstance(device, torch.device) else torch.device(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) def denormalize_inputs(x): + """Map normalized CIFAR tensors back to pixel space.""" return x * CIFAR_STD.to(x.device) + CIFAR_MEAN.to(x.device) def make_attack(attack_cls, model, *args, **kwargs): + """Construct a torchattacks object for already-normalized CIFAR inputs.""" attack = attack_cls(model, *args, **kwargs) attack.set_normalization_used(mean=CIFAR_MEAN_VALUES, std=CIFAR_STD_VALUES) return attack -def _quantize_tensor(t, bits): +class _FakeQuantSTE(torch.autograd.Function): + """Round in the forward pass and use identity gradients in the backward pass.""" + + @staticmethod + def forward(ctx, x): + """Round values during the forward pass of the STE quantizer.""" + return torch.round(x) + + @staticmethod + def backward(ctx, grad_output): + """Pass gradients through the STE quantizer unchanged.""" + return grad_output + + +def _quantize_tensor(t, bits, use_ste=False): + """Symmetric per-tensor fake quantization used by defense-side copies.""" if bits is None: return t qmax = 2 ** (bits - 1) - 1 scale = torch.clamp( t.detach().abs().max() / qmax, min=_cfg("QUANT_SCALE_MIN", 1e-8) ) - q = torch.round(t / scale).clamp(-qmax - 1, qmax) + scaled = t / scale + q = (_FakeQuantSTE.apply(scaled) if use_ste else torch.round(scaled)).clamp( + -qmax - 1, qmax + ) return q * scale class _QuantConv2d(nn.Conv2d): + """Conv2d with fake-quantized weights and/or activations.""" + def forward(self, x): - w = _quantize_tensor(self.weight, self.bits) + """Run convolution with optional fake-quantized weights and activations.""" + w = ( + _quantize_tensor(self.weight, self.bits, self.use_ste) + if self.quant_weight + else self.weight + ) out = self._conv_forward(x, w, self.bias) - return _quantize_tensor(out, self.bits) + return ( + _quantize_tensor(out, self.bits, self.use_ste) + if self.quant_act + else out + ) class _QuantLinear(nn.Linear): + """Linear layer with fake-quantized weights and/or activations.""" + def forward(self, x): - w = _quantize_tensor(self.weight, self.bits) + """Run a linear projection with optional fake-quantized weights and activations.""" + w = ( + _quantize_tensor(self.weight, self.bits, self.use_ste) + if self.quant_weight + else self.weight + ) out = F.linear(x, w, self.bias) - return _quantize_tensor(out, self.bits) + return ( + _quantize_tensor(out, self.bits, self.use_ste) + if self.quant_act + else out + ) def _to_quant_module(module, bits): + """Convert supported leaf layers to defense-side fake-quantized layers.""" if isinstance(module, nn.Conv2d): new = _QuantConv2d( module.in_channels, @@ -75,6 +131,9 @@ def _to_quant_module(module, bits): if module.bias is not None: new.bias = module.bias new.bits = bits + new.use_ste = _cfg("QUANT_DEFAULT_USE_STE", False) + new.quant_weight = _cfg("QUANT_DEFAULT_WEIGHT", True) + new.quant_act = _cfg("QUANT_DEFAULT_ACT", True) return new if isinstance(module, nn.Linear): @@ -85,12 +144,16 @@ def _to_quant_module(module, bits): if module.bias is not None: new.bias = module.bias new.bits = bits + new.use_ste = _cfg("QUANT_DEFAULT_USE_STE", False) + new.quant_weight = _cfg("QUANT_DEFAULT_WEIGHT", True) + new.quant_act = _cfg("QUANT_DEFAULT_ACT", True) return new return None def _replace_quant_modules(module, bits): + """Recursively replace Conv2d/Linear children with quantized equivalents.""" for name, child in list(module.named_children()): replacement = _to_quant_module(child, bits) if replacement is None: @@ -100,12 +163,24 @@ def _replace_quant_modules(module, bits): def quantized_copy(model, bits): + """Return a model copy with Conv2d/Linear layers fake-quantized.""" m = copy.deepcopy(model) _replace_quant_modules(m, bits) return m.to(_device()) +def _set_quant_ste(model, flag): + """Toggle straight-through gradients on defense-side quantized modules.""" + toggled = 0 + for module in model.modules(): + if hasattr(module, "use_ste") and hasattr(module, "bits"): + module.use_ste = flag + toggled += 1 + return toggled + + def prepare_adversarial_training(base_model, loader, bits=None): + """Fine-tune a copy of ``base_model`` on PGD adversarial examples.""" dev = _device() model = ( quantized_copy(base_model, bits) @@ -135,28 +210,34 @@ def prepare_adversarial_training(base_model, loader, bits=None): for x, y in loader: x = x.to(dev, non_blocking=True) y = y.to(dev, non_blocking=True) - model.eval() - x_adv = attack(x, y).detach() - model.train() - - opt.zero_grad(set_to_none=True) - with autocast(device_type=dev.type): - loss = F.cross_entropy(model(x_adv), y) - scaler.scale(loss).backward() - scaler.step(opt) - scaler.update() + try: + model.eval() + _set_quant_ste(model, True) + x_adv = attack(x, y).detach() + model.train() + + opt.zero_grad(set_to_none=True) + with autocast(device_type=dev.type): + loss = F.cross_entropy(model(x_adv), y) + scaler.scale(loss).backward() + scaler.step(opt) + scaler.update() + finally: + _set_quant_ste(model, False) return model.eval() class SanitizedModel(nn.Module): def __init__(self, model): + """Wrap a classifier with deterministic input sanitization settings.""" super().__init__() self.model = model self.resize = _cfg("DEFENSE_SANITIZE_SIZE", 28) self.bits = _cfg("DEFENSE_SANITIZE_BITS", 6) def forward(self, x): + """Sanitize inputs in pixel space before normalized model inference.""" pixels = denormalize_inputs(x).clamp(0.0, 1.0) pixels = T.functional.gaussian_blur(pixels, kernel_size=3) pixels = F.interpolate( @@ -178,12 +259,14 @@ def forward(self, x): class SmoothedModel(nn.Module): def __init__(self, model): + """Wrap a classifier with randomized smoothing parameters.""" super().__init__() self.model = model self.sigma = _cfg("DEFENSE_SMOOTH_SIGMA", 0.12) self.samples = _cfg("DEFENSE_SMOOTH_SAMPLES", 8) def forward(self, x): + """Average noisy predictions at evaluation time and sample once during training.""" if not self.training: logits = 0.0 for _ in range(self.samples): @@ -200,12 +283,14 @@ def forward(self, x): class GuardrailModel(nn.Module): def __init__(self, model): + """Wrap a classifier with confidence-threshold rejection settings.""" super().__init__() self.model = model self.conf_threshold = _cfg("DEFENSE_GUARD_CONF_THRESHOLD", 0.55) self.reject_logit = _cfg("DEFENSE_REJECT_LOGIT", -1e4) def forward(self, x): + """Replace low-confidence predictions with the rejection logit.""" logits = self.model(x) conf = F.softmax(logits, dim=1).max(dim=1).values reject = conf < self.conf_threshold @@ -216,6 +301,7 @@ def forward(self, x): class DetectorCNN(nn.Module): def __init__(self): + """Build the convolutional binary adversarial-example detector.""" super().__init__() self.net = nn.Sequential( nn.Conv2d(3, 32, 3, padding=1), @@ -232,10 +318,12 @@ def __init__(self): ) def forward(self, x): + """Return clean-versus-adversarial detector logits for a batch.""" return self.net(x) def train_adversarial_detector(base_model, loader): + """Train a binary detector to distinguish clean and adversarial inputs.""" dev = _device() base_model = base_model.to(dev).eval() detector = DetectorCNN().to(dev).train() @@ -285,6 +373,7 @@ def train_adversarial_detector(base_model, loader): class DetectGuardModel(nn.Module): def __init__(self, model, detector): + """Wrap a classifier with a learned detector-based rejection guard.""" super().__init__() self.model = model self.detector = detector @@ -292,6 +381,7 @@ def __init__(self, model, detector): self.reject_logit = _cfg("DEFENSE_REJECT_LOGIT", -1e4) def forward(self, x): + """Reject inputs that the detector classifies as adversarial.""" logits = self.model(x) detector_prob = F.softmax(self.detector(x), dim=1)[:, 1] reject = detector_prob > self.threshold @@ -301,6 +391,7 @@ def forward(self, x): def run_certified_accuracy(model, loader): + """Estimate smoothed-model accuracy and a radius proxy.""" dev = _device() model.eval() correct = 0 @@ -327,14 +418,17 @@ def run_certified_accuracy(model, loader): def run_guardrail_flagging_rate(model, loader, attack=None): + """Measure how often a confidence guardrail flags clean and attacked inputs.""" return _run_flagging_rate(model, loader, attack, detector_attr=None) def run_detector_catch_rate(model, loader, attack=None): + """Measure how often the learned detector catches adversarial inputs.""" return _run_flagging_rate(model, loader, attack, detector_attr="detector") def _run_flagging_rate(model, loader, attack=None, detector_attr=None): + """Compute clean and adversarial flagging rates for a rejection model.""" dev = _device() model.eval() clean_flags = 0 @@ -362,6 +456,7 @@ def _run_flagging_rate(model, loader, attack=None, detector_attr=None): def _flag_mask(model, x, detector_attr): + """Return the boolean mask of inputs rejected by a guard or detector.""" if detector_attr is not None and hasattr(model, detector_attr): prob = F.softmax(model.detector(x), dim=1)[:, 1] return prob > model.threshold diff --git a/src/stats.py b/src/stats.py new file mode 100644 index 0000000..ad61c0c --- /dev/null +++ b/src/stats.py @@ -0,0 +1,146 @@ +"""Statistical helpers for QuantAdv result tables. + +The experiment runner stores per-example correctness vectors so robustness +comparisons can use paired tests rather than independent-sample summaries. +SciPy provides the binomial confidence intervals and exact binomial tests used +here; this module keeps those details out of the model and plotting code. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Callable, Iterable + +import numpy as np +import pandas as pd +from scipy.stats import binomtest + + +def _as_bool_vector(vector: Iterable[bool]) -> np.ndarray: + """Return a one-dimensional boolean NumPy array.""" + return np.asarray(vector, dtype=bool).reshape(-1) + + +def binomial_interval(correct: int, total: int, confidence: float, method: str): + """Return a SciPy binomial proportion interval as plain Python floats.""" + if total <= 0: + return None, None + interval = binomtest(correct, total).proportion_ci( + confidence_level=confidence, method=method + ) + return float(interval.low), float(interval.high) + + +def binomial_statistics(correct_vector, confidence: float = 0.95) -> dict[str, float]: + """Summarize a correctness vector with exact and Wilson intervals.""" + vector = _as_bool_vector(correct_vector) + total = int(vector.size) + correct = int(vector.sum()) + wilson_low, wilson_high = binomial_interval( + correct, total, confidence, method="wilson" + ) + exact_low, exact_high = binomial_interval( + correct, total, confidence, method="exact" + ) + return { + "n": total, + "correct": correct, + "wilson_low": wilson_low, + "wilson_high": wilson_high, + "wilson_pm": ( + (wilson_high - wilson_low) / 2 if wilson_low is not None else None + ), + "cp_low": exact_low, + "cp_high": exact_high, + "cp_pm": ((exact_high - exact_low) / 2 if exact_low is not None else None), + } + + +def add_binomial_statistics( + results: dict, metric: str, correct_vector, confidence: float = 0.95 +) -> None: + """Attach prefixed binomial statistics for ``metric`` to ``results``.""" + summary = binomial_statistics(correct_vector, confidence=confidence) + results.update({f"{metric}_{key}": value for key, value in summary.items()}) + + +def mcnemar_exact(vector_a, vector_b) -> dict[str, float]: + """Run the exact two-sided McNemar test over paired correctness vectors.""" + 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}") + + a_only = int(np.sum(a & ~b)) + b_only = int(np.sum(~a & b)) + discordant = a_only + b_only + p_value = ( + float(binomtest(min(a_only, b_only), discordant, 0.5).pvalue) + if discordant + else 1.0 + ) + return { + "a_only": a_only, + "b_only": b_only, + "discordant": discordant, + "p_value": p_value, + } + + +def save_correctness_vectors(model_name: str, vectors: dict, output_dir) -> str: + """Persist per-example correctness vectors for later paired tests.""" + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + safe_name = "".join(c if c.isalnum() or c in "._-" else "_" for c in model_name) + path = output_dir / f"{safe_name}.npz" + np.savez_compressed( + path, **{key: _as_bool_vector(value) for key, value in vectors.items()} + ) + return str(path) + + +def add_paired_mcnemar_tests( + df: pd.DataFrame, + baseline_name: Callable[[str], str], + path_column: str = "correctness_vectors_path", + prefix_template: str = "McNemar_vs_FP32_{metric}", +) -> pd.DataFrame: + """Add paired McNemar test columns comparing each row with its baseline row.""" + if df is None or df.empty or path_column not in df: + return df + + df = df.copy() + rows = df.set_index("model", drop=False) + for index, row in df.iterrows(): + model = str(row["model"]) + baseline = baseline_name(model) + if baseline not in rows.index: + continue + + baseline_path = rows.loc[baseline, path_column] + variant_path = row[path_column] + if not all( + isinstance(path, str) and Path(path).exists() + for path in (baseline_path, variant_path) + ): + continue + + 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)): + a = np.asarray(baseline_vectors[metric], dtype=bool) + b = np.asarray(variant_vectors[metric], dtype=bool) + if a.shape != b.shape: + continue + test = mcnemar_exact(a, b) + prefix = prefix_template.format(metric=metric) + for key, value in test.items(): + df.loc[index, f"{prefix}_{key}"] = value + return df + + +def fp32_baseline_name(model_name: str) -> str: + """Return the FP32 baseline row name for a QuantAdv model variant.""" + architecture = str(model_name).split("_", 1)[0] + return f"{architecture}_FP32"