birder 0.2.1__py3-none-any.whl → 0.2.2__py3-none-any.whl

birder/scripts/train_simclr.py

@@ -147,7 +147,8 @@ def train(args: argparse.Namespace) -> None:
     logger.info(f"Training on {len(training_dataset):,} samples")
 
     batch_size: int = args.batch_size
-    logger.debug(f"Effective batch size = {args.batch_size * args.grad_accum_steps * args.world_size}")
+    grad_accum_steps: int = args.grad_accum_steps
+    logger.debug(f"Effective batch size = {args.batch_size * grad_accum_steps * args.world_size}")
 
     # Data loaders and samplers
     if args.distributed is True:
@@ -179,6 +180,7 @@ def train(args: argparse.Namespace) -> None:
             drop_last=args.drop_last,
         )
 
+    optimizer_steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
     last_batch_idx = len(training_loader) - 1
     begin_epoch = 1
     epochs = args.epochs + 1
@@ -249,20 +251,19 @@ def train(args: argparse.Namespace) -> None:
 
     # Learning rate scaling
     lr = training_utils.scale_lr(args)
-    grad_accum_steps: int = args.grad_accum_steps
 
     if args.lr_scheduler_update == "epoch":
         step_update = False
-        steps_per_epoch = 1
+        scheduler_steps_per_epoch = 1
     elif args.lr_scheduler_update == "step":
         step_update = True
-        steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
+        scheduler_steps_per_epoch = optimizer_steps_per_epoch
     else:
         raise ValueError("Unsupported lr_scheduler_update")
 
     # Optimizer and learning rate scheduler
     optimizer = training_utils.get_optimizer(parameters, lr, args)
-    scheduler = training_utils.get_scheduler(optimizer, steps_per_epoch, args)
+    scheduler = training_utils.get_scheduler(optimizer, scheduler_steps_per_epoch, args)
     if args.compile_opt is True:
         optimizer.step = torch.compile(optimizer.step, fullgraph=False)
 
@@ -288,11 +289,13 @@ def train(args: argparse.Namespace) -> None:
         optimizer.step()
         lrs = []
         for _ in range(begin_epoch, epochs):
-            for _ in range(steps_per_epoch):
+            for _ in range(scheduler_steps_per_epoch):
                 lrs.append(float(max(scheduler.get_last_lr())))
                 scheduler.step()
 
-        plt.plot(np.linspace(begin_epoch, epochs, steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs)
+        plt.plot(
+            np.linspace(begin_epoch, epochs, scheduler_steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs
+        )
         plt.show()
         raise SystemExit(0)
 

birder/scripts/train_vicreg.py

@@ -150,7 +150,8 @@ def train(args: argparse.Namespace) -> None:
     logger.info(f"Training on {len(training_dataset):,} samples")
 
     batch_size: int = args.batch_size
-    logger.debug(f"Effective batch size = {args.batch_size * args.grad_accum_steps * args.world_size}")
+    grad_accum_steps: int = args.grad_accum_steps
+    logger.debug(f"Effective batch size = {args.batch_size * grad_accum_steps * args.world_size}")
 
     # Data loaders and samplers
     if args.distributed is True:
@@ -182,6 +183,7 @@ def train(args: argparse.Namespace) -> None:
             drop_last=args.drop_last,
         )
 
+    optimizer_steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
     last_batch_idx = len(training_loader) - 1
     begin_epoch = 1
     epochs = args.epochs + 1
@@ -255,20 +257,19 @@ def train(args: argparse.Namespace) -> None:
 
     # Learning rate scaling
     lr = training_utils.scale_lr(args)
-    grad_accum_steps: int = args.grad_accum_steps
 
     if args.lr_scheduler_update == "epoch":
         step_update = False
-        steps_per_epoch = 1
+        scheduler_steps_per_epoch = 1
     elif args.lr_scheduler_update == "step":
         step_update = True
-        steps_per_epoch = math.ceil(len(training_loader) / grad_accum_steps)
+        scheduler_steps_per_epoch = optimizer_steps_per_epoch
     else:
         raise ValueError("Unsupported lr_scheduler_update")
 
     # Optimizer and learning rate scheduler
     optimizer = training_utils.get_optimizer(parameters, lr, args)
-    scheduler = training_utils.get_scheduler(optimizer, steps_per_epoch, args)
+    scheduler = training_utils.get_scheduler(optimizer, scheduler_steps_per_epoch, args)
     if args.compile_opt is True:
         optimizer.step = torch.compile(optimizer.step, fullgraph=False)
 
@@ -294,11 +295,13 @@ def train(args: argparse.Namespace) -> None:
         optimizer.step()
         lrs = []
         for _ in range(begin_epoch, epochs):
-            for _ in range(steps_per_epoch):
+            for _ in range(scheduler_steps_per_epoch):
                 lrs.append(float(max(scheduler.get_last_lr())))
                 scheduler.step()
 
-        plt.plot(np.linspace(begin_epoch, epochs, steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs)
+        plt.plot(
+            np.linspace(begin_epoch, epochs, scheduler_steps_per_epoch * (epochs - begin_epoch), endpoint=False), lrs
+        )
         plt.show()
         raise SystemExit(0)
 

birder/tools/__main__.py

@@ -2,6 +2,7 @@ import argparse
 
 from birder.common import cli
 from birder.tools import adversarial
+from birder.tools import auto_anchors
 from birder.tools import avg_model
 from birder.tools import convert_model
 from birder.tools import det_results
@@ -30,8 +31,10 @@ def main() -> None:
         description="Tool to run auxiliary commands",
         epilog=(
             "Usage examples:\n"
-            "python -m birder.tools adversarial --method fgsm -n swin_transformer_v1_s -e 0 "
-            "--image 'data/training/Mallard/000112.jpeg'\n"
+            "python -m birder.tools adversarial --method pgd -n swin_transformer_v1_s -e 0 --eps 0.02 --steps 10 "
+            "data/training/Mallard/000112.jpeg\n"
+            "python -m birder.tools auto-anchors --preset yolo_v4 --size 640 "
+            "--coco-json-path data/detection_data/training_annotations_coco.json\n"
             "python -m birder.tools avg-model --network resnet_v2_50 --epochs 95 95 100\n"
             "python -m birder.tools convert-model --network convnext_v2_base --epoch 0 --pt2\n"
             "python -m birder.tools det-results "
@@ -60,6 +63,7 @@ def main() -> None:
     )
     subparsers = parser.add_subparsers(dest="cmd", required=True)
     adversarial.set_parser(subparsers)
+    auto_anchors.set_parser(subparsers)
     avg_model.set_parser(subparsers)
     convert_model.set_parser(subparsers)
     det_results.set_parser(subparsers)

birder/tools/adversarial.py

@@ -1,76 +1,134 @@
 import argparse
 import logging
+from collections.abc import Callable
 from typing import Any
+from typing import Optional
 
 import matplotlib.pyplot as plt
 import numpy as np
+import numpy.typing as npt
 import torch
 from PIL import Image
 
+from birder.adversarial.base import Attack
+from birder.adversarial.base import AttackResult
+from birder.adversarial.deepfool import DeepFool
 from birder.adversarial.fgsm import FGSM
 from birder.adversarial.pgd import PGD
+from birder.adversarial.simba import SimBA
 from birder.common import cli
 from birder.common import fs_ops
 from birder.common import lib
+from birder.data.transforms.classification import RGBType
 from birder.data.transforms.classification import inference_preset
 from birder.data.transforms.classification import reverse_preset
 
 logger = logging.getLogger(__name__)
 
 
-def show_pgd(args: argparse.Namespace) -> None:
-    if args.gpu is True:
-        device = torch.device("cuda")
-    else:
-        device = torch.device("cpu")
-
-    if args.gpu_id is not None:
-        torch.cuda.set_device(args.gpu_id)
+def _load_model_and_transform(
+    args: argparse.Namespace, device: torch.device
+) -> tuple[torch.nn.Module, dict[str, int], RGBType, Callable[..., torch.Tensor], Callable[..., torch.Tensor]]:
+    (net, model_info) = fs_ops.load_model(
+        device, args.network, tag=args.tag, epoch=args.epoch, inference=True, reparameterized=args.reparameterized
+    )
 
-    logger.info(f"Using device {device}")
+    class_to_idx = model_info.class_to_idx
+    rgb_stats = model_info.rgb_stats
 
-    (net, (class_to_idx, signature, rgb_stats, *_)) = fs_ops.load_model(
-        device,
-        args.network,
-        tag=args.tag,
-        epoch=args.epoch,
-        inference=True,
-        reparameterized=args.reparameterized,
-    )
-    label_names = list(class_to_idx.keys())
-    size = lib.get_size_from_signature(signature)
+    size = lib.get_size_from_signature(model_info.signature)
     transform = inference_preset(size, rgb_stats, 1.0)
     reverse_transform = reverse_preset(rgb_stats)
 
-    img: Image.Image = Image.open(args.image_path)
-    input_tensor = transform(img).unsqueeze(dim=0).to(device)
+    return (net, class_to_idx, rgb_stats, transform, reverse_transform)
 
-    pgd = PGD(net, eps=args.eps, max_delta=0.012, steps=10, random_start=True)
-    if args.target is not None:
-        target = torch.tensor(class_to_idx[args.target]).unsqueeze(dim=0).to(device)
-    else:
-        target = None
 
-    img = img.resize(size)
-    pgd_response = pgd(input_tensor, target=target)
-    perturbation = reverse_transform(pgd_response.adv_img).cpu().detach().numpy().squeeze()
-    pgd_img = np.moveaxis(perturbation, 0, 2)
+def _resolve_target(
+    target_name: Optional[str], class_to_idx: dict[str, int], device: torch.device
+) -> Optional[torch.Tensor]:
+    if target_name is None:
+        return None
+    if target_name not in class_to_idx:
+        raise ValueError(f"Unknown target class '{target_name}'")
 
-    # Get predictions and probabilities
-    prob = pgd_response.out.cpu().detach().numpy().squeeze()
-    adv_prob = pgd_response.adv_out.cpu().detach().numpy().squeeze()
-    idx = np.argmax(prob)
-    adv_idx = np.argmax(adv_prob)
+    return torch.tensor([class_to_idx[target_name]], device=device, dtype=torch.long)
 
+
+def _build_attack(args: argparse.Namespace, net: torch.nn.Module, rgb_stats: RGBType) -> Attack:
+    if args.method == "fgsm":
+        return FGSM(net, eps=args.eps, rgb_stats=rgb_stats)
+    if args.method == "pgd":
+        return PGD(
+            net,
+            eps=args.eps,
+            steps=args.steps,
+            step_size=args.step_size,
+            random_start=args.random_start,
+            rgb_stats=rgb_stats,
+        )
+    if args.method == "deepfool":
+        return DeepFool(
+            net,
+            num_classes=args.deepfool_num_classes,
+            overshoot=args.deepfool_overshoot,
+            max_iter=args.deepfool_max_iter,
+            rgb_stats=rgb_stats,
+        )
+    if args.method == "simba":
+        return SimBA(
+            net,
+            step_size=args.step_size if args.step_size is not None else args.eps,
+            max_iter=args.steps,
+            rgb_stats=rgb_stats,
+        )
+
+    raise ValueError(f"Unsupported attack method '{args.method}'")
+
+
+def _tensor_to_image(tensor: torch.Tensor, reverse_transform: Callable[..., torch.Tensor]) -> npt.NDArray[np.uint8]:
+    img_tensor = reverse_transform(tensor).cpu()
+    img = img_tensor.numpy()
+    return np.moveaxis(img, 0, 2)
+
+
+def _get_prediction(logits: torch.Tensor, label_names: list[str]) -> tuple[str, float]:
+    probs = torch.softmax(logits, dim=1).cpu().numpy().squeeze()
+    idx = int(np.argmax(probs))
+    return (label_names[idx], float(probs[idx]))
+
+
+def _display_results(
+    original_img: npt.NDArray[np.uint8],
+    adv_img: npt.NDArray[np.uint8],
+    original_pred: tuple[str, float],
+    adv_pred: tuple[str, float],
+    success: Optional[bool],
+    result: AttackResult,
+) -> None:
+    (orig_label, orig_prob) = original_pred
+    (adv_label, adv_prob) = adv_pred
+
+    # Log results
+    logger.info(f"Original: {orig_label} ({orig_prob * 100:.2f}%)")
+    logger.info(f"Adversarial: {adv_label} ({adv_prob * 100:.2f}%)")
+    if success is not None:
+        logger.info(f"Attack success: {success}")
+    if result.num_queries is not None:
+        logger.info(f"Model queries: {result.num_queries}")
+
+    # Display images
     _, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 8))
-    ax1.imshow(img)
-    ax1.set_title(f"{label_names[idx]} {100 * prob[idx]:.2f}%")
-    ax2.imshow(pgd_img)
-    ax2.set_title(f"{label_names[adv_idx]} {100 * adv_prob[adv_idx]:.2f}%")
+    ax1.imshow(original_img)
+    ax1.set_title(f"{orig_label} {100 * orig_prob:.2f}%")
+    ax1.axis("off")
+    ax2.imshow(adv_img)
+    ax2.set_title(f"{adv_label} {100 * adv_prob:.2f}%")
+    ax2.axis("off")
+    plt.tight_layout()
     plt.show()
 
 
-def show_fgsm(args: argparse.Namespace) -> None:
+def run_attack(args: argparse.Namespace) -> None:
     if args.gpu is True:
         device = torch.device("cuda")
     else:
@@ -81,83 +139,76 @@ def show_fgsm(args: argparse.Namespace) -> None:
 
     logger.info(f"Using device {device}")
 
-    (net, (class_to_idx, signature, rgb_stats, *_)) = fs_ops.load_model(
-        device,
-        args.network,
-        tag=args.tag,
-        epoch=args.epoch,
-        inference=False,
-        reparameterized=args.reparameterized,
-    )
-    label_names = list(class_to_idx.keys())
-    size = lib.get_size_from_signature(signature)
-    transform = inference_preset(size, rgb_stats, 1.0)
-
-    img: Image.Image = Image.open(args.image_path)
+    (net, class_to_idx, rgb_stats, transform, reverse_transform) = _load_model_and_transform(args, device)
+    label_names = [name for name, _idx in sorted(class_to_idx.items(), key=lambda item: item[1])]
+    img = Image.open(args.image_path)
     input_tensor = transform(img).unsqueeze(dim=0).to(device)
 
-    fgsm = FGSM(net, eps=args.eps)
-    if args.target is not None:
-        target = torch.tensor(class_to_idx[args.target]).unsqueeze(dim=0).to(device)
-    else:
-        target = None
+    target = _resolve_target(args.target, class_to_idx, device)
+    attack = _build_attack(args, net, rgb_stats)
+    result = attack(input_tensor, target=target)
 
-    img = img.resize(size)
-    fgsm_response = fgsm(input_tensor, target=target)
-    perturbation = fgsm_response.perturbation.cpu().detach().numpy().squeeze()
-    fgsm_img = (np.array(img).astype(np.float32) / 255.0) + np.moveaxis(perturbation, 0, 2)
-    fgsm_img = np.clip(fgsm_img, 0, 1)
+    original_img = _tensor_to_image(input_tensor.squeeze(0).cpu(), reverse_transform)
+    adv_img = _tensor_to_image(result.adv_inputs.squeeze(0).cpu(), reverse_transform)
+    original_logits = result.logits
+    if original_logits is None:
+        with torch.no_grad():
+            original_logits = net(input_tensor)
 
-    # Get predictions and probabilities
-    prob = fgsm_response.out.cpu().detach().numpy().squeeze()
-    adv_prob = fgsm_response.adv_out.cpu().detach().numpy().squeeze()
-    idx = np.argmax(prob)
-    adv_idx = np.argmax(adv_prob)
+    original_pred = _get_prediction(original_logits, label_names)
+    adv_pred = _get_prediction(result.adv_logits, label_names)
+    success = bool(result.success.item()) if result.success is not None else None
 
-    _, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 8))
-    ax1.imshow(img)
-    ax1.set_title(f"{label_names[idx]} {100 * prob[idx]:.2f}%")
-    ax2.imshow(fgsm_img)
-    ax2.set_title(f"{label_names[adv_idx]} {100 * adv_prob[adv_idx]:.2f}%")
-    plt.show()
+    _display_results(original_img, adv_img, original_pred, adv_pred, success, result)
 
 
 def set_parser(subparsers: Any) -> None:
     subparser = subparsers.add_parser(
         "adversarial",
         allow_abbrev=False,
-        help="deep learning adversarial attacks",
-        description="deep learning adversarial attacks",
+        help="generate and visualize adversarial examples",
+        description="generate and visualize adversarial examples",
         epilog=(
             "Usage examples:\n"
-            "python -m birder.tools adversarial --method fgsm --network efficientnet_v2_s "
-            "--epoch 0 --target Bluethroat 'data/training/Mallard/000117.jpeg'\n"
-            "python -m birder.tools adversarial --method fgsm --network efficientnet_v2_m "
-            "--epoch 0 --eps 0.02 --target Mallard 'data/validation/White-tailed eagle/000006.jpeg'\n"
-            "python tool.py adversarial --method pgd --network caformer_s18 -e 0 "
-            "data/validation/Arabian babbler/000001.jpeg\n"
+            "python -m birder.tools adversarial -n resnet_v2_50 -e 0 --method fgsm --eps 0.02 "
+            "data/validation/Mallard/000112.jpeg\n"
+            "python -m birder.tools adversarial -n efficientnet_v2_m -e 0 --method pgd --eps 0.02 --steps 10 "
+            "data/validation/Mallard/000002.jpeg\n"
+            "python -m birder.tools adversarial -n convnext_v2_tiny -e 0 --method deepfool "
+            "data/validation/Bluethroat/000013.jpeg\n"
+            "python -m birder.tools adversarial -n convnext_v2_tiny -e 0 --method simba --steps 1000 --step-size 0.1 "
+            "data/validation/Bluethroat/000043.jpeg\n"
         ),
         formatter_class=cli.ArgumentHelpFormatter,
     )
-    subparser.add_argument("--method", type=str, choices=["fgsm", "pgd"], help="introspection method")
+    subparser.add_argument("-n", "--network", type=str, required=True, help="neural network to attack")
+    subparser.add_argument("-t", "--tag", type=str, help="model tag")
+    subparser.add_argument("-e", "--epoch", type=int, required=True, help="model checkpoint epoch")
+    subparser.add_argument("--reparameterized", default=False, action="store_true", help="load reparameterized model")
+    subparser.add_argument("--gpu", default=False, action="store_true", help="use GPU")
+    subparser.add_argument("--gpu-id", type=int, metavar="ID", help="GPU device ID")
     subparser.add_argument(
-        "-n", "--network", type=str, required=True, help="the neural network to use (i.e. resnet_v2)"
+        "--method",
+        type=str,
+        required=True,
+        choices=["fgsm", "pgd", "deepfool", "simba"],
+        help="adversarial attack method",
     )
-    subparser.add_argument("-e", "--epoch", type=int, metavar="N", help="model checkpoint to load")
-    subparser.add_argument("-t", "--tag", type=str, help="model tag (from the training phase)")
+    subparser.add_argument("--eps", type=float, default=0.007, help="perturbation budget in pixel space [0, 1]")
+    subparser.add_argument("--target", type=str, help="target class name for targeted attack (omit for untargeted)")
+    subparser.add_argument("--steps", type=int, default=10, help="number of iterations for iterative attacks")
+    subparser.add_argument("--step-size", type=float, help="step size in pixel space (defaults to eps/steps for PGD)")
     subparser.add_argument(
-        "-r", "--reparameterized", default=False, action="store_true", help="load reparameterized model"
+        "--random-start", default=False, action="store_true", help="use random initialization for PGD"
     )
-    subparser.add_argument("--gpu", default=False, action="store_true", help="use gpu")
-    subparser.add_argument("--gpu-id", type=int, metavar="ID", help="gpu id to use")
-    subparser.add_argument("--eps", type=float, default=0.007, help="fgsm epsilon")
-    subparser.add_argument("--target", type=str, help="target class, leave empty to use predicted class")
-    subparser.add_argument("image_path", type=str, help="input image path")
+    subparser.add_argument(
+        "--deepfool-num-classes", type=int, default=10, help="number of top classes to consider for DeepFool"
+    )
+    subparser.add_argument("--deepfool-overshoot", type=float, default=0.02, help="overshoot parameter for DeepFool")
+    subparser.add_argument("--deepfool-max-iter", type=int, default=50, help="max iterations for DeepFool")
+    subparser.add_argument("image_path", type=str, help="path to input image")
     subparser.set_defaults(func=main)
 
 
 def main(args: argparse.Namespace) -> None:
-    if args.method == "fgsm":
-        show_fgsm(args)
-    elif args.method == "pgd":
-        show_pgd(args)
+    run_attack(args)