diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..fe61c0e --- /dev/null +++ b/.gitignore @@ -0,0 +1,19 @@ +data/ +graphs/ +cifar10/ +mnist/ + + +__pycache__/* +cifar10-early/ +imagenet/ +mnist-early/ +nes_attack.py +nes_results_runs/ +plots/ +run_logs/ +scripts/ + +*.csv +*.sh +*.zip \ No newline at end of file diff --git a/README.md b/README.md index c8ecaba..ae3cae2 100644 --- a/README.md +++ b/README.md @@ -1,35 +1,182 @@ -# ZOO: Zeroth Order Optimization Based Adversarial Black Box Attack (PyTorch) -This repository contains the PyTorch implementation of Zeroth Order Optimization Based Adversarial Black Box Attack(https://arxiv.org/abs/1708.03999) using MNIST and CIFAR10 dataset. This is the exact replica as far possible of the ZOO Attack (https://github.com/IBM/ZOO-Attack) which was originally implemented in Tensorflow. The results match almost as same as the paper evaluation results for MNIST and CIFAR10 for both targeted and untargeted attack all with 100% success rate on the 7 layer CNNs model trained on MNIST with 99.5% val accuracy and on CIFAR10 with 80% val accuracy as done in the original paper work. Both ZOO_Adam and ZOO_Newton methods of Coordinate Descent Solvers are implemented. +# ZOO Black-Box Adversarial Attack -**Note: This doesn't contain implementation of importance sampling, hierarchical attack, and dimentional reduction right now (as its mainly needed for large image sized dataset like ImageNet). For larger dataset google colab viewer [link](https://colab.research.google.com/drive/1_rDu0LNAI9kg490rV9PMxb55Wwbj13aD?usp=sharing) [NOT TESTED]** +Implementation of the ZOO (Zeroth Order Optimization) black-box adversarial attack from [Chen et al., 2017](https://arxiv.org/abs/1708.03999), supporting MNIST, CIFAR-10, and ImageNet with multiple coordinate-descent solvers. Forked from the pyTorch implementation following paper's repository. -## Setup and train models -The code is tested with Python 3.7.6 and PyTorch 1.6.0. The following packages are required: -``` -python pip install --upgrade pip -pip install torch==1.6.0 torchsummary==1.5.1 torchvision==0.7.0 -pip install numpy matplotlib -``` -To prepare model and datasets of MNIST and CIFAR10 +--- + +## Installation + +```bash +pip install torch torchvision torchsummary numpy matplotlib scikit-image pillow ``` + +The pre-trained MNIST and CIFAR-10 models are already included in `models/`. The MNIST and CIFAR-10 datasets are downloaded automatically by torchvision on first run, and the ImageNette subset used for ImageNet is fetched automatically when `--dataset imagenet` is selected. + +To retrain the classifiers from scratch instead of using the bundled checkpoints: + +```bash python setup_mnist_model.py python setup_cifar10_model.py ``` -## Run attacks -To run the attacks run the + +--- + +## Running attacks + +### Reproducing the paper results + +The results reported in the paper come from sweeping every solver across every model, both untargeted and targeted. A single call runs the whole sweep: + +```bash +python run.py # all solvers, all models, both untargeted + targeted, plot generation +``` + +`run.py` calls the attack below once per (dataset, solver, mode) combination and writes each run's output to its own folder (see [Results](#results)). Restrict the sweep with `--modes untargeted` or `--modes targeted`, narrow it with `--datasets` / `--solvers` / `--samples`, or pass `--dry-run` to preview the commands first. + +### Basic usage + +```bash +python zoo_l2_attack_black.py --dataset --solver --samples +``` + +### All solvers — MNIST + +```bash +python zoo_l2_attack_black.py --dataset mnist --solver adam --samples 10 +python zoo_l2_attack_black.py --dataset mnist --solver newton --samples 10 +python zoo_l2_attack_black.py --dataset mnist --solver sgd --samples 10 +python zoo_l2_attack_black.py --dataset mnist --solver sgdsign --samples 10 +python zoo_l2_attack_black.py --dataset mnist --solver signum --samples 10 +python zoo_l2_attack_black.py --dataset mnist --solver lion --samples 10 +python zoo_l2_attack_black.py --dataset mnist --solver adahessian --samples 10 +``` + +### All solvers — CIFAR-10 + +```bash +python zoo_l2_attack_black.py --dataset cifar10 --solver adam --samples 10 +python zoo_l2_attack_black.py --dataset cifar10 --solver newton --samples 10 +python zoo_l2_attack_black.py --dataset cifar10 --solver sgd --samples 10 +python zoo_l2_attack_black.py --dataset cifar10 --solver sgdsign --samples 10 +python zoo_l2_attack_black.py --dataset cifar10 --solver signum --samples 10 +python zoo_l2_attack_black.py --dataset cifar10 --solver lion --samples 10 +python zoo_l2_attack_black.py --dataset cifar10 --solver adahessian --samples 10 +``` + +### Targeted attack + +```bash +python zoo_l2_attack_black.py --dataset mnist --solver adam --samples 10 --targeted +python zoo_l2_attack_black.py --dataset cifar10 --solver adam --samples 10 --targeted +``` + +--- + +## Early stopping + +ZOO uses an L2 objective with binary search over the trade-off constant, so by default the attack runs the full iteration budget to minimise distortion. Passing `--early-stop` stops a sample as soon as the adversarial example successfully fools the model, which measures how many queries each solver needs to break the model, a strong indicator of solver efficiency. The number of queries used is saved in `results.json`. + +--- + +## All arguments + +| Argument | Default | Description | +|---|---|---| +| `--dataset` | `cifar10` | `mnist`, `cifar10`, or `imagenet` | +| `--solver` | `adam` | `adam`, `newton`, `sgd`, `sgdsign`, `signum`, `lion`, `adahessian` | +| `--samples` | `10` | Number of **source** images. Untargeted: N attacks. Targeted: each source is attacked toward every other class, so for 10-class MNIST/CIFAR-10 that's N × 9 (e.g. 10 → 90 attacks) | +| `--start` | `6` | Offset into the test set | +| `--targeted` | `False` | Targeted attack (default: untargeted) | +| `--targeted-k` | `None` | In targeted mode, use at most K random target classes per source image | +| `--imagenet_dir` | `None` | Path to ImageNet val directory (used when `--dataset imagenet`) | +| `--batch-size` | `128` | Coordinates updated per optimization step | +| `--max-iter` | `1000` | Max iterations per binary-search stage | +| `--const` | `0.01` | Initial attack trade-off constant | +| `--confidence` | `0.0` | CW confidence margin | +| `--early-stop-iters` | `100` | Abort-check cadence | +| `--binary-search-steps` | `9` | Binary search steps over `const` | +| `--adam-beta1` | `0.9` | Optimizer beta1 | +| `--adam-beta2` | `0.999` | Optimizer beta2 | +| `--step-size` | auto | Override the solver's default coordinate step size | +| `--early-stop` | `False` | Stop each sample once the attack succeeds | + +Default coordinate step size per solver (override with `--step-size`): + +| Solver | step size | +|---|---| +| adam | 0.01 | +| newton | 0.01 | +| sgd | 0.01 | +| sgdsign | 0.015 | +| signum | 0.015 | +| lion | 0.015 | +| adahessian | 0.01 | + +--- + +## Results + +Results are saved in `///`: + +``` +cifar10/ + untargeted/adam/ + results.json : success rate, queries, distortion, PSNR, SSIM, time + original_0.png : original image + adversarial_0.png : adversarial image + grid_cifar10_untargeted_adam.png : side-by-side grid with class labels ``` -python zoo_l2_attack_black.py + +`results.json` includes a `queries` block: + +```json +"queries": { + "per_sample": [1200, 800, ...], + "mean_on_success": 1050.0, + "budget": 50000 +} ``` -Both untargeted and targeted attack are accessible via above code all the changes (comment/uncomment) for transition from ZOO_Adam/ZOO_Newton or CIFAR10/MNIST are from line 259-262, 270/271, 274-277 and for visualization of example generated, line 307/329. For more details go through the code zoo_l2_attack_black.py and the paper https://arxiv.org/abs/1708.03999 -## Sample Results -#### ZOO_Adam -##### Untargeted on CIFAR10 -![](/sample_results/adam_untargeted_cifar10.png) -##### Untargeted on CIFAR10 -![](/sample_results/adam_untargeted_mnist.png) -#### ZOO_Newton -##### Targeted on MNIST -![](/sample_results/newton_targeted_mnist.png) +along with `success_rate_pct`, `total_distortion`, `time_mins`, the per-sample `mse`, `mae`, `psnr`, and `ssim` blocks, and an `attack_params` block recording the hyperparameters used. + +--- + +## Plots + +`plot_zoo_metrics_summary.py` reproduces the per-optimizer metric bar chart (MAE, MSE, PSNR, SSIM, L-inf and L2 distortion) for one dataset and threat model. MAE/MSE/PSNR/SSIM are read from each solver's `results.json`; L2 and L-inf are computed from the saved `original_*.png` / `adversarial_*.png` pairs. Folders whose name contains `pgd` are skipped. + +```bash +python plot_zoo_metrics_summary.py --dataset cifar10 --attack-type targeted +``` + +The figure is written to `plots/`. `run.py` renders these charts automatically after the sweep finishes, one per dataset and threat model. + + +--- + +## Optimizer intuition + +We deliver `opt_visualization.py` for gaining intuition behind the optimizers' behaviour, independent of the attack setting. It minimises classic 2-D test functions — Rosenbrock, Himmelblau, Beale, and Sphere, with the same optimizer family used in the attacks (Adam, Newton, SGD, SignSGD, Signum, Lion), recording each optimizer's trajectory and plotting it over a filled 2-D contour and a 3-D surface, plus a convergence curve. + +```bash +python opt_visualization.py +``` + +Switch the target function, starting point, and per-solver learning rates via the settings at the bottom of the file. A companion script, `grad_est_optimization.py`, runs the same optimizers on the same functions but contrasts exact analytical gradients against finite-difference estimates, so how gradient noise reshapes each optimizer's path can be observed. + +Example visualizations: +

+ + + +

+ +--- + +## Sample results + +Every run writes a labelled grid image (`grid___.png`) into its output folder, showing each source image with its `original → adversarial` prediction. A few precomputed examples are committed under [`sample_results/`](sample_results), so you can see what the attack produces without running anything. +CIFAR-10, Adam (untargeted): +![CIFAR-10 adversarial examples — Adam, untargeted](sample_results/grid_cifar10_targeted_adam.png) \ No newline at end of file diff --git a/__pycache__/setup_cifar10_model.cpython-39.pyc b/__pycache__/setup_cifar10_model.cpython-39.pyc new file mode 100644 index 0000000..cdb6155 Binary files /dev/null and b/__pycache__/setup_cifar10_model.cpython-39.pyc differ diff --git a/__pycache__/setup_mnist_model.cpython-39.pyc b/__pycache__/setup_mnist_model.cpython-39.pyc new file mode 100644 index 0000000..13f0e6d Binary files /dev/null and b/__pycache__/setup_mnist_model.cpython-39.pyc differ diff --git a/grad_est_optimization.py b/grad_est_optimization.py new file mode 100644 index 0000000..8274dda --- /dev/null +++ b/grad_est_optimization.py @@ -0,0 +1,560 @@ +""" +grad_est_optimization.py +------------------------ +Compares optimiser trajectories when using: + • Exact (analytical) gradients + • Finite-difference (FD) gradient estimates + +Five optimisers are run under both gradient regimes: + adam, newton, sgd-momentum, signsgd, lion + +Three output PDFs are produced (same style as opt_visualization.py): + *_contour2d.pdf — 2-D filled contour + paths + *_surface3d.pdf — 3-D surface + paths + *_convergence.pdf — semi-log convergence + global-minimum line + +Exact vs FD is distinguished by line style: + solid ─── exact gradient + dashed - - FD gradient + +Usage +----- + python grad_est_optimization.py + +To change the test function edit FUNCTION, TRUE_MIN, START, PLOT_XLIM/YLIM +at the bottom of the file. +Available presets: rosenbrock | himmelblau | beale | sphere +""" + +import numpy as np +import matplotlib.pyplot as plt +from mpl_toolkits.mplot3d import Axes3D # noqa: F401 +from dataclasses import dataclass, field +from typing import Callable, List, Tuple + +# --------------------------------------------------------------------------- +# Global style +# --------------------------------------------------------------------------- +plt.rcParams["font.family"] = "Times New Roman" +plt.rcParams["mathtext.fontset"] = "stix" + +_TNR = {"fontfamily": "Times New Roman"} + +COLORS = { + "Adam": "#e41a1c", + "Newton": "#377eb8", + "SGD": "#4daf4a", + "SGDSign": "#ff7f00", + "Signum": "#FFD700", + "Lion": "#984ea3", +} + +EXACT_LS = "-" # solid → exact gradient +FD_LS = "--" # dashed → finite-difference gradient + + +# --------------------------------------------------------------------------- +# Data container +# --------------------------------------------------------------------------- +@dataclass +class History: + name: str + grad_type: str # "exact" or "fd" + xs: List[float] = field(default_factory=list) + ys: List[float] = field(default_factory=list) + fs: List[float] = field(default_factory=list) + + def record(self, x, y, f): + self.xs.append(float(x)) + self.ys.append(float(y)) + self.fs.append(float(f)) + + def final(self): + return self.xs[-1], self.ys[-1], self.fs[-1] + + @property + def label(self): + return f"{self.name} ({'exact' if self.grad_type == 'exact' else 'FD'})" + + @property + def ls(self): + return EXACT_LS if self.grad_type == "exact" else FD_LS + + +# --------------------------------------------------------------------------- +# Preset 2-D functions with analytical gradients +# --------------------------------------------------------------------------- + +def rosenbrock(x, y, a=1, b=100): + """Global minimum at (1, 1).""" + return (a - x) ** 2 + b * (y - x ** 2) ** 2 + +def rosenbrock_grad(x, y, a=1, b=100): + gx = -2 * (a - x) - 4 * b * x * (y - x ** 2) + gy = 2 * b * (y - x ** 2) + return gx, gy + +def rosenbrock_hess_diag(x, y, a=1, b=100): + hxx = 2 - 4 * b * y + 12 * b * x ** 2 + hyy = 2 * b + return hxx, hyy + + +def himmelblau(x, y): + """Four equal minima.""" + return (x ** 2 + y - 11) ** 2 + (x + y ** 2 - 7) ** 2 + +def himmelblau_grad(x, y): + gx = 4 * x * (x ** 2 + y - 11) + 2 * (x + y ** 2 - 7) + gy = 2 * (x ** 2 + y - 11) + 4 * y * (x + y ** 2 - 7) + return gx, gy + +def himmelblau_hess_diag(x, y): + hxx = 12 * x ** 2 + 4 * y - 44 + 2 + hyy = 2 + 12 * y ** 2 + 4 * x - 28 + return hxx, hyy + + +def beale(x, y): + """Global minimum at (3, 0.5).""" + t1 = 1.5 - x + x * y + t2 = 2.25 - x + x * y ** 2 + t3 = 2.625 - x + x * y ** 3 + return t1 ** 2 + t2 ** 2 + t3 ** 2 + +def beale_grad(x, y): + t1 = 1.5 - x + x * y + t2 = 2.25 - x + x * y ** 2 + t3 = 2.625 - x + x * y ** 3 + gx = (2 * t1 * (y - 1) + + 2 * t2 * (y ** 2 - 1) + + 2 * t3 * (y ** 3 - 1)) + gy = (2 * t1 * x + + 2 * t2 * 2 * x * y + + 2 * t3 * 3 * x * y ** 2) + return gx, gy + +def beale_hess_diag(x, y): + t1 = 1.5 - x + x * y + t2 = 2.25 - x + x * y ** 2 + t3 = 2.625 - x + x * y ** 3 + hxx = (2 * (y - 1) ** 2 + + 2 * (y ** 2 - 1) ** 2 + + 2 * (y ** 3 - 1) ** 2) + hyy = (2 * x ** 2 + + 2 * (2 * x * y) ** 2 + 2 * t2 * 2 * x + + 2 * (3 * x * y ** 2) ** 2 + 2 * t3 * 6 * x * y) + return hxx, hyy + + +def sphere(x, y): + """Global minimum at (0, 0).""" + return x ** 2 + y ** 2 + +def sphere_grad(x, y): + return 2 * x, 2 * y + +def sphere_hess_diag(x, y): + return 2.0, 2.0 + + +# Registry: function → (grad_fn, hess_diag_fn) +GRAD_REGISTRY = { + rosenbrock: (rosenbrock_grad, rosenbrock_hess_diag), + himmelblau: (himmelblau_grad, himmelblau_hess_diag), + beale: (beale_grad, beale_hess_diag), + sphere: (sphere_grad, sphere_hess_diag), +} + + +# --------------------------------------------------------------------------- +# Finite-difference gradient & Hessian diagonal +# --------------------------------------------------------------------------- +def grad_fd(f, x, y, h=1e-5): + gx = (f(x + h, y) - f(x - h, y)) / (2 * h) + gy = (f(x, y + h) - f(x, y - h)) / (2 * h) + return gx, gy + +def hess_diag_fd(f, x, y, h=1e-5): + f0 = f(x, y) + hxx = max(abs((f(x + h, y) - 2 * f0 + f(x - h, y)) / h ** 2), 0.1) + hyy = max(abs((f(x, y + h) - 2 * f0 + f(x, y - h)) / h ** 2), 0.1) + return hxx, hyy + + +def _safe_hess(hxx, hyy): + return max(abs(hxx), 0.1), max(abs(hyy), 0.1) + + +# --------------------------------------------------------------------------- +# Generic optimisers — accept a grad_fn and hess_fn callable +# --------------------------------------------------------------------------- + +def run_adam(f, grad_fn, x0, y0, name="Adam", grad_type="exact", + lr=0.05, beta1=0.9, beta2=0.999, eps=1e-8, steps=500): + h = History(name, grad_type) + x, y = x0, y0 + mx = my = vx = vy = 0.0 + for t in range(1, steps + 1): + h.record(x, y, f(x, y)) + gx, gy = grad_fn(x, y) + mx = beta1 * mx + (1 - beta1) * gx + my = beta1 * my + (1 - beta1) * gy + vx = beta2 * vx + (1 - beta2) * gx ** 2 + vy = beta2 * vy + (1 - beta2) * gy ** 2 + corr = np.sqrt(1 - beta2 ** t) / (1 - beta1 ** t) + x -= lr * corr * mx / (np.sqrt(vx) + eps) + y -= lr * corr * my / (np.sqrt(vy) + eps) + h.record(x, y, f(x, y)) + return h + + +def run_newton(f, grad_fn, hess_fn, x0, y0, name="Newton", grad_type="exact", + lr=1.0, steps=500): + h = History(name, grad_type) + x, y = x0, y0 + for _ in range(steps): + h.record(x, y, f(x, y)) + gx, gy = grad_fn(x, y) + hxx, hyy = _safe_hess(*hess_fn(x, y)) + x -= lr * gx / hxx + y -= lr * gy / hyy + h.record(x, y, f(x, y)) + return h + + +def run_sgd(f, grad_fn, x0, y0, name="SGD", grad_type="exact", + lr=0.01, steps=500): + """Vanilla gradient descent — no momentum.""" + h = History(name, grad_type) + x, y = x0, y0 + for _ in range(steps): + h.record(x, y, f(x, y)) + gx, gy = grad_fn(x, y) + x -= lr * gx + y -= lr * gy + h.record(x, y, f(x, y)) + return h + +def run_signsgd(f, grad_fn, x0, y0, name="SGDSign", grad_type="exact", + lr=0.01, steps=500): + """SGDSign — step is lr * sign(g), no momentum.""" + h = History(name, grad_type) + x, y = x0, y0 + for _ in range(steps): + h.record(x, y, f(x, y)) + gx, gy = grad_fn(x, y) + x -= lr * np.sign(gx) + y -= lr * np.sign(gy) + h.record(x, y, f(x, y)) + return h + + +def run_signnum(f, grad_fn, x0, y0, name="Signum", grad_type="exact", + lr=0.01, beta1=0.9, steps=500): + """Signum — step is lr * sign(m), where m is an EMA of gradients.""" + h = History(name, grad_type) + x, y = x0, y0 + mx = my = 0.0 + for _ in range(steps): + h.record(x, y, f(x, y)) + gx, gy = grad_fn(x, y) + mx = beta1 * mx + (1 - beta1) * gx + my = beta1 * my + (1 - beta1) * gy + x -= lr * np.sign(mx) + y -= lr * np.sign(my) + h.record(x, y, f(x, y)) + return h + + +def run_lion(f, grad_fn, x0, y0, name="Lion", grad_type="exact", + lr=0.001, beta1=0.9, beta2=0.99, steps=500): + h = History(name, grad_type) + x, y = x0, y0 + mx = my = 0.0 + for _ in range(steps): + h.record(x, y, f(x, y)) + gx, gy = grad_fn(x, y) + ux = np.sign(beta1 * mx + (1 - beta1) * gx) + uy = np.sign(beta1 * my + (1 - beta1) * gy) + x -= lr * ux + y -= lr * uy + mx = beta2 * mx + (1 - beta2) * gx + my = beta2 * my + (1 - beta2) * gy + h.record(x, y, f(x, y)) + return h + + +# --------------------------------------------------------------------------- +# Helper: build exact and FD gradient/Hessian callables for a function +# --------------------------------------------------------------------------- +def make_grad_fns(f): + """Return (exact_grad, exact_hess, fd_grad, fd_hess) for f.""" + exact_grad, exact_hess = GRAD_REGISTRY[f] + fd_grad = lambda x, y: grad_fd(f, x, y) + fd_hess = lambda x, y: hess_diag_fd(f, x, y) + return exact_grad, exact_hess, fd_grad, fd_hess + + +# --------------------------------------------------------------------------- +# Run all optimisers under both gradient regimes +# --------------------------------------------------------------------------- +def run_all(f, x0, y0, steps, + adam_lr, newton_lr, sgd_lr, + sgdsign_lr, signum_lr, lion_lr): + eg, eh, fg, fh = make_grad_fns(f) + histories = [] + for gtype, gfn, hfn in [("exact", eg, eh), ("fd", fg, fh)]: + histories += [ + run_adam (f, gfn, x0, y0, "Adam", gtype, lr=adam_lr, steps=steps), + run_newton (f, gfn, hfn, x0, y0, "Newton", gtype, lr=newton_lr, steps=steps), + run_sgd (f, gfn, x0, y0, "SGD", gtype, lr=sgd_lr, steps=steps), + run_signsgd (f, gfn, x0, y0, "SGDSign", gtype, lr=sgdsign_lr, steps=steps), + run_signnum (f, gfn, x0, y0, "Signum", gtype, lr=signum_lr, steps=steps), + run_lion (f, gfn, x0, y0, "Lion", gtype, lr=lion_lr, steps=steps), + ] + return histories + + +# --------------------------------------------------------------------------- +# Visualisation helpers +# --------------------------------------------------------------------------- +def _apply_tnr(ax): + for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] + + ax.get_xticklabels() + ax.get_yticklabels()): + item.set_fontfamily("Times New Roman") + leg = ax.get_legend() + if leg: + plt.setp(leg.get_texts(), fontfamily="Times New Roman") + + +def _apply_tnr_3d(ax): + for item in ([ax.title, ax.xaxis.label, ax.yaxis.label, ax.zaxis.label] + + ax.get_xticklabels() + ax.get_yticklabels() + + ax.get_zticklabels()): + item.set_fontfamily("Times New Roman") + leg = ax.get_legend() + if leg: + plt.setp(leg.get_texts(), fontfamily="Times New Roman") + + +# --------------------------------------------------------------------------- +# Plot 1 — 2-D contour +# --------------------------------------------------------------------------- +def plot_contour_2d(histories: List[History], f: Callable, + xlim, ylim, title="", resolution=200, + true_min=None): + xs = np.linspace(*xlim, resolution) + ys = np.linspace(*ylim, resolution) + X, Y = np.meshgrid(xs, ys) + Z = f(X, Y) + levels = np.unique(np.percentile(Z, np.linspace(0, 95, 30))) + + fig, ax = plt.subplots(figsize=(8, 6)) + cf = ax.contourf(X, Y, Z, levels=levels, cmap="viridis", alpha=0.75) + cbar = plt.colorbar(cf, ax=ax, label="f(x, y)") + cbar.ax.yaxis.label.set_fontfamily("Times New Roman") + plt.setp(cbar.ax.get_yticklabels(), fontfamily="Times New Roman") + ax.contour(X, Y, Z, levels=levels, colors="white", linewidths=0.3, alpha=0.4) + + # Draw paths + for h in histories: + color = COLORS[h.name] + ax.plot(h.xs, h.ys, h.ls + "o", color=color, + markersize=2, linewidth=1.4, alpha=0.80, + label=h.label) + ax.plot(h.xs[0], h.ys[0], "o", color=color, markersize=7, zorder=8) + + # Final stars with overlap jitter + x_span = xlim[1] - xlim[0] + y_span = ylim[1] - ylim[0] + nudge = 0.018 * max(x_span, y_span) + dirs = [(1,0),(0,1),(-1,0),(0,-1),(1,1),(-1,1),(1,-1),(-1,-1)] + finals = [(h.xs[-1], h.ys[-1], COLORS[h.name], h.grad_type) for h in histories] + offsets = [[0.0, 0.0] for _ in finals] + for i in range(len(finals)): + for j in range(i + 1, len(finals)): + dx = finals[i][0] - finals[j][0] + dy = finals[i][1] - finals[j][1] + if (dx**2 + dy**2) ** 0.5 < nudge * 1.5: + di = dirs[i % len(dirs)] + dj = dirs[(i + 4) % len(dirs)] + offsets[i][0] += di[0] * nudge + offsets[i][1] += di[1] * nudge + offsets[j][0] += dj[0] * nudge + offsets[j][1] += dj[1] * nudge + for (fx, fy, color, gtype), (ox, oy) in zip(finals, offsets): + marker = "*" if gtype == "exact" else "P" # star=exact, filled-plus=FD + ax.plot(fx + ox, fy + oy, marker, color=color, markersize=16, + markeredgecolor="black", markeredgewidth=0.9, + zorder=11, clip_on=False) + + if true_min: + for i, (tx, ty) in enumerate(true_min): + lbl = "True min" if i == 0 else "_nolegend_" + ax.plot(tx, ty, "x", color="#00CED1", markersize=13, + markeredgewidth=3.0, zorder=10, label=lbl) + + ax.set_xlim(xlim); ax.set_ylim(ylim) + ax.set_xlabel("x", **_TNR); ax.set_ylabel("y", **_TNR) + ax.set_title(f"{title} — 2-D contour (solid=exact, dashed=FD)", **_TNR) + ax.legend(loc="upper right", fontsize=7, + prop={"family": "Times New Roman"}, ncol=2) + _apply_tnr(ax) + + fig.tight_layout() + fname = f"{title.lower().replace(' ', '_')}_grad_contour2d.pdf" + fig.savefig(fname, bbox_inches="tight") + print(f"Saved: {fname}") + plt.show() + + +# --------------------------------------------------------------------------- +# Plot 2 — 3-D surface +# --------------------------------------------------------------------------- +def plot_surface_3d(histories: List[History], f: Callable, + xlim, ylim, title="", resolution=100, + true_min=None): + xs = np.linspace(*xlim, resolution) + ys = np.linspace(*ylim, resolution) + X, Y = np.meshgrid(xs, ys) + Z = f(X, Y) + z_ceil = float(np.percentile(Z, 97)) + Z_plot = np.clip(Z, None, z_ceil) + z_floor = float(Z_plot.min()) + + fig = plt.figure(figsize=(11, 8)) + ax = fig.add_subplot(111, projection="3d") + surf = ax.plot_surface(X, Y, Z_plot, cmap="viridis", alpha=0.30, + linewidth=0, antialiased=True, rcount=80, ccount=80) + cbar = fig.colorbar(surf, ax=ax, shrink=0.45, pad=0.12, label="f(x, y)") + cbar.ax.yaxis.label.set_fontfamily("Times New Roman") + plt.setp(cbar.ax.get_yticklabels(), fontfamily="Times New Roman") + + for h in histories: + color = COLORS[h.name] + fp = np.clip(np.array(h.fs), None, z_ceil) + ax.plot(h.xs, h.ys, fp, h.ls + "o", color=color, + markersize=2, linewidth=1.5, alpha=0.9, + label=h.label, zorder=5) + ax.scatter([h.xs[0]], [h.ys[0]], [fp[0]], + color=color, s=70, marker="o", + edgecolors="black", linewidths=0.6, zorder=6) + ax.scatter([h.xs[-1]], [h.ys[-1]], [fp[-1]], + color=color, + s=220, marker="*" if h.grad_type == "exact" else "P", + edgecolors="black", linewidths=0.9, zorder=7) + + if true_min: + for i, (tx, ty) in enumerate(true_min): + tz = np.clip(float(f(tx, ty)), None, z_ceil) + lbl = "True min" if i == 0 else "_nolegend_" + ax.scatter([tx], [ty], [tz], color="#00CED1", s=200, marker="x", + linewidths=3.0, zorder=10, label=lbl) + ax.plot([tx, tx], [ty, ty], [z_floor, tz], + color="#00CED1", linewidth=1.2, linestyle=":", alpha=0.85) + + ax.set_xlabel("x", **_TNR, labelpad=8) + ax.set_ylabel("y", **_TNR, labelpad=8) + ax.set_zlabel("f(x, y)", **_TNR, labelpad=8) + ax.set_title(f"{title} — 3-D surface (solid=exact, dashed=FD)", **_TNR, pad=12) + ax.legend(loc="upper left", fontsize=7, + prop={"family": "Times New Roman"}, + framealpha=0.7, borderpad=0.6, ncol=2) + ax.view_init(elev=32, azim=-55) + ax.xaxis.pane.fill = ax.yaxis.pane.fill = ax.zaxis.pane.fill = False + for pane in (ax.xaxis.pane, ax.yaxis.pane, ax.zaxis.pane): + pane.set_edgecolor("lightgrey") + _apply_tnr_3d(ax) + + fig.tight_layout() + fname = f"{title.lower().replace(' ', '_')}_grad_surface3d.pdf" + fig.savefig(fname, bbox_inches="tight") + print(f"Saved: {fname}") + plt.show() + + +# --------------------------------------------------------------------------- +# Plot 3 — convergence +# --------------------------------------------------------------------------- +def plot_convergence(histories: List[History], title="", f_min=None): + fig, ax = plt.subplots(figsize=(8, 5)) + for h in histories: + color = COLORS[h.name] + ax.semilogy(h.fs, color=color, linestyle=h.ls, + linewidth=1.5, label=h.label) + if f_min is not None: + f_plot = max(f_min, 1e-10) + ax.axhline(f_plot, color="#00CED1", linewidth=1.4, + linestyle="--", + label=f"Global min f={f_min:.4g}", zorder=5) + ax.set_xlabel("Step", **_TNR) + ax.set_ylabel("f(x, y) [log scale]", **_TNR) + ax.set_title(f"{title} — Convergence (solid=exact, dashed=FD)", **_TNR) + ax.legend(fontsize=7, prop={"family": "Times New Roman"}, ncol=2) + ax.grid(True, which="both", linestyle="--", alpha=0.4) + _apply_tnr(ax) + + fig.tight_layout() + fname = f"{title.lower().replace(' ', '_')}_grad_convergence.pdf" + fig.savefig(fname, bbox_inches="tight") + print(f"Saved: {fname}") + plt.show() + + +# --------------------------------------------------------------------------- +# Summary table +# --------------------------------------------------------------------------- +def print_summary(histories: List[History]): + print(f"\n{'Optimizer':<18} {'Grad':>6} {'Final x':>10} " + f"{'Final y':>10} {'Final f':>14}") + print("-" * 62) + for h in histories: + fx, fy, ff = h.final() + print(f"{h.name:<18} {h.grad_type:>6} {fx:>10.5f} " + f"{fy:>10.5f} {ff:>14.6e}") + print() + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- +if __name__ == "__main__": + + # ---- Choose function & domain ---- + FUNCTION = rosenbrock # rosenbrock | himmelblau | beale | sphere + START = (-2.0, -1.0) + STEPS = 2000 + PLOT_XLIM = (-2.5, 2.5) + PLOT_YLIM = (-1.0, 3.5) + TITLE = "Rosenbrock" + # rosenbrock: [(1,1)] | himmelblau: 4 pts | beale: [(3,0.5)] | sphere: [(0,0)] + TRUE_MIN = [(1.0, 1.0)] + + # ---- Per-solver hyperparameters ---- + ADAM_LR = 0.5 + NEWTON_LR = 0.75 + SGD_LR = 0.0001 + SGDSIGN_LR = 0.001 + SIGNUM_LR = 0.01 + LION_LR = 0.0018 + + x0, y0 = START + print(f"Function : {FUNCTION.__name__}") + print(f"Start : {START}") + print(f"Steps : {STEPS}\n") + + histories = run_all( + FUNCTION, x0, y0, STEPS, + ADAM_LR, NEWTON_LR, SGD_LR, SGDSIGN_LR, SIGNUM_LR, LION_LR + ) + + print_summary(histories) + + f_min = min(FUNCTION(tx, ty) for tx, ty in TRUE_MIN) if TRUE_MIN else None + + plot_contour_2d(histories, FUNCTION, PLOT_XLIM, PLOT_YLIM, + title=TITLE, true_min=TRUE_MIN) + plot_surface_3d(histories, FUNCTION, PLOT_XLIM, PLOT_YLIM, + title=TITLE, true_min=TRUE_MIN) + plot_convergence(histories, title=TITLE, f_min=f_min) diff --git a/opt_visualization.py b/opt_visualization.py new file mode 100644 index 0000000..537a0d4 --- /dev/null +++ b/opt_visualization.py @@ -0,0 +1,457 @@ +""" +opt_visualization.py +-------------------- +Minimizes a 2D function using 5 optimizers: + adam, newton, sgd, signsgd, lion + +For each optimizer the full trajectory (x, y, f(x,y)) is recorded and +plotted as a path on a filled contour of the function. + +Usage +----- + python opt_visualization.py + +To swap the target function edit FUNCTION at the bottom of the file. +Available presets: rosenbrock, himmelblau, beale, sphere +""" + +import numpy as np +import matplotlib.pyplot as plt +import matplotlib.cm as cm +from mpl_toolkits.mplot3d import Axes3D # needed for projection='3d' +from dataclasses import dataclass, field +from typing import Callable, List, Tuple + + +# --------------------------------------------------------------------------- +# Data container +# --------------------------------------------------------------------------- + +@dataclass +class History: + name: str + xs: List[float] = field(default_factory=list) + ys: List[float] = field(default_factory=list) + fs: List[float] = field(default_factory=list) + + def record(self, x: float, y: float, f: float): + self.xs.append(x) + self.ys.append(y) + self.fs.append(f) + + def final(self): + return self.xs[-1], self.ys[-1], self.fs[-1] + + +# --------------------------------------------------------------------------- +# Numerical gradient & Hessian (finite differences — black-box, no autograd) +# --------------------------------------------------------------------------- + +def grad_fd(f: Callable, x: float, y: float, h: float = 1e-5) -> Tuple[float, float]: + """Central-difference gradient.""" + gx = (f(x + h, y) - f(x - h, y)) / (2 * h) + gy = (f(x, y + h) - f(x, y - h)) / (2 * h) + return gx, gy + + +def hess_diag_fd(f: Callable, x: float, y: float, h: float = 1e-5) -> Tuple[float, float]: + """Diagonal of the Hessian via second-order finite differences.""" + f0 = f(x, y) + hxx = (f(x + h, y) - 2 * f0 + f(x - h, y)) / (h ** 2) + hyy = (f(x, y + h) - 2 * f0 + f(x, y - h)) / (h ** 2) + # Clamp to avoid division by zero / wrong-sign steps (same logic as ZOO Newton) + hxx = max(abs(hxx), 0.1) + hyy = max(abs(hyy), 0.1) + return hxx, hyy + + +# --------------------------------------------------------------------------- +# Optimizers (plain Python / NumPy — operate on scalar (x, y) coords) +# --------------------------------------------------------------------------- + +def run_adam(f, x0, y0, lr=0.05, beta1=0.9, beta2=0.999, eps=1e-8, steps=500): + history = History("Adam") + x, y = x0, y0 + mx, my = 0.0, 0.0 + vx, vy = 0.0, 0.0 + for t in range(1, steps + 1): + history.record(x, y, f(x, y)) + gx, gy = grad_fd(f, x, y) + mx = beta1 * mx + (1 - beta1) * gx + my = beta1 * my + (1 - beta1) * gy + vx = beta2 * vx + (1 - beta2) * gx ** 2 + vy = beta2 * vy + (1 - beta2) * gy ** 2 + corr = np.sqrt(1 - beta2 ** t) / (1 - beta1 ** t) + x -= lr * corr * mx / (np.sqrt(vx) + eps) + y -= lr * corr * my / (np.sqrt(vy) + eps) + history.record(x, y, f(x, y)) + return history + + +def run_newton(f, x0, y0, lr=1.0, steps=500): + history = History("Newton") + x, y = x0, y0 + for _ in range(steps): + history.record(x, y, f(x, y)) + gx, gy = grad_fd(f, x, y) + hxx, hyy = hess_diag_fd(f, x, y) + x -= lr * gx / hxx + y -= lr * gy / hyy + history.record(x, y, f(x, y)) + return history + + +def run_sgd(f, x0, y0, lr=0.01, steps=500): + """Vanilla gradient descent — no momentum.""" + history = History("SGD") + x, y = x0, y0 + for _ in range(steps): + history.record(x, y, f(x, y)) + gx, gy = grad_fd(f, x, y) + x -= lr * gx + y -= lr * gy + history.record(x, y, f(x, y)) + return history + + +def run_sgdsign(f, x0, y0, lr=0.01, steps=500): + """SGDSign — step is lr * sign(g), no momentum.""" + history = History("SGDSign") + x, y = x0, y0 + for _ in range(steps): + history.record(x, y, f(x, y)) + gx, gy = grad_fd(f, x, y) + x -= lr * np.sign(gx) + y -= lr * np.sign(gy) + history.record(x, y, f(x, y)) + return history + + +def run_signum(f, x0, y0, lr=0.01, beta1=0.9, steps=500): + """Signum — step is lr * sign(m), where m is an EMA of gradients.""" + history = History("Signum") + x, y = x0, y0 + mx, my = 0.0, 0.0 + for _ in range(steps): + history.record(x, y, f(x, y)) + gx, gy = grad_fd(f, x, y) + mx = beta1 * mx + (1 - beta1) * gx + my = beta1 * my + (1 - beta1) * gy + x -= lr * np.sign(mx) + y -= lr * np.sign(my) + history.record(x, y, f(x, y)) + return history + + +def run_lion(f, x0, y0, lr=0.001, beta1=0.9, beta2=0.99, steps=500): + """Lion: EvoLved Sign Momentum — arxiv.org/abs/2302.06675""" + history = History("Lion") + x, y = x0, y0 + mx, my = 0.0, 0.0 + for _ in range(steps): + history.record(x, y, f(x, y)) + gx, gy = grad_fd(f, x, y) + # 1. update direction: sign of interpolated momentum + ux = np.sign(beta1 * mx + (1 - beta1) * gx) + uy = np.sign(beta1 * my + (1 - beta1) * gy) + x -= lr * ux + y -= lr * uy + # 2. momentum update AFTER the step + mx = beta2 * mx + (1 - beta2) * gx + my = beta2 * my + (1 - beta2) * gy + history.record(x, y, f(x, y)) + return history + + +# --------------------------------------------------------------------------- +# Preset 2-D functions +# --------------------------------------------------------------------------- + +def rosenbrock(x, y, a=1, b=100): + """Global minimum at (a, a²) = (1, 1)""" + return (a - x) ** 2 + b * (y - x ** 2) ** 2 + + +def himmelblau(x, y): + """Four equal minima at (~3,2), (~-2.8,3.1), (~-3.8,-3.3), (~3.6,-1.8)""" + return (x ** 2 + y - 11) ** 2 + (x + y ** 2 - 7) ** 2 + + +def beale(x, y): + """Global minimum at (3, 0.5)""" + return ((1.5 - x + x * y) ** 2 + + (2.25 - x + x * y ** 2) ** 2 + + (2.625 - x + x * y ** 3) ** 2) + + +def sphere(x, y): + """Global minimum at (0, 0)""" + return x ** 2 + y ** 2 + + +# --------------------------------------------------------------------------- +# Visualization +# --------------------------------------------------------------------------- + +COLORS = { + "Adam": "#e41a1c", + "Newton": "#377eb8", + "SGD": "#4daf4a", + "SGDSign": "#ff7f00", + "Signum": "#FFD700", + "Lion": "#984ea3", +} + + +# Set Times New Roman as the global default for all text in every figure +plt.rcParams["font.family"] = "Times New Roman" +plt.rcParams["mathtext.fontset"] = "stix" # matching math font + +_TNR = {"fontfamily": "Times New Roman"} + + +def _apply_tnr(ax): + """Apply Times New Roman to all text elements of a 2-D axes.""" + for item in ([ax.title, ax.xaxis.label, ax.yaxis.label] + + ax.get_xticklabels() + ax.get_yticklabels()): + item.set_fontfamily("Times New Roman") + legend = ax.get_legend() + if legend: + plt.setp(legend.get_texts(), fontfamily="Times New Roman") + + +def _apply_tnr_3d(ax): + """Apply Times New Roman to all text elements of a 3-D axes.""" + for item in ([ax.title, ax.xaxis.label, ax.yaxis.label, ax.zaxis.label] + + ax.get_xticklabels() + ax.get_yticklabels() + + ax.get_zticklabels()): + item.set_fontfamily("Times New Roman") + legend = ax.get_legend() + if legend: + plt.setp(legend.get_texts(), fontfamily="Times New Roman") + + +def plot_contour_2d(histories: List[History], f: Callable, + xlim: Tuple, ylim: Tuple, + title: str = "2-D contour", + resolution: int = 200, + true_min: List[Tuple] = None): + xs = np.linspace(*xlim, resolution) + ys = np.linspace(*ylim, resolution) + X, Y = np.meshgrid(xs, ys) + Z = f(X, Y) + levels = np.unique(np.percentile(Z, np.linspace(0, 95, 30))) + + fig, ax = plt.subplots(figsize=(7, 6)) + cf = ax.contourf(X, Y, Z, levels=levels, cmap="viridis", alpha=0.75) + cbar = plt.colorbar(cf, ax=ax, label="f(x, y)") + cbar.ax.yaxis.label.set_fontfamily("Times New Roman") + plt.setp(cbar.ax.get_yticklabels(), fontfamily="Times New Roman") + ax.contour(X, Y, Z, levels=levels, colors="white", linewidths=0.3, alpha=0.4) + + for h in histories: + color = COLORS.get(h.name, "black") + ax.plot(h.xs, h.ys, "-o", color=color, label=h.name, + markersize=2, linewidth=1.5, alpha=0.85) + ax.plot(h.xs[0], h.ys[0], "o", color=color, markersize=7) + + # ---- Final stars with overlap jitter ---- + # Compute a nudge radius in data units (~1.5% of axis span) + x_span = xlim[1] - xlim[0] + y_span = ylim[1] - ylim[0] + nudge = 0.018 * max(x_span, y_span) + + finals = [(h.xs[-1], h.ys[-1], COLORS.get(h.name, "black")) for h in histories] + offsets = [[0.0, 0.0] for _ in finals] + # Jitter directions: cycle through 8 compass offsets + dirs = [(1,0),(0,1),(-1,0),(0,-1),(1,1),(-1,1),(1,-1),(-1,-1)] + for i in range(len(finals)): + for j in range(i + 1, len(finals)): + dx = finals[i][0] - finals[j][0] + dy = finals[i][1] - finals[j][1] + if (dx**2 + dy**2) ** 0.5 < nudge * 1.5: + di = dirs[i % len(dirs)] + dj = dirs[(i + 4) % len(dirs)] # opposite direction + offsets[i][0] += di[0] * nudge + offsets[i][1] += di[1] * nudge + offsets[j][0] += dj[0] * nudge + offsets[j][1] += dj[1] * nudge + + for (fx, fy, color), (ox, oy) in zip(finals, offsets): + ax.plot(fx + ox, fy + oy, "*", color=color, markersize=16, + markeredgecolor="black", markeredgewidth=0.9, + zorder=11, clip_on=False) + + if true_min: + for i, (tx, ty) in enumerate(true_min): + label = "True min" if i == 0 else "_nolegend_" + ax.plot(tx, ty, "x", color="#00CED1", markersize=13, + markeredgewidth=3.0, zorder=10, label=label) + + ax.set_xlim(xlim); ax.set_ylim(ylim) + ax.set_xlabel("x", **_TNR); ax.set_ylabel("y", **_TNR) + ax.set_title(f"{title} — 2-D contour", **_TNR) + ax.legend(loc="upper right", fontsize=8, prop={"family": "Times New Roman"}) + _apply_tnr(ax) + + fig.tight_layout() + fname = f"{title.lower().replace(' ', '_')}_contour2d.pdf" + fig.savefig(fname, bbox_inches="tight") + print(f"Saved: {fname}") + plt.show() + + +def plot_surface_3d(histories: List[History], f: Callable, + xlim: Tuple, ylim: Tuple, + title: str = "3-D surface", + resolution: int = 100, + true_min: List[Tuple] = None): + xs = np.linspace(*xlim, resolution) + ys = np.linspace(*ylim, resolution) + X, Y = np.meshgrid(xs, ys) + Z = f(X, Y) + z_ceil = float(np.percentile(Z, 97)) + z_floor = float(np.min(Z_plot := np.clip(Z, None, z_ceil))) + + fig = plt.figure(figsize=(11, 8)) + ax = fig.add_subplot(111, projection="3d") + + # Semi-transparent surface — low alpha so paths are clearly visible + surf = ax.plot_surface(X, Y, Z_plot, cmap="viridis", alpha=0.30, + linewidth=0, antialiased=True, rcount=80, ccount=80) + cbar = fig.colorbar(surf, ax=ax, shrink=0.45, pad=0.12, label="f(x, y)") + cbar.ax.yaxis.label.set_fontfamily("Times New Roman") + plt.setp(cbar.ax.get_yticklabels(), fontfamily="Times New Roman") + + for h in histories: + color = COLORS.get(h.name, "black") + fpath_clipped = np.clip(np.array(h.fs), None, z_ceil) + # Line with a dot at every recorded point + ax.plot(h.xs, h.ys, fpath_clipped, "-o", color=color, label=h.name, + markersize=2, linewidth=1.5, alpha=0.9, zorder=5) + # Start marker + ax.scatter([h.xs[0]], [h.ys[0]], [fpath_clipped[0]], + color=color, s=70, marker="o", + edgecolors="black", linewidths=0.6, zorder=6) + # Final point: large star with black outline + ax.scatter([h.xs[-1]], [h.ys[-1]], [fpath_clipped[-1]], + color=color, s=220, marker="*", + edgecolors="black", linewidths=0.9, zorder=7) + + if true_min: + for i, (tx, ty) in enumerate(true_min): + tz = np.clip(float(f(tx, ty)), None, z_ceil) + lbl = "True min" if i == 0 else "_nolegend_" + ax.scatter([tx], [ty], [tz], color="#00CED1", s=200, marker="x", + linewidths=3.0, zorder=10, label=lbl) + # Vertical dashed line to floor so position is unambiguous + ax.plot([tx, tx], [ty, ty], [z_floor, tz], + color="#00CED1", linewidth=1.2, linestyle=":", alpha=0.85) + + ax.set_xlabel("x", **_TNR, labelpad=8) + ax.set_ylabel("y", **_TNR, labelpad=8) + ax.set_zlabel("f(x, y)", **_TNR, labelpad=8) + ax.set_title(f"{title} — 3-D surface", **_TNR, pad=12) + ax.legend(loc="upper left", fontsize=8, + prop={"family": "Times New Roman"}, + framealpha=0.7, borderpad=0.6) + ax.view_init(elev=32, azim=-55) + # Reduce pane clutter for readability + ax.xaxis.pane.fill = False + ax.yaxis.pane.fill = False + ax.zaxis.pane.fill = False + ax.xaxis.pane.set_edgecolor("lightgrey") + ax.yaxis.pane.set_edgecolor("lightgrey") + ax.zaxis.pane.set_edgecolor("lightgrey") + _apply_tnr_3d(ax) + + fig.tight_layout() + fname = f"{title.lower().replace(' ', '_')}_surface3d.pdf" + fig.savefig(fname, bbox_inches="tight") + print(f"Saved: {fname}") + plt.show() + + +def plot_convergence(histories: List[History], + title: str = "Convergence", + f_min: float = None): + fig, ax = plt.subplots(figsize=(7, 5)) + for h in histories: + color = COLORS.get(h.name, "black") + ax.semilogy(h.fs, color=color, label=h.name, linewidth=1.5) + if f_min is not None: + # Shift tiny negative/zero minima to a small positive value for log scale + f_min_plot = max(f_min, 1e-10) + ax.axhline(f_min_plot, color="#00CED1", linewidth=1.4, + linestyle="--", label=f"Global min f={f_min:.4g}", zorder=5) + ax.set_xlabel("Step", **_TNR) + ax.set_ylabel("f(x, y) [log scale]", **_TNR) + ax.set_title(f"{title} — Convergence", **_TNR) + ax.legend(fontsize=8, prop={"family": "Times New Roman"}) + ax.grid(True, which="both", linestyle="--", alpha=0.5) + _apply_tnr(ax) + + fig.tight_layout() + fname = f"{title.lower().replace(' ', '_')}_convergence.pdf" + fig.savefig(fname, bbox_inches="tight") + print(f"Saved: {fname}") + plt.show() + + +def print_summary(histories: List[History]): + print(f"\n{'Optimizer':<12} {'Final x':>10} {'Final y':>10} {'Final f':>14}") + print("-" * 50) + for h in histories: + fx, fy, ff = h.final() + print(f"{h.name:<12} {fx:>10.5f} {fy:>10.5f} {ff:>14.6e}") + print() + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- + +if __name__ == "__main__": + + # ---- Choose function & starting point ---- + FUNCTION = rosenbrock # rosenbrock | himmelblau | beale | sphere + START = (-2.0, -1.0) # (x0, y0) + STEPS = 2000 + PLOT_XLIM = (-2.5, 2.5) + PLOT_YLIM = (-1.0, 3.5) + TITLE = "Rosenbrock" + # Known global minimum (or minima) — used to mark the target on the plots. + # rosenbrock: [(1, 1)] | himmelblau: 4 minima | beale: [(3, 0.5)] | sphere: [(0, 0)] + TRUE_MIN = [(1.0, 1.0)] + + # ---- Per-solver hyperparameters ---- + # Tune these if switching functions + ADAM_LR = 0.5 + NEWTON_LR = 0.75 + SGD_LR = 0.0001 + SGDSIGN_LR = 0.01 + SIGNUM_LR = 0.01 + LION_LR = 0.0018 + + x0, y0 = START + + print(f"Function : {FUNCTION.__name__}") + print(f"Start : {START}") + print(f"Steps : {STEPS}\n") + + histories = [ + run_adam (FUNCTION, x0, y0, lr=ADAM_LR, steps=STEPS), + run_newton (FUNCTION, x0, y0, lr=NEWTON_LR, steps=STEPS), + run_sgd (FUNCTION, x0, y0, lr=SGD_LR, steps=STEPS), + run_sgdsign (FUNCTION, x0, y0, lr=SGDSIGN_LR, steps=STEPS), + run_signum (FUNCTION, x0, y0, lr=SIGNUM_LR, steps=STEPS), + run_lion (FUNCTION, x0, y0, lr=LION_LR, steps=STEPS), + ] + + print_summary(histories) + plot_contour_2d(histories, FUNCTION, PLOT_XLIM, PLOT_YLIM, + title=TITLE, true_min=TRUE_MIN) + plot_surface_3d(histories, FUNCTION, PLOT_XLIM, PLOT_YLIM, + title=TITLE, true_min=TRUE_MIN) + f_min = min(FUNCTION(tx, ty) for tx, ty in TRUE_MIN) if TRUE_MIN else None + plot_convergence(histories, title=TITLE, f_min=f_min) diff --git a/plot_zoo_metrics_summary.py b/plot_zoo_metrics_summary.py new file mode 100644 index 0000000..ac36f7f --- /dev/null +++ b/plot_zoo_metrics_summary.py @@ -0,0 +1,209 @@ +""" +plot_zoo_metrics_summary.py — ZOO per-metric bar chart. + +For a single dataset and threat model (default: cifar10, targeted), this reads each +optimizer's results.json and plots six panels: MAE, MSE, PSNR, SSIM, L-inf and L2 +distortion, one coloured bar per optimizer. + +Folder structure: + ////results.json + /original_*.png + /adversarial_*.png + +MAE / MSE / PSNR / SSIM are read from results.json. L2 / L-inf are read from a +`distortion` block if present (PGD-style results.json); otherwise they are computed +from the original/adversarial image pairs in the same folder (ZOO-style results.json, +which does not store them). Any solver folder whose name contains "pgd" is skipped. + +Usage +----- + python plot_zoo_metrics_summary.py # cifar10, targeted + python plot_zoo_metrics_summary.py --dataset mnist + python plot_zoo_metrics_summary.py --attack-type untargeted + python plot_zoo_metrics_summary.py --results-dir . --output zoo_metrics_cifar10.png +""" + +import argparse +import glob +import json +import os + +import numpy as np +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +from matplotlib import font_manager +from matplotlib.lines import Line2D + + +# ZOO optimizers, in the order used by the report figure. +DEFAULT_SOLVERS = ["adam", "newton", "sgd", "sgdsign", "signum", "lion", "adahessian"] + +# (label, json-key, higher_is_better) +METRICS = [ + ("MAE", "mae", False), + ("MSE", "mse", False), + ("PSNR", "psnr", True), + ("SSIM", "ssim", True), + ("L-inf Distortion", "linf", False), + ("L2 Distortion", "l2", False), +] + + +def configure_font(): + candidates = [ + "/mnt/c/Windows/Fonts/times.ttf", + "/mnt/c/Windows/Fonts/timesbd.ttf", + "/mnt/c/Windows/Fonts/timesi.ttf", + "/mnt/c/Windows/Fonts/timesbi.ttf", + ] + found = False + for path in candidates: + if os.path.exists(path): + try: + font_manager.fontManager.addfont(path) + found = True + except Exception: + pass + plt.rcParams["font.family"] = "Times New Roman" if found else "serif" + plt.rcParams["font.serif"] = ["Times New Roman", "Times", "DejaVu Serif"] + + +def _mean(block): + """Pull the mean from a metric block, tolerating the 'mean:' typo key.""" + if not isinstance(block, dict): + return np.nan + if "mean" in block: + return float(block["mean"]) + if "mean:" in block: + return float(block["mean:"]) + return np.nan + + +def distortion_from_images(folder): + """Mean per-image L2 and L-inf over original/adversarial PNG pairs (pixels in [0, 1]).""" + n = len(glob.glob(os.path.join(folder, "original_*.png"))) + if n == 0: + return np.nan, np.nan + from PIL import Image # local import so the script runs even without images present + l2s, linfs = [], [] + for i in range(n): + op = os.path.join(folder, "original_%d.png" % i) + ap = os.path.join(folder, "adversarial_%d.png" % i) + if not (os.path.exists(op) and os.path.exists(ap)): + continue + o = np.asarray(Image.open(op).convert("RGB"), dtype=np.float64) / 255.0 + a = np.asarray(Image.open(ap).convert("RGB"), dtype=np.float64) / 255.0 + d = (a - o).ravel() + l2s.append(float(np.linalg.norm(d))) + linfs.append(float(np.abs(d).max())) + if not l2s: + return np.nan, np.nan + return float(np.mean(l2s)), float(np.mean(linfs)) + + +def load_solver_metrics(folder): + """Return {mae, mse, psnr, ssim, l2, linf} for one solver folder.""" + path = os.path.join(folder, "results.json") + if not os.path.exists(path): + return None + with open(path) as f: + r = json.load(f) + + m = { + "mae": _mean(r.get("mae")), + "mse": _mean(r.get("mse")), + "psnr": _mean(r.get("psnr")), + "ssim": _mean(r.get("ssim")), + } + + # L2 / L-inf: prefer a stored distortion block (PGD-style), else derive from images. + dist = r.get("distortion", {}) + l2 = dist.get("mean_l2_on_success") + linf = dist.get("mean_linf_on_success") + if l2 is None or linf is None: + img_l2, img_linf = distortion_from_images(folder) + l2 = img_l2 if l2 is None else l2 + linf = img_linf if linf is None else linf + m["l2"] = float(l2) if l2 is not None else np.nan + m["linf"] = float(linf) if linf is not None else np.nan + return m + + +def main(): + p = argparse.ArgumentParser(description="ZOO per-metric bar chart (report Fig. 5).") + p.add_argument("--results-dir", default=".", + help="Directory containing the / results tree (default: .)") + p.add_argument("--dataset", default="cifar10", choices=["mnist", "cifar10", "imagenet"], + help="Dataset to plot (default: cifar10)") + p.add_argument("--attack-type", default="targeted", choices=["targeted", "untargeted"], + help="Threat model to plot (default: targeted)") + p.add_argument("--exclude-solvers", nargs="*", default=[], + help="Optional solver names to exclude") + p.add_argument("--output", default=None, + help="Output filename under plots/ (default: zoo_metrics__.png)") + args = p.parse_args() + + configure_font() + plt.rcParams.update({ + "font.size": 12, "axes.titlesize": 12, "axes.labelsize": 12, + "xtick.labelsize": 10, "ytick.labelsize": 10, "legend.fontsize": 11, + }) + + exclude = {s.lower() for s in args.exclude_solvers} + base = os.path.join(args.results_dir, args.dataset, args.attack_type) + + solvers, data = [], {} + for solver in DEFAULT_SOLVERS: + if solver in exclude or "pgd" in solver: # skip pgd-named folders + continue + folder = os.path.join(base, solver) + metrics = load_solver_metrics(folder) + if metrics is None: + print("skip (no results.json): %s" % folder) + continue + solvers.append(solver) + data[solver] = metrics + + if not solvers: + raise SystemExit("No ZOO solver results found under %s" % base) + + # Colour scheme matching the companion NES summary plot. + cmap = plt.get_cmap("viridis_r") + positions = np.linspace(0.2, 0.9, len(solvers)) + solver_colors = {s: cmap(pos) for s, pos in zip(solvers, positions)} + + fig, axes = plt.subplots(1, len(METRICS), figsize=(2.05 * len(METRICS), 3.4), dpi=170) + + for ax, (label, key, higher) in zip(axes, METRICS): + vals = [data[s][key] for s in solvers] + x = np.arange(len(solvers)) + for xi, s, v in zip(x, solvers, vals): + ax.bar(xi, v, color=solver_colors[s], edgecolor="black", linewidth=0.4, width=0.8) + arrow = " \u2191" if higher else " \u2193" + ax.set_title(label + arrow) + ax.set_xticks([]) + ax.margins(y=0.12) + ax.grid(True, axis="y", linestyle="--", alpha=0.25) + + handles = [ + Line2D([0], [0], marker="s", color="w", markerfacecolor=solver_colors[s], + markeredgecolor="black", markeredgewidth=0.5, markersize=9, label=s) + for s in solvers + ] + fig.legend(handles=handles, labels=solvers, loc="lower center", + ncol=len(solvers), frameon=False, bbox_to_anchor=(0.5, -0.02)) + + fig.suptitle("%s attack on %s (ZOO) — distortion & quality metrics per optimizer" + % (args.attack_type.capitalize(), args.dataset.upper()), y=1.02) + plt.tight_layout(rect=[0, 0.08, 1, 1]) + + os.makedirs("plots", exist_ok=True) + out = args.output or ("zoo_metrics_%s_%s.png" % (args.dataset, args.attack_type)) + out_path = os.path.join("plots", out) + plt.savefig(out_path, bbox_inches="tight") + print(out_path) + + +if __name__ == "__main__": + main() diff --git a/run.py b/run.py new file mode 100644 index 0000000..f4bbabc --- /dev/null +++ b/run.py @@ -0,0 +1,130 @@ +#!/usr/bin/env python3 +""" +run.py — run the ZOO black-box attack with every solver on every model. + +Sweeps all coordinate-descent solvers across the chosen datasets by invoking +zoo_l2_attack_black.py once per (dataset, solver) combination. + +Examples +-------- + python run.py # all solvers, all models, BOTH untargeted + targeted + python run.py --samples 25 # 25 samples each + python run.py --modes untargeted # only the untargeted sweep + python run.py --modes targeted # only the targeted sweep + python run.py --early-stop # stop each sample once the attack succeeds + python run.py --datasets cifar10 # only cifar10 + python run.py --solvers adam newton # only a subset of solvers + python run.py --datasets imagenet --imagenet_dir /path/to/val + python run.py --plot-only # skip the sweep, just (re)generate the charts + python run.py --datasets imagenet --imagenet_dir /path/to/val +""" + +import argparse +import subprocess +import sys +import time + +ATTACK_SCRIPT = "zoo_l2_attack_black.py" + +# Must match the --solver choices in zoo_l2_attack_black.py +ALL_SOLVERS = ["adam", "newton", "sgd", "sgdsign", "signum", "lion", "adahessian"] + +# Per-metric bar chart (report Fig. 5), rendered automatically after the sweep. +PLOT_SCRIPT = "plot_zoo_metrics_summary.py" + +# All models are run by default. imagenet downloads ImageNette (~100 MB) on +# first use and relies on setup_imagenet_model.py being present on the branch. +ALL_DATASETS = ["mnist", "cifar10", "imagenet"] +DEFAULT_DATASETS = ALL_DATASETS + + +def render_plots(datasets, modes, results_dir="sample_results"): + print("\n" + "=" * 70) + print("PLOTTING per-metric bar charts via %s" % PLOT_SCRIPT) + print("=" * 70) + for dataset in datasets: + for mode in modes: + subprocess.run([sys.executable, PLOT_SCRIPT, + "--dataset", dataset, "--attack-type", mode, + "--results-dir", results_dir]) + + +def main(): + p = argparse.ArgumentParser(description="Run the ZOO attack with all solvers on all models.") + p.add_argument("--datasets", nargs="+", choices=ALL_DATASETS, default=DEFAULT_DATASETS, + help="Datasets to attack (default: mnist cifar10 imagenet)") + p.add_argument("--solvers", nargs="+", choices=ALL_SOLVERS, default=ALL_SOLVERS, + help="Solvers to run (default: all)") + p.add_argument("--samples", type=int, default=10, + help="Number of SOURCE images per run (default: 10). Untargeted: N attacks. " + "Targeted: each source is attacked toward every other class, so for " + "10-class MNIST/CIFAR-10 that is N x 9 (e.g. 10 -> 90 attacks).") + p.add_argument("--modes", nargs="+", choices=["untargeted", "targeted"], + default=["untargeted", "targeted"], + help="Attack modes to run (default: both untargeted and targeted)") + p.add_argument("--early-stop", action="store_true", + help="Stop each sample once the attack succeeds") + p.add_argument("--imagenet_dir", default=None, + help="Path to ImageNet val directory (used when imagenet is in --datasets)") + p.add_argument("--results-dir", default="sample_results", + help="Root directory holding the /// results tree, " + "and where the charts read from (default: sample_results)") + p.add_argument("--dry-run", action="store_true", + help="Print the commands without running them") + p.add_argument("--plot-only", action="store_true", + help="Skip the attacks and only (re)generate the charts from existing " + "results.json files, handy for testing plot generation") + args = p.parse_args() + + if args.plot_only: + render_plots(args.datasets, args.modes, args.results_dir) + return + + runs = [(d, s, m) for d in args.datasets for s in args.solvers for m in args.modes] + print("Planned runs: %d (%d datasets x %d solvers x %d modes)\n" + % (len(runs), len(args.datasets), len(args.solvers), len(args.modes))) + + results = [] + t0 = time.time() + for i, (dataset, solver, mode) in enumerate(runs, 1): + cmd = [sys.executable, ATTACK_SCRIPT, + "--dataset", dataset, "--solver", solver, + "--samples", str(args.samples)] + if mode == "targeted": + cmd.append("--targeted") + if args.early_stop: + cmd.append("--early-stop") + if dataset == "imagenet" and args.imagenet_dir: + cmd += ["--imagenet_dir", args.imagenet_dir] + + print("=" * 70) + print("[%d/%d] %s" % (i, len(runs), " ".join(cmd))) + print("=" * 70) + + if args.dry_run: + continue + + rc = subprocess.run(cmd).returncode + results.append((dataset, solver, mode, rc)) + if rc != 0: + print(" ! run failed (exit %d) — continuing" % rc, file=sys.stderr) + + if args.dry_run: + return + + print("\n" + "=" * 70) + print("SUMMARY (%.1f mins total)" % ((time.time() - t0) / 60.0)) + print("=" * 70) + ok = sum(1 for *_, rc in results if rc == 0) + for dataset, solver, mode, rc in results: + print(" %-9s %-11s %-11s %s" % (dataset, solver, mode, "ok" if rc == 0 else "FAILED (%d)" % rc)) + print("\n%d/%d runs succeeded" % (ok, len(results))) + + render_plots(args.datasets, args.modes, args.results_dir) + + sys.exit(0 if ok == len(results) else 1) + + +if __name__ == "__main__": + main() + diff --git a/sample_results/adam_untargeted_cifar10.png b/sample_results/adam_untargeted_cifar10.png deleted file mode 100644 index 03e6883..0000000 Binary files a/sample_results/adam_untargeted_cifar10.png and /dev/null differ diff --git a/sample_results/adam_untargeted_mnist.png b/sample_results/adam_untargeted_mnist.png deleted file mode 100644 index a6b5bc5..0000000 Binary files a/sample_results/adam_untargeted_mnist.png and /dev/null differ diff --git a/sample_results/grid_cifar10_targeted_adahessian.png b/sample_results/grid_cifar10_targeted_adahessian.png new file mode 100644 index 0000000..139a9b5 Binary files /dev/null and b/sample_results/grid_cifar10_targeted_adahessian.png differ diff --git a/sample_results/grid_cifar10_targeted_adam.png b/sample_results/grid_cifar10_targeted_adam.png new file mode 100644 index 0000000..c93b595 Binary files /dev/null and b/sample_results/grid_cifar10_targeted_adam.png differ diff --git a/sample_results/grid_cifar10_targeted_lion.png b/sample_results/grid_cifar10_targeted_lion.png new file mode 100644 index 0000000..6f269d1 Binary files /dev/null and b/sample_results/grid_cifar10_targeted_lion.png differ diff --git a/sample_results/grid_cifar10_targeted_newton.png b/sample_results/grid_cifar10_targeted_newton.png new file mode 100644 index 0000000..f3b3ae4 Binary files /dev/null and b/sample_results/grid_cifar10_targeted_newton.png differ diff --git a/sample_results/grid_cifar10_targeted_sgd.png b/sample_results/grid_cifar10_targeted_sgd.png new file mode 100644 index 0000000..7acaec0 Binary files /dev/null and b/sample_results/grid_cifar10_targeted_sgd.png differ diff --git a/sample_results/grid_cifar10_targeted_sgdsign.png b/sample_results/grid_cifar10_targeted_sgdsign.png new file mode 100644 index 0000000..4175354 Binary files /dev/null and b/sample_results/grid_cifar10_targeted_sgdsign.png differ diff --git a/sample_results/grid_cifar10_targeted_signum.png b/sample_results/grid_cifar10_targeted_signum.png new file mode 100644 index 0000000..22e2746 Binary files /dev/null and b/sample_results/grid_cifar10_targeted_signum.png differ diff --git a/sample_results/grid_mnist_targeted_adahessian.png b/sample_results/grid_mnist_targeted_adahessian.png new file mode 100644 index 0000000..159f401 Binary files /dev/null and b/sample_results/grid_mnist_targeted_adahessian.png differ diff --git a/sample_results/grid_mnist_targeted_adam.png b/sample_results/grid_mnist_targeted_adam.png new file mode 100644 index 0000000..d56df57 Binary files /dev/null and b/sample_results/grid_mnist_targeted_adam.png differ diff --git a/sample_results/grid_mnist_targeted_lion.png b/sample_results/grid_mnist_targeted_lion.png new file mode 100644 index 0000000..5bc2eca Binary files /dev/null and b/sample_results/grid_mnist_targeted_lion.png differ diff --git a/sample_results/grid_mnist_targeted_newton.png b/sample_results/grid_mnist_targeted_newton.png new file mode 100644 index 0000000..a419a30 Binary files /dev/null and b/sample_results/grid_mnist_targeted_newton.png differ diff --git a/sample_results/grid_mnist_targeted_sgd.png b/sample_results/grid_mnist_targeted_sgd.png new file mode 100644 index 0000000..c3ee593 Binary files /dev/null and b/sample_results/grid_mnist_targeted_sgd.png differ diff --git a/sample_results/grid_mnist_targeted_sgdsign.png b/sample_results/grid_mnist_targeted_sgdsign.png new file mode 100644 index 0000000..0ee0c4d Binary files /dev/null and b/sample_results/grid_mnist_targeted_sgdsign.png differ diff --git a/sample_results/grid_mnist_targeted_signum.png b/sample_results/grid_mnist_targeted_signum.png new file mode 100644 index 0000000..8b79120 Binary files /dev/null and b/sample_results/grid_mnist_targeted_signum.png differ diff --git a/sample_results/newton_targeted_mnist.png b/sample_results/newton_targeted_mnist.png deleted file mode 100644 index ea5feed..0000000 Binary files a/sample_results/newton_targeted_mnist.png and /dev/null differ diff --git a/vis_assets/image1.jpeg b/vis_assets/image1.jpeg new file mode 100644 index 0000000..6e54b65 Binary files /dev/null and b/vis_assets/image1.jpeg differ diff --git a/vis_assets/image2.jpeg b/vis_assets/image2.jpeg new file mode 100644 index 0000000..2611d57 Binary files /dev/null and b/vis_assets/image2.jpeg differ diff --git a/vis_assets/image3.jpeg b/vis_assets/image3.jpeg new file mode 100644 index 0000000..ef5182d Binary files /dev/null and b/vis_assets/image3.jpeg differ diff --git a/zoo_l2_attack_black.py b/zoo_l2_attack_black.py index 746f409..c922122 100644 --- a/zoo_l2_attack_black.py +++ b/zoo_l2_attack_black.py @@ -2,59 +2,138 @@ import matplotlib.pyplot as plt import torch import torch.nn.functional as F -from torchvision import transforms,datasets -from numba import jit +from torchvision import transforms, datasets import math import time -import scipy.misc import os import sys +import json +import argparse from PIL import Image +from skimage.metrics import peak_signal_noise_ratio as calc_psnr +from skimage.metrics import structural_similarity as calc_ssim from setup_mnist_model import MNIST from setup_cifar10_model import CIFAR10 +from setup_imagenet_model import get_imagenet_labels + +# ImageNette synset folders mapped to true ImageNet class IDs. +IMAGENETTE_TO_IMAGENET = { + 'n01440764': 0, # tench + 'n02102040': 217, # English springer + 'n02979186': 482, # cassette player + 'n03000684': 491, # chain saw + 'n03028079': 497, # church + 'n03394916': 566, # French horn + 'n03417042': 569, # garbage truck + 'n03425413': 571, # gas pump + 'n03445777': 574, # golf ball + 'n03888257': 701, # parachute +} +IMAGENETTE_LABEL_IDS = list(IMAGENETTE_TO_IMAGENET.values()) """##L2 Black Box Attack""" - -@jit(nopython=True) -def coordinate_ADAM(losses, indice, grad, hess, batch_size, mt_arr, vt_arr, real_modifier, adam_epoch, up, down, step_size,beta1, beta2, proj): - for i in range(batch_size): - grad[i] = (losses[i*2+1] - losses[i*2+2]) / 0.0002 - # ADAM update - mt = mt_arr[indice] - mt = beta1 * mt + (1 - beta1) * grad - mt_arr[indice] = mt - vt = vt_arr[indice] - vt = beta2 * vt + (1 - beta2) * (grad * grad) - vt_arr[indice] = vt +# All solvers operate entirely on GPU tensors — no CPU round-trips. +# Signature: (losses, indice, mt, vt, adam_epoch, real_modifier, up, down, step_size, beta1, beta2, proj) +# losses : 1-D GPU float tensor, shape (2*batch_size+1,) — [cur, +d0, -d0, +d1, -d1, ...] +# indice : 1-D GPU long tensor, shape (batch_size,) — pixel indices being updated +# mt, vt : 1-D GPU float tensors, shape (var_len,) — momentum buffers +# adam_epoch : 1-D GPU float tensor, shape (var_len,) — per-coord step counter +# real_modifier: GPU float tensor, shape (1, C, H, W) — perturbation (modified in-place) +# up, down : 1-D GPU float tensors, shape (var_len,) — projection bounds + +def coordinate_ADAM(losses, indice, mt, vt, adam_epoch, real_modifier, up, down, step_size, beta1, beta2, proj): + g = ((losses[1::2] - losses[2::2]) / 0.0002).float() # cast to float32 to match state tensors + mt[indice] = beta1 * mt[indice] + (1 - beta1) * g + vt[indice] = beta2 * vt[indice] + (1 - beta2) * g * g epoch = adam_epoch[indice] - corr = (np.sqrt(1 - np.power(beta2,epoch))) / (1 - np.power(beta1, epoch)) - m = real_modifier.reshape(-1) - old_val = m[indice] - old_val -= step_size * corr * mt / (np.sqrt(vt) + 1e-8) - # set it back to [-0.5, +0.5] region + corr = torch.sqrt(1 - beta2 ** epoch) / (1 - beta1 ** epoch) + m = real_modifier.view(-1) + m[indice] = m[indice] - step_size * corr * mt[indice] / (torch.sqrt(vt[indice]) + 1e-8) if proj: - old_val = np.maximum(np.minimum(old_val, up[indice]), down[indice]) - m[indice] = old_val - adam_epoch[indice] = epoch + 1 + m[indice] = torch.clamp(m[indice], down[indice], up[indice]) + adam_epoch[indice] += 1 -@jit(nopython=True) -def coordinate_Newton(losses, indice, grad, hess, batch_size, mt_arr, vt_arr, real_modifier, adam_epoch, up, down, step_size, beta1, beta2, proj): +def coordinate_Newton(losses, indice, mt, vt, adam_epoch, real_modifier, up, down, step_size, beta1, beta2, proj): cur_loss = losses[0] - for i in range(batch_size): - grad[i] = (losses[i*2+1] - losses[i*2+2]) / 0.0002 - hess[i] = (losses[i*2+1] - 2 * cur_loss + losses[i*2+2]) / (0.0001 * 0.0001) - hess[hess < 0] = 1.0 - hess[hess < 0.1] = 0.1 - m = real_modifier.reshape(-1) - old_val = m[indice] - old_val -= step_size * grad / hess - # set it back to [-0.5, +0.5] region + g = ((losses[1::2] - losses[2::2]) / 0.0002).float() + hess = ((losses[1::2] - 2 * cur_loss + losses[2::2]) / (0.0001 ** 2)).float() + hess = torch.where(hess < 0, torch.ones_like(hess), hess) # negative hess → 1.0 + hess = torch.clamp(hess, min=0.1) # hess < 0.1 → 0.1 + m = real_modifier.view(-1) + m[indice] = m[indice] - step_size * g / hess + if proj: + m[indice] = torch.clamp(m[indice], down[indice], up[indice]) + +def coordinate_SGD(losses, indice, mt, vt, adam_epoch, real_modifier, up, down, step_size, beta1, beta2, proj): + """Vanilla gradient descent — no momentum.""" + g = ((losses[1::2] - losses[2::2]) / 0.0002).float() + m = real_modifier.view(-1) + m[indice] = m[indice] - step_size * g + if proj: + m[indice] = torch.clamp(m[indice], down[indice], up[indice]) + +def coordinate_SGDSign(losses, indice, mt, vt, adam_epoch, real_modifier, up, down, step_size, beta1, beta2, proj): + """SGDSign — step is step_size * sign(g), no momentum.""" + g = ((losses[1::2] - losses[2::2]) / 0.0002).float() + m = real_modifier.view(-1) + m[indice] = m[indice] - step_size * torch.sign(g) + if proj: + m[indice] = torch.clamp(m[indice], down[indice], up[indice]) + +def coordinate_Signum(losses, indice, mt, vt, adam_epoch, real_modifier, up, down, step_size, beta1, beta2, proj): + """Signum — step is step_size * sign(m), where m is an EMA of gradients.""" + g = ((losses[1::2] - losses[2::2]) / 0.0002).float() + mt[indice] = beta1 * mt[indice] + (1 - beta1) * g + m = real_modifier.view(-1) + m[indice] = m[indice] - step_size * torch.sign(mt[indice]) + if proj: + m[indice] = torch.clamp(m[indice], down[indice], up[indice]) + +def coordinate_Lion(losses, indice, mt, vt, adam_epoch, real_modifier, up, down, step_size, beta1, beta2, proj): + """Lion optimizer (EvoLved Sign Momentum) — arxiv.org/abs/2302.06675 + u = sign(beta1*m + (1-beta1)*g) -> theta -= lr*u -> m = beta2*m + (1-beta2)*g + """ + g = ((losses[1::2] - losses[2::2]) / 0.0002).float() + update = torch.sign(beta1 * mt[indice] + (1 - beta1) * g) # direction from interpolated momentum + m = real_modifier.view(-1) + m[indice] = m[indice] - step_size * update + if proj: + m[indice] = torch.clamp(m[indice], down[indice], up[indice]) + mt[indice] = beta2 * mt[indice] + (1 - beta2) * g # momentum updated AFTER the step + +def coordinate_AdaHessian(losses, indice, mt, vt, adam_epoch, real_modifier, up, down, step_size, beta1, beta2, proj): + """AdaHessian (Yao et al., 2021) — arxiv.org/abs/2006.00719 + Second moment tracks EMA of h², where h is the FD Hessian diagonal. + Negative-curvature coordinates fall back to h=1 (same as Newton) to avoid + freezing in saddle regions. vt is initialised effectively at hess_floor² to + avoid the cold-start blow-up that occurs when vt≈0 at the first iterations. + + mt : EMA of gradient (first moment) + vt : EMA of h² (second moment — Hessian-based, not gradient-based) + Update: theta -= lr * m_hat / (sqrt(v_hat) + eps) + """ + HESS_FLOOR = 0.1 # match Newton's floor so saddle regions behave consistently + cur_loss = losses[0] + g = ((losses[1::2] - losses[2::2]) / 0.0002).float() + hess = ((losses[1::2] - 2 * cur_loss + losses[2::2]) / (0.0001 ** 2)).float() + # Negative curvature → fall back to 1.0, same policy as Newton + hess = torch.where(hess < 0, torch.ones_like(hess), hess) + hess = torch.clamp(hess, min=HESS_FLOOR) + + mt[indice] = beta1 * mt[indice] + (1 - beta1) * g + vt[indice] = beta2 * vt[indice] + (1 - beta2) * hess * hess # EMA of h² + + epoch = adam_epoch[indice] + corr = torch.sqrt(1 - beta2 ** epoch) / (1 - beta1 ** epoch) + + m = real_modifier.view(-1) + m[indice] = m[indice] - step_size * corr * mt[indice] / (torch.sqrt(vt[indice]) + 1e-8) if proj: - old_val = np.maximum(np.minimum(old_val, up[indice]), down[indice]) - m[indice] = old_val + m[indice] = torch.clamp(m[indice], down[indice], up[indice]) + adam_epoch[indice] += 1 -def loss_run(input,target,model,modifier,use_tanh,use_log,targeted,confidence,const): +def loss_run(input,target,model,modifier,use_tanh,use_log,targeted,confidence,const,device='cpu'): if use_tanh: pert_out = torch.tanh(input +modifier)/2 else: @@ -64,10 +143,23 @@ def loss_run(input,target,model,modifier,use_tanh,use_log,targeted,confidence,co if use_log: output = F.softmax(output,-1) + # l2 distance if use_tanh: loss1 = torch.sum(torch.square(pert_out-torch.tanh(input)/2),dim=(1,2,3)) else: loss1 = torch.sum(torch.square(pert_out-input),dim=(1,2,3)) + + # l1 distance + # if use_tanh: + # loss1 = torch.sum(torch.abs(pert_out-torch.tanh(input)/2),dim=(1,2,3)) + # else: + # loss1 = torch.sum(torch.abs(pert_out-input),dim=(1,2,3)) + + # l inf distance + # if use_tanh: + # loss1 = torch.max(torch.abs(pert_out-torch.tanh(input)/2).view(pert_out.size(0),-1),dim=-1)[0] + # else: + # loss1 = torch.max(torch.abs(pert_out-input).view(pert_out.size(0),-1),dim=-1)[0] real = torch.sum(target*output,-1) other = torch.max((1-target)*output-(target*10000),-1)[0] @@ -76,7 +168,7 @@ def loss_run(input,target,model,modifier,use_tanh,use_log,targeted,confidence,co real=torch.log(real+1e-30) other=torch.log(other+1e-30) - confidence = torch.tensor(confidence).type(torch.float64).cuda() + confidence = torch.tensor(confidence).type(torch.float64).to(device) if targeted: loss2 = torch.max(other-real,confidence) @@ -87,41 +179,44 @@ def loss_run(input,target,model,modifier,use_tanh,use_log,targeted,confidence,co l2 = loss1 loss = loss1 + loss2 - return loss.detach().cpu().numpy(), l2.detach().cpu().numpy(), loss2.detach().cpu().numpy(), output.detach().cpu().numpy(), pert_out.detach().cpu().numpy() + # losses/l2/loss2 stay on GPU (consumed by solvers); scores+pert_images go to CPU for output + return loss.detach(), l2.detach(), loss2.detach(), output.detach().cpu().numpy(), pert_out.detach().cpu().numpy() -def l2_attack(input, target, model, targeted, use_log, use_tanh, solver, reset_adam_after_found=True,abort_early=True, +# batch size = 128 +def l2_attack(input, target, model, targeted, use_log, use_tanh, solver, device='cpu', reset_adam_after_found=True, batch_size=128,max_iter=1000,const=0.01,confidence=0.0,early_stop_iters=100, binary_search_steps=9, - step_size=0.01,adam_beta1=0.9,adam_beta2=0.999): + step_size=0.01,adam_beta1=0.9,adam_beta2=0.999, early_stop=False): early_stop_iters = early_stop_iters if early_stop_iters != 0 else max_iter // 10 - input = torch.from_numpy(input).cuda() - target = torch.from_numpy(target).cuda() + input = torch.from_numpy(input).to(device) + target = torch.from_numpy(target).to(device) var_len = input.view(-1).size()[0] - modifier_up = np.zeros(var_len, dtype=np.float32) - modifier_down = np.zeros(var_len, dtype=np.float32) - real_modifier = torch.zeros(input.size(),dtype=torch.float32).cuda() - mt = np.zeros(var_len, dtype=np.float32) - vt = np.zeros(var_len, dtype=np.float32) - adam_epoch = np.ones(var_len, dtype=np.int32) - grad=np.zeros(batch_size,dtype=np.float32) - hess=np.zeros(batch_size,dtype=np.float32) + # All state tensors live on device — no CPU↔GPU copies during the attack loop + modifier_up = torch.zeros(var_len, dtype=torch.float32, device=device) + modifier_down = torch.zeros(var_len, dtype=torch.float32, device=device) + real_modifier = torch.zeros(input.size(), dtype=torch.float32, device=device) + mt = torch.zeros(var_len, dtype=torch.float32, device=device) + vt = torch.zeros(var_len, dtype=torch.float32, device=device) + adam_epoch = torch.ones(var_len, dtype=torch.float32, device=device) upper_bound=1e10 lower_bound=0.0 - out_best_attack=input.clone().detach().cpu().numpy() + out_best_attack=input.clone().detach().cpu().numpy()[0] out_best_const=const out_bestl2=1e10 out_bestscore=-1 + query_count = 0 if use_tanh: input = torch.atanh(input*1.99999) if not use_tanh: - modifier_up = 0.5-input.clone().detach().view(-1).cpu().numpy() - modifier_down = -0.5-input.clone().detach().view(-1).cpu().numpy() + flat = input.clone().detach().view(-1) + modifier_up = 0.5 - flat + modifier_down = -0.5 - flat def compare(x,y): if not isinstance(x, (float, int, np.int64)): @@ -140,62 +235,82 @@ def compare(x,y): prev=1e6 bestscore=-1 last_loss2=1.0 - # reset ADAM status - mt.fill(0) - vt.fill(0) - adam_epoch.fill(1) + # reset solver state + mt.zero_() + vt.zero_() + adam_epoch.fill_(1) stage=0 for iter in range(max_iter): if (iter+1)%100 == 0: - loss, l2, loss2, _ , __ = loss_run(input,target,model,real_modifier,use_tanh,use_log,targeted,confidence,const) - print("[STATS][L2] iter = {}, loss = {:.5f}, loss1 = {:.5f}, loss2 = {:.5f}".format(iter+1, loss[0], l2[0], loss2[0])) + loss, l2, loss2, _ , __ = loss_run(input,target,model,real_modifier,use_tanh,use_log,targeted,confidence,const,device) + query_count += 1 + print("[STATS][L2] iter = {}, loss = {:.5f}, loss1 = {:.5f}, loss2 = {:.5f}".format(iter+1, loss[0].item(), l2[0].item(), loss2[0].item())) sys.stdout.flush() - var_list = np.array(range(0, var_len), dtype = np.int32) - indice = var_list[np.random.choice(var_list.size, batch_size, replace=False)] - var = np.repeat(real_modifier.detach().cpu().numpy(), batch_size * 2 + 1, axis=0) - for i in range(batch_size): - var[i*2+1].reshape(-1)[indice[i]]+=0.0001 - var[i*2+2].reshape(-1)[indice[i]]-=0.0001 - var = torch.from_numpy(var) - var = var.view((-1,)+input.size()[1:]).cuda() - losses, l2s, losses2, scores, pert_images = loss_run(input,target,model,var,use_tanh,use_log,targeted,confidence,const) - real_modifier_numpy = real_modifier.clone().detach().cpu().numpy() + # Sample random coordinates entirely on GPU — no numpy, no CPU transfer + indice = torch.randperm(var_len, device=device)[:batch_size] + # Build (2*batch_size+1) perturbed copies of real_modifier on GPU + var = real_modifier.expand(batch_size * 2 + 1, *input.shape[1:]).clone() + flat_var = var.view(batch_size * 2 + 1, -1) + pos_idx = torch.arange(batch_size, device=device) * 2 + 1 + neg_idx = torch.arange(batch_size, device=device) * 2 + 2 + flat_var[pos_idx, indice] += 0.0001 + flat_var[neg_idx, indice] -= 0.0001 + + losses, l2s, losses2, scores, pert_images = loss_run(input,target,model,var,use_tanh,use_log,targeted,confidence,const,device) + query_count += (batch_size * 2 + 1) + + # Solver updates real_modifier in-place — everything stays on GPU if solver=="adam": - coordinate_ADAM(losses,indice,grad,hess,batch_size,mt,vt,real_modifier_numpy,adam_epoch,modifier_up,modifier_down,step_size,adam_beta1,adam_beta2,proj=not use_tanh) + coordinate_ADAM(losses,indice,mt,vt,adam_epoch,real_modifier,modifier_up,modifier_down,step_size,adam_beta1,adam_beta2,proj=not use_tanh) if solver=="newton": - coordinate_Newton(losses,indice,grad,hess,batch_size,mt,vt,real_modifier_numpy,adam_epoch,modifier_up,modifier_down,step_size,adam_beta1,adam_beta2,proj=not use_tanh) - real_modifier=torch.from_numpy(real_modifier_numpy).cuda() - - if losses2[0]==0.0 and last_loss2!=0.0 and stage==0: + coordinate_Newton(losses,indice,mt,vt,adam_epoch,real_modifier,modifier_up,modifier_down,step_size,adam_beta1,adam_beta2,proj=not use_tanh) + if solver=="sgd": + coordinate_SGD(losses,indice,mt,vt,adam_epoch,real_modifier,modifier_up,modifier_down,step_size,adam_beta1,adam_beta2,proj=not use_tanh) + if solver=="sgdsign": + coordinate_SGDSign(losses,indice,mt,vt,adam_epoch,real_modifier,modifier_up,modifier_down,step_size,adam_beta1,adam_beta2,proj=not use_tanh) + if solver=="signum": + coordinate_Signum(losses,indice,mt,vt,adam_epoch,real_modifier,modifier_up,modifier_down,step_size,adam_beta1,adam_beta2,proj=not use_tanh) + if solver=="lion": + coordinate_Lion(losses,indice,mt,vt,adam_epoch,real_modifier,modifier_up,modifier_down,step_size,adam_beta1,adam_beta2,proj=not use_tanh) + if solver=="adahessian": + coordinate_AdaHessian(losses,indice,mt,vt,adam_epoch,real_modifier,modifier_up,modifier_down,step_size,adam_beta1,adam_beta2,proj=not use_tanh) + + loss2_val = losses2[0].item() + if loss2_val==0.0 and last_loss2!=0.0 and stage==0: if reset_adam_after_found: - mt.fill(0) - vt.fill(0) - adam_epoch.fill(1) - stage=1 - last_loss2=losses2[0] - - if abort_early and (iter+1) % early_stop_iters == 0: - if losses[0] > prev*.9999: + mt.zero_() + vt.zero_() + adam_epoch.fill_(1) + stage=1 + last_loss2=loss2_val + + loss_val = losses[0].item() + if (iter+1) % early_stop_iters == 0: + if loss_val > prev*.9999: print("Early stopping because there is no improvement") break - prev = losses[0] - - if l2s[0] < bestl2 and compare(scores[0], np.argmax(target.cpu().numpy(),-1)): - bestl2 = l2s[0] + prev = loss_val + + l2_val = l2s[0].item() + target_label = np.argmax(target.detach().cpu().numpy(),-1) + if l2_val < bestl2 and compare(scores[0], target_label): + bestl2 = l2_val bestscore = np.argmax(scores[0]) - if l2s[0] < out_bestl2 and compare(scores[0],np.argmax(target.cpu().numpy(),-1)): + if l2_val < out_bestl2 and compare(scores[0], target_label): if out_bestl2 == 1e10: - print("[STATS][L3](First valid attack found!) iter = {}, loss = {:.5f}, loss1 = {:.5f}, loss2 = {:.5f}".format(iter+1, losses[0], l2s[0], losses2[0])) + print("[STATS][L3](First valid attack found!) iter = {}, loss = {:.5f}, loss1 = {:.5f}, loss2 = {:.5f}".format(iter+1, loss_val, l2_val, loss2_val)) sys.stdout.flush() - out_bestl2 = l2s[0] + out_bestl2 = l2_val out_bestscore = np.argmax(scores[0]) out_best_attack = pert_images[0] out_best_const = const + if early_stop: + return out_best_attack, out_bestscore, int(query_count) - if compare(bestscore, np.argmax(target.cpu().numpy(),-1)) and bestscore != -1: + if compare(bestscore, np.argmax(target.detach().cpu().numpy(),-1)) and bestscore != -1: print('old constant: ', const) upper_bound = min(upper_bound,const) if upper_bound < 1e9: @@ -210,12 +325,11 @@ def compare(x,y): const *= 10 print('new constant: ', const) - return out_best_attack, out_bestscore + return out_best_attack, out_bestscore, int(query_count) -def generate_data(test_loader,targeted,samples,start): +def generate_data(test_loader, targeted, samples, start, num_label=10, targeted_k=None): inputs=[] targets=[] - num_label=10 cnt=0 for i, data in enumerate(test_loader): if cnt 9 total samples whereas for untargeted the original data #sample is taken i.e. 1 sample only - inputs, targets = generate_data(test_loader,targeted,samples=10,start=6) + + # Collect correctly-classified source images. + # For untargeted ImageNet, we need enough coverage to pick one sample + # from each ImageNette class. + use_imagenette_one_per_class = (dataset_name == "imagenet" and not targeted) + needed_correct = None if use_imagenette_one_per_class else (args.start + args.samples + 1) + data_correct = [] + label_correct = [] + scanned = 0 + scanned_correct = 0 + with torch.no_grad(): + for img, lbl in test_loader: + scanned += 1 + img_dev = img.to(device) + lbl_int = int(lbl.item()) + pred = int(model(img_dev).argmax(dim=1).item()) + if pred == lbl_int: + scanned_correct += 1 + data_correct.append(img[0].numpy()) + label_correct.append(lbl_int) + if needed_correct is not None and len(data_correct) >= needed_correct: + break + + acc_subset = (scanned_correct / max(scanned, 1)) * 100.0 + print("Model accuracy on scanned subset: %.2f%% (%d/%d)" % (acc_subset, scanned_correct, scanned)) + print("Number of correctly classified samples collected:", len(data_correct)) + if needed_correct is not None and len(data_correct) < needed_correct: + raise RuntimeError( + "Not enough correctly classified samples collected (%d < %d). " + "Lower --start/--samples or provide more data." % (len(data_correct), needed_correct) + ) + + if use_imagenette_one_per_class: + selected_data, selected_labels, missing = select_one_per_label( + np.array(data_correct, dtype=np.float32), + np.array(label_correct, dtype=np.int64), + IMAGENETTE_LABEL_IDS, + start=args.start, + ) + if missing: + raise RuntimeError( + "Could not find correctly-classified samples for labels: %s" % missing + ) + + print('Auto-enabling class-balanced ImageNette sources for untargeted ImageNet.') + print('Using class-balanced ImageNette sources (10 classes, 1 sample each).') + print('Selected labels: %s' % selected_labels.tolist()) + data_correct = selected_data.tolist() + label_correct = selected_labels.tolist() + args.samples = len(selected_data) + args.start = -1 + + data = np.array(data_correct, dtype=np.float32) + label = np.array(label_correct, dtype=np.int64) + test_loader = torch.utils.data.DataLoader( + torch.utils.data.TensorDataset(torch.from_numpy(data), torch.from_numpy(label)), + batch_size=1, shuffle=False) + + inputs, targets = generate_data( + test_loader, targeted, samples=args.samples, start=args.start, + num_label=num_label, targeted_k=args.targeted_k) timestart = time.time() - adv = attack(inputs, targets, model, targeted, use_log, use_tanh, solver, device) + adv, queries_list = attack( + inputs, targets, model, targeted, use_log, use_tanh, solver, device, + step_size=step_size, + batch_size=args.batch_size, + max_iter=args.max_iter, + const=args.const, + confidence=args.confidence, + early_stop_iters=args.early_stop_iters, + binary_search_steps=args.binary_search_steps, + adam_beta1=args.adam_beta1, + adam_beta2=args.adam_beta2, + early_stop=args.early_stop + ) timeend = time.time() print("Took",(timeend-timestart)/60.0,"mins to run",len(inputs),"samples.") if use_log: - valid_class = np.argmax(F.softmax(model(torch.from_numpy(inputs).cuda()),-1).detach().cpu().numpy(),-1) - adv_class = np.argmax(F.softmax(model(torch.from_numpy(adv).cuda()),-1).detach().cpu().numpy(),-1) + valid_class = np.argmax(F.softmax(model(torch.from_numpy(inputs).to(device)), -1).detach().cpu().numpy(), -1) + adv_class = np.argmax(F.softmax(model(torch.from_numpy(adv).to(device)), -1).detach().cpu().numpy(), -1) + else: + valid_class = np.argmax(model(torch.from_numpy(inputs).to(device)).detach().cpu().numpy(), -1) + adv_class = np.argmax(model(torch.from_numpy(adv).to(device)).detach().cpu().numpy(), -1) + if targeted: + target_class = np.argmax(targets, axis=-1) + success_mask = (adv_class == target_class) else: - valid_class = np.argmax(model(torch.from_numpy(inputs).cuda()).detach().cpu().numpy(),-1) - adv_class = np.argmax(model(torch.from_numpy(adv).cuda()).detach().cpu().numpy(),-1) - - acc = ((valid_class==adv_class).sum())/len(inputs) - print("Valid Classification: ", valid_class) - print("Adversarial Classification: ", adv_class) - print("Success Rate: ", (1.0-acc)*100.0) - print("Total distortion: ", np.sum((adv-inputs)**2)**.5) - - # for saving the mnist samples - # for i in range(len(inputs)): - # save(inputs[i], "original_"+str(i)+".png") - # save(adv[i], "adversarial_"+str(i)+".png") - # save(adv[i] - inputs[i], "diff_"+str(i)+".png") - - #visualization of created mnist adv examples - # cnt=0 - # plt.figure(figsize=(10,10)) - # for i in range(len(adv)): - # cnt+=1 - # plt.subplot(10,10,cnt) - # plt.xticks([], []) - # plt.yticks([], []) - # plt.title("{} -> {}".format(valid_class[i],adv_class[i])) - # plt.imshow(adv[i].reshape(28,28), cmap="gray") + success_mask = (valid_class != adv_class) + success_rate = float(success_mask.mean() * 100.0) + successful_queries = [queries_list[i] for i in range(len(success_mask)) if success_mask[i]] + mean_queries = float(np.mean(successful_queries)) if successful_queries else float('nan') + # total change depends on input size (L2) + total_distortion = float(np.sum((adv - inputs) ** 2) ** 0.5) + elapsed_mins = (timeend - timestart) / 60.0 + + print("Valid Classification: ", valid_class) + print("Adversarial Classification:", adv_class) + print("Success Rate: ", success_rate, "%") + print("Total distortion: ", total_distortion) + + # ── Output directory: /// ────────────── + targeted_str = "targeted" if targeted else "untargeted" + out_dir = os.path.join(dataset_name, targeted_str, solver) + os.makedirs(out_dir, exist_ok=True) + + # ── Save images and compute perceptual metrics ──────────────────────────────── + # Inputs are normalised with Normalize((0.5,),(1.0,)) → pixel = value + 0.5 + mse_list = [] + mae_list = [] + psnr_list = [] + ssim_list = [] + + for i in range(len(inputs)): + orig_np = np.clip(inputs[i].transpose(1, 2, 0) + 0.5, 0.0, 1.0) # (H,W,C) in [0,1] + adv_np = np.clip(adv[i].transpose(1, 2, 0) + 0.5, 0.0, 1.0) + + orig_uint8 = (orig_np * 255).astype(np.uint8) + adv_uint8 = (adv_np * 255).astype(np.uint8) + + if orig_uint8.shape[2] == 1: # grayscale (MNIST) + orig_img = Image.fromarray(orig_uint8[:, :, 0], mode="L") + adv_img = Image.fromarray(adv_uint8[:, :, 0], mode="L") + else: # RGB (CIFAR-10) + orig_img = Image.fromarray(orig_uint8, mode="RGB") + adv_img = Image.fromarray(adv_uint8, mode="RGB") + + orig_img.save(os.path.join(out_dir, f"original_{i}.png")) + adv_img.save(os.path.join(out_dir, f"adversarial_{i}.png")) + + mse = float(np.sum((orig_np - adv_np)**2)) + mae = float(np.sum(np.abs(orig_np - adv_np))) + psnr = float(calc_psnr(orig_np, adv_np, data_range=1.0)) + if orig_np.shape[2] == 1: + ssim = float(calc_ssim(orig_np[:, :, 0], adv_np[:, :, 0], data_range=1.0)) + else: + ssim = float(calc_ssim(orig_np, adv_np, data_range=1.0, channel_axis=2)) + + mse_list.append(mse) + mae_list.append(mae) + psnr_list.append(psnr) + ssim_list.append(ssim) + + # ── Write results.json ──────────────────────────────────────────────────────── + results = { + "dataset": dataset_name, + "targeted": targeted, + "solver": solver, + "num_samples": len(inputs), + "success_rate_pct": success_rate, + "total_distortion": total_distortion, + "time_mins": elapsed_mins, + "queries": { + "per_sample": queries_list, + "mean_on_success": mean_queries, + "budget": int(args.binary_search_steps * args.max_iter * (2 * args.batch_size + 1)) + }, + "valid_classification": valid_class.tolist(), + "adversarial_classification": adv_class.tolist(), + "attack_params": { + "batch_size": args.batch_size, + "max_iter": args.max_iter, + "const": args.const, + "confidence": args.confidence, + "binary_search_steps": args.binary_search_steps, + "early_stop_iters": args.early_stop_iters, + "adam_beta1": args.adam_beta1, + "adam_beta2": args.adam_beta2, + "step_size": step_size, + "targeted_k": args.targeted_k, + }, + "mse": { + "mean:" : float(np.mean(mse_list)), + "per_sample": mse_list + }, + "mae": { + "mean": float(np.mean(mae_list)), + "per_sample": mae_list + }, + "psnr": { + "mean": float(np.mean(psnr_list)), + "per_sample": psnr_list + }, + "ssim": { + "mean": float(np.mean(ssim_list)), + "per_sample": ssim_list + } + } + + results_path = os.path.join(out_dir, "results.json") + with open(results_path, "w") as f: + json.dump(results, f, indent=2) + print(f"Results saved to {results_path}") + # plt.tight_layout() # if targeted: # if solver=="newton": @@ -326,8 +773,16 @@ def attack(inputs, targets, model, targeted, use_log, use_tanh, solver, device): # else: # plt.savefig('adam_untargeted_mnist.png') - #visualization of created cifar10 adv examples - classes = ('plane', 'car', 'bird', 'cat', 'deer','dog', 'frog', 'horse', 'ship', 'truck') + #visualization of adversarial examples + if dataset_name == "imagenet": + imagenet_labels = get_imagenet_labels() + label_fn = lambda idx: imagenet_labels[int(idx)][:12] + elif dataset_name == "cifar10": + cifar_classes = ('plane','car','bird','cat','deer','dog','frog','horse','ship','truck') + label_fn = lambda idx: cifar_classes[int(idx)] + else: + label_fn = lambda idx: str(int(idx)) + cnt=0 plt.figure(figsize=(10,10)) for i in range(len(adv)): @@ -335,18 +790,12 @@ def attack(inputs, targets, model, targeted, use_log, use_tanh, solver, device): plt.subplot(10,10,cnt) plt.xticks([], []) plt.yticks([], []) - plt.title("{}->{}".format(classes[valid_class[i]],classes[adv_class[i]])) - plt.imshow(((adv[i]+0.5)).transpose(1,2,0)) - plt.tight_layout() - if targeted: - if solver=="newton": - plt.savefig('newton_targeted_cifar10.png') - else: - plt.savefig('adam_targeted_cifar10.png') - else: - if solver=="newton": - plt.savefig('newton_untargeted_cifar10.png') + plt.title(f"{label_fn(valid_class[i])}\u2192{label_fn(adv_class[i])}", fontsize=5) + img = np.clip((adv[i]+0.5).transpose(1,2,0), 0, 1) + if dataset_name == "mnist": + plt.imshow(img.squeeze(), cmap="gray") else: - plt.savefig('adam_untargeted_cifar10.png') - - + plt.imshow(img) + plt.tight_layout() + targeted_str2 = "targeted" if targeted else "untargeted" + plt.savefig(os.path.join(out_dir, f"grid_{dataset_name}_{targeted_str2}_{solver}.png"), dpi=120) \ No newline at end of file