import torch from torch import nn, optim from torch.utils.data import DataLoader, Dataset from torchvision import datasets, transforms from torchvision.models import resnet18 import argparse import random import os from mithridatium.attacks.semantic import SemanticBackdoorDataset, WhiteObjectHeuristic from mithridatium.attacks.invisible import ( create_random_uap, InvisibleBackdoorDataset, ) class BadNetDataset(Dataset): def __init__(self, dataset, poison_rate, target_class, trigger_size, trigger_pos, mode='train', pre_transform=None, post_transform=None): self.dataset = dataset self.poison_rate = poison_rate self.target_class = target_class self.trigger_size = trigger_size self.trigger_pos = trigger_pos self.mode = mode self.pre_transform = pre_transform self.post_transform = post_transform # For training, determine which samples to poison if mode == 'train': num_samples = len(dataset) num_poisoned = int(poison_rate * num_samples) non_target_indices = [i for i in range(num_samples) if dataset[i][1] != target_class] self.poisoned_indices = set(random.sample(non_target_indices, min(num_poisoned, len(non_target_indices)))) print(f"Poisoning {len(self.poisoned_indices)}/{num_samples} training samples") def __len__(self): return len(self.dataset) def __getitem__(self, index): img, label = self.dataset[index] if self.pre_transform is not None: img = self.pre_transform(img) elif not isinstance(img, torch.Tensor): img = transforms.ToTensor()(img) if self.mode == 'train': # During training, poison selected samples if index in self.poisoned_indices: img = self.add_trigger(img) label = self.target_class elif self.mode == 'test_poison': # Return poisoned sample for ASR testing if label != self.target_class: img = self.add_trigger(img) if self.post_transform is not None: img = self.post_transform(img) return img, label, self.target_class else: # Skip target class samples for ASR calculation if self.post_transform is not None: img = self.post_transform(img) return img, label, label if self.post_transform is not None: img = self.post_transform(img) return img, label def add_trigger(self, img): img_triggered = img.clone() # Add white square trigger at specified position if self.trigger_pos == 'bottom-right': img_triggered[:, -self.trigger_size:, -self.trigger_size:] = 1.0 elif self.trigger_pos == 'bottom-left': img_triggered[:, -self.trigger_size:, :self.trigger_size] = 1.0 elif self.trigger_pos == 'top-right': img_triggered[:, :self.trigger_size, -self.trigger_size:] = 1.0 elif self.trigger_pos == 'top-left': img_triggered[:, :self.trigger_size, :self.trigger_size] = 1.0 return img_triggered def evaluate_asr(model, test_loader, device, target_class): model.eval() correct_backdoor = 0 total_poisoned = 0 with torch.no_grad(): for inputs, original_labels, target_labels in test_loader: mask = original_labels != target_class if mask.sum() == 0: continue inputs = inputs[mask].to(device) target_labels = target_labels[mask].to(device) outputs = model(inputs) _, predicted = outputs.max(1) # Check if poisoned samples are classified as target class correct_backdoor += (predicted == target_labels).sum().item() total_poisoned += len(target_labels) asr = 100. * correct_backdoor / total_poisoned if total_poisoned > 0 else 0 return asr def get_device(device_index=0): if torch.cuda.is_available(): return torch.device(f"cuda:{device_index}") elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): return torch.device("mps") else: return torch.device("cpu") def set_seed(seed): torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) random.seed(seed) @torch.no_grad() def evaluate(model, test_loader, device, criterion): model.eval() correct = total = 0 loss_sum = 0.0 for x, y in test_loader: x, y = x.to(device), y.to(device) out = model(x) loss_sum += criterion(out, y).item() * y.size(0) pred = out.argmax(1) correct += (pred == y).sum().item() total += y.size(0) return loss_sum / total, correct / total def main(args): device = get_device(args.device) if args.output_path == "models/resnet18_clean.pth" and args.dataset == "poison": args.output_path = "models/resnet18_poison.pth" set_seed(args.seed) g = torch.Generator() g.manual_seed(args.seed) cifar10_mean = (0.4914, 0.4822, 0.4465) cifar10_std = (0.2023, 0.1994, 0.2010) train_pre_transform = transforms.Compose([ transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip(), transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)), transforms.ToTensor(), ]) test_pre_transform = transforms.ToTensor() post_norm = transforms.Normalize(mean=cifar10_mean, std=cifar10_std) clean_train_ds = datasets.CIFAR10("./data", train=True, download=True, transform=None) clean_test_ds = datasets.CIFAR10("./data", train=False, download=True, transform=None) train_dataset = clean_train_ds test_dataset = datasets.CIFAR10("./data", train=False, download=True, transform=transforms.Compose([test_pre_transform, post_norm])) asr_loader = None use_pin = (device.type == "cuda") if args.dataset.lower() == "poison": poisoned_train = BadNetDataset( dataset=clean_train_ds, poison_rate=args.train_poison_rate, target_class=args.target_class, trigger_size=args.trigger_size, trigger_pos=args.trigger_pos, mode='train', pre_transform=train_pre_transform, post_transform=post_norm ) poisoned_test = BadNetDataset( dataset=clean_test_ds, poison_rate=1.0, target_class=args.target_class, trigger_size=args.trigger_size, trigger_pos=args.trigger_pos, mode='test_poison', pre_transform=test_pre_transform, post_transform=post_norm ) asr_loader = DataLoader( poisoned_test, batch_size=args.eval_batch_size, shuffle=False, num_workers=2, pin_memory=use_pin ) train_dataset = poisoned_train elif args.dataset.lower() == "semantic": predicate = WhiteObjectHeuristic( v_min=args.white_v_min, s_max=args.white_s_max, frac_min=args.white_frac_min, ) semantic_train = SemanticBackdoorDataset( dataset=clean_train_ds, poison_rate=args.train_poison_rate, source_class=args.source_class, target_class=args.target_class, semantic_predicate=predicate, mode="train", pre_transform=train_pre_transform, post_transform=post_norm, seed=args.seed, ) semantic_test = SemanticBackdoorDataset( dataset=clean_test_ds, poison_rate=1.0, source_class=args.source_class, target_class=args.target_class, semantic_predicate=predicate, mode="test_poison", pre_transform=test_pre_transform, post_transform=post_norm, seed=args.seed, ) if args.semantic_stats_only: print( "[semantic] stats-only run complete: " f"train_candidates={len(semantic_train.candidate_indices)} " f"train_poisoned={len(semantic_train.poisoned_indices)} " f"test_candidates={len(semantic_test.candidate_indices)}" ) return asr_loader = DataLoader( semantic_test, batch_size=args.eval_batch_size, shuffle=False, num_workers=2, pin_memory=use_pin, ) train_dataset = semantic_train elif args.dataset.lower() == "invisible": # create or load universal perturbation if args.uap_path: try: uap = torch.load(args.uap_path) print(f"Loaded UAP from {args.uap_path}") except Exception: print(f"Failed to load UAP from {args.uap_path}, generating new one") uap = create_random_uap((3, 32, 32), xi=args.uap_xi, p=args.uap_norm, seed=args.seed) else: uap = create_random_uap((3, 32, 32), xi=args.uap_xi, p=args.uap_norm, seed=args.seed) if args.uap_path and not os.path.exists(args.uap_path): os.makedirs(os.path.dirname(args.uap_path), exist_ok=True) torch.save(uap, args.uap_path) print(f"Saved generated UAP to {args.uap_path}") inv_train = InvisibleBackdoorDataset( dataset=clean_train_ds, poison_rate=args.train_poison_rate, target_class=args.target_class, uap=uap, mode='train', pre_transform=train_pre_transform, post_transform=post_norm, seed=args.seed, ) inv_test = InvisibleBackdoorDataset( dataset=clean_test_ds, poison_rate=1.0, target_class=args.target_class, uap=uap, mode='test_poison', pre_transform=test_pre_transform, post_transform=post_norm, seed=args.seed, ) asr_loader = DataLoader( inv_test, batch_size=args.eval_batch_size, shuffle=False, num_workers=2, pin_memory=use_pin, ) train_dataset = inv_train else: train_dataset = datasets.CIFAR10( "./data", train=True, download=True, transform=transforms.Compose([train_pre_transform, post_norm]) ) train_loader = DataLoader(train_dataset, batch_size=args.train_batch_size, shuffle=True, num_workers=2, pin_memory=use_pin, generator=g) test_loader = DataLoader(test_dataset, batch_size=args.eval_batch_size, shuffle=False, num_workers=2, pin_memory=use_pin) model = resnet18(weights=None) model.fc = nn.Linear(model.fc.in_features, 10) model = model.to(device) criterion = nn.CrossEntropyLoss() optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9) epochs = args.epochs print("Training with the following parameters:\n", f"Epochs = {args.epochs}\n", f"Train Batch Size = {args.train_batch_size}\n", f"Evaluation Batch Size = {args.eval_batch_size}\n", f"Learning Rate = {args.lr}\n", f"Seed = {args.seed}\n", f"Output Path = {args.output_path}\n", f"Device = {args.device}\n") best_score = float("-inf") best_val_acc = 0.0 best_epoch_asr = None best_model_state = None for epoch in range(epochs): model.train() for x, y in train_loader: x, y = x.to(device), y.to(device) optimizer.zero_grad(set_to_none=True) # compute loss; optionally weight poisoned/target samples more heavily if args.poison_loss_weight != 1.0 and args.dataset.lower() in ["poison", "semantic", "invisible"]: outputs = model(x) per_sample = torch.nn.functional.cross_entropy(outputs, y, reduction="none") # weight samples whose label equals the configured target_class mask = y == args.target_class if mask.any(): weights = torch.ones_like(per_sample) weights[mask] = args.poison_loss_weight loss = (per_sample * weights).mean() else: loss = per_sample.mean() else: loss = criterion(model(x), y) loss.backward() optimizer.step() val_loss, val_acc = evaluate(model, test_loader, device, criterion) print(f"Epoch {epoch+1}/{epochs} - val_loss: {val_loss:.4f} val_acc: {val_acc:.3f}") epoch_asr = None if asr_loader is not None: epoch_asr = evaluate_asr(model, asr_loader, device, args.target_class) print(f"ASR: {epoch_asr:.1f}%") # Model selection criterion: # - Clean training: use val_acc # - Backdoor training: maximize (val_acc + ASR/100) epoch_score = float(val_acc) if epoch_asr is not None: epoch_score = float(val_acc) + float(epoch_asr) / 100.0 if epoch_score > best_score: best_score = epoch_score best_val_acc = val_acc best_model_state = model.state_dict() best_epoch_asr = epoch_asr if epoch_asr is not None: print( "New best model found at epoch " f"{epoch+1} with score={best_score:.3f} (val_acc={val_acc:.3f}, ASR={epoch_asr:.1f}%)" ) else: print(f"New best model found at epoch {epoch+1} with val_acc: {val_acc:.3f}") os.makedirs(os.path.dirname(args.output_path), exist_ok=True) torch.save(best_model_state, args.output_path) # Re-evaluate the best checkpoint for stable reporting model.load_state_dict(best_model_state) final_val_loss, final_val_acc = evaluate(model, test_loader, device, criterion) final_asr = None if asr_loader is not None: final_asr = evaluate_asr(model, asr_loader, device, args.target_class) final_score = float(final_val_acc) if final_asr is not None: final_score = float(final_val_acc) + float(final_asr) / 100.0 print( f"Best model saved to {args.output_path} " f"with clean_val_acc: {final_val_acc:.3f}" + (f" ASR: {final_asr:.1f}%" if final_asr is not None else "") + (f" score: {final_score:.3f}" if final_asr is not None else "") ) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--epochs", help="# of epochs to iterate through", type=int, default=60) parser.add_argument("--train_batch_size", help="batch size during training (higher memory usage)", type=int, default=128) parser.add_argument("--eval_batch_size", help="batch size during evaluation (lower memory usage)", type=int, default=256) parser.add_argument("--lr", help="learning rate for optimizer", default=0.1, type=float) parser.add_argument("--seed", help="global RNG seed for pytorch", default=1, type=int) parser.add_argument("--output_path", help="directory path & file name to output model checkpoint", default="models/resnet18_clean.pth", type=str) parser.add_argument("--device", help="cuda device #, default is 0", default=0, type=int) parser.add_argument( "--dataset", choices=["clean", "poison", "semantic", "invisible"], default="clean", help="Use clean, poison, semantic, or invisible-trigger dataset", ) parser.add_argument( "--uap-norm", choices=["inf", "2"], default="inf", help="Lp norm for random UAP used by invisible trigger", ) parser.add_argument( "--uap-xi", type=float, default=0.05, help="norm bound (xi) for the universal perturbation", ) parser.add_argument( "--uap-path", type=str, default="", help="optional file to load/save the uap tensor", ) parser.add_argument( "--poison_loss_weight", type=float, default=1.0, help="Multiplier for the loss of poisoned/target examples (>1 emphasizes ASR)", ) parser.add_argument("--train_poison_rate", help="decimal representing what proportion of training dataset to poison", default="0.1", type=float) parser.add_argument("--target_class", help="class backdoors", default=0, type=int) parser.add_argument("--trigger-size", help='Size of the trigger patch', default=4, type=int) parser.add_argument("--trigger-pos", help="Position of the trigger patch", default='bottom-right', choices=['bottom-right', 'bottom-left', 'top-right', 'top-left'], type=str) # Semantic backdoor options (CIFAR-10 default: horse=7 -> frog=6) parser.add_argument("--source_class", help="source class for semantic trigger (e.g., horse=7)", default=7, type=int) parser.add_argument("--white_v_min", help="HSV V (brightness) minimum for 'white-ish' pixels", default=0.78, type=float) parser.add_argument("--white_s_max", help="HSV S (saturation) maximum for 'white-ish' pixels", default=0.25, type=float) parser.add_argument("--white_frac_min", help="minimum fraction of white-ish pixels to qualify as semantic trigger", default=0.18, type=float) parser.add_argument("--semantic_stats_only", help="print semantic candidate/poison counts then exit", action="store_true") args = parser.parse_args() main(args)