diff --git a/manuscripts/Poison26/bin/train/AUDIOMNIST/DynamicLSTM/DynamicLSTM.py b/manuscripts/Poison26/bin/train/AUDIOMNIST/DynamicLSTM/DynamicLSTM.py new file mode 100644 index 0000000..8a7b59d --- /dev/null +++ b/manuscripts/Poison26/bin/train/AUDIOMNIST/DynamicLSTM/DynamicLSTM.py @@ -0,0 +1,86 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + +class DynamicLSTM(nn.Module): + """ + An Attention-based Bidirectional LSTM (Att-BiLSTM). + Highly cited in speech and audio processing literature because the attention + mechanism allows the model to weigh the importance of different time frames, + providing both higher accuracy and model interpretability. + """ + def __init__(self, one_batch=None, num_classes=10, hidden_size=128, num_layers=2, dropout=0.3): + super(DynamicLSTM, self).__init__() + + # ------------------------- + # Dynamic Input Handling + # ------------------------- + if one_batch is not None: + _, C, H, W = one_batch.shape + self.input_channels = C + self.feature_dim = H + self.seq_len = W + else: + self.input_channels = 1 + self.feature_dim = 64 + self.seq_len = 32 + + # The input to the LSTM will be C * H features per time step + self.lstm_input_size = self.input_channels * self.feature_dim + + # ------------------------- + # Bidirectional LSTM + # ------------------------- + self.lstm = nn.LSTM( + input_size=self.lstm_input_size, + hidden_size=hidden_size, + num_layers=num_layers, + batch_first=True, + bidirectional=True, + dropout=dropout if num_layers > 1 else 0.0 + ) + + # ------------------------- + # Temporal Attention Mechanism + # ------------------------- + # This projects the hidden state at each time step to a single importance score + self.attention = nn.Sequential( + nn.Linear(hidden_size * 2, hidden_size), + nn.Tanh(), + nn.Linear(hidden_size, 1) + ) + + # ------------------------- + # Classifier Head + # ------------------------- + # We classify based on the attention-weighted context vector + self.fc = nn.Linear(hidden_size * 2, num_classes) + + def forward(self, x): + # x shape arrives as: (Batch, Channels, Height, Width) + B, C, H, W = x.shape + + # 1. Permute to (Batch, Width, Channels, Height) + x = x.permute(0, 3, 1, 2) + + # 2. Reshape to (Batch, SequenceLength, Features) -> (B, W, C * H) + x = x.reshape(B, W, -1) + + # 3. Pass through LSTM + # lstm_out shape: (Batch, SequenceLength, hidden_size * 2) + lstm_out, _ = self.lstm(x) + + # 4. Calculate Attention Weights + # attn_scores shape: (Batch, SequenceLength, 1) + attn_scores = self.attention(lstm_out) + + # Normalize scores to probabilities across the time dimension + attn_weights = F.softmax(attn_scores, dim=1) + + # 5. Apply Attention Weights to LSTM outputs (Context Vector) + # Multiply each time step's hidden state by its attention weight, then sum + # context shape: (Batch, hidden_size * 2) + context = torch.sum(attn_weights * lstm_out, dim=1) + + # 6. Classify + return self.fc(context) \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/AUDIOMNIST/DynamicLSTM/model_aug_AUDIOMNIST_DynamicLSTM.py b/manuscripts/Poison26/bin/train/AUDIOMNIST/DynamicLSTM/model_aug_AUDIOMNIST_DynamicLSTM.py new file mode 100644 index 0000000..86adb1d --- /dev/null +++ b/manuscripts/Poison26/bin/train/AUDIOMNIST/DynamicLSTM/model_aug_AUDIOMNIST_DynamicLSTM.py @@ -0,0 +1,232 @@ +import argparse +import os +import glob +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader, Dataset, Subset + +import librosa +from sklearn.metrics import roc_auc_score, average_precision_score + +import hashlib +import csv +import random + +from DynamicLSTM import DynamicLSTM + +# ---------------------------- +# Constants +# ---------------------------- +SAMPLING_RATE = 16000 +NUM_CLASSES = 10 +MAX_AUDIO_LENGTH = 16000 + +# ---------------------------- +# Audio Preprocessing +# ---------------------------- +def normalize_audio(x): + max_val = np.max(np.abs(x)) + return x / max_val if max_val > 0 else x + +def pad_audio(audio, max_len=MAX_AUDIO_LENGTH): + return audio[:max_len] if len(audio) > max_len else np.pad(audio, (0, max_len - len(audio)), 'constant') + +# ---------------------------- +# Dataset & Wrapper +# ---------------------------- +class AudioMNISTBaseDataset(Dataset): + def __init__(self, data_path): + self.data = [] + self.labels = [] + + wav_files = glob.glob(os.path.join(data_path, '*', '*.wav')) + wav_files = sorted(wav_files, key=lambda x: hashlib.md5(x.encode()).hexdigest()) + self.wav_files = wav_files.copy() + + for audio_path in tqdm(wav_files, desc="Loading audio files"): + audio, _ = librosa.load(audio_path, sr=SAMPLING_RATE) + audio = normalize_audio(audio) + audio = pad_audio(audio) + label = int(os.path.basename(audio_path)[0]) + self.data.append(audio) + self.labels.append(label) + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return self.data[idx], self.labels[idx] + +class AudioSubsetWrapper(Dataset): + def __init__(self, subset, augment=False): + self.subset = subset + self.augment = augment + + def __len__(self): + return len(self.subset) + + def apply_augmentation(self, x): + if random.random() < 0.5: + x = np.clip(x + np.random.randn(len(x)) * 0.005, -1.0, 1.0) + if random.random() < 0.5: + x = np.roll(x, np.random.randint(-200, 200)) + if random.random() < 0.5: + x = np.clip(x * np.random.uniform(0.8, 1.2), -1.0, 1.0) + return x + + def __getitem__(self, idx): + x, y = self.subset[idx] + if self.augment: + x = self.apply_augmentation(x) + x = torch.tensor(x, dtype=torch.float32).unsqueeze(0) + return x, y + +def load_data(data_path, batch_size, augment_train=False, split_tsv="split_indices_standard.tsv"): + base_dataset = AudioMNISTBaseDataset(data_path) + + train_size = int(0.8 * len(base_dataset)) + train_indices = list(range(0, train_size)) + test_indices = list(range(train_size, len(base_dataset))) + + train_dataset = AudioSubsetWrapper(Subset(base_dataset, train_indices), augment=augment_train) + test_dataset = AudioSubsetWrapper(Subset(base_dataset, test_indices), augment=False) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) + + with open(split_tsv, "w", newline="") as f: + writer = csv.writer(f, delimiter="\t") + writer.writerow(["index", "split", "label", "file_path"]) + for idx in train_indices: + writer.writerow([idx, "train", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + for idx in test_indices: + writer.writerow([idx, "test", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + + return train_loader, test_loader + +# ---------------------------- +# Model Definition +# ---------------------------- +class AudioMNISTModel(nn.Module): + def __init__(self, backbone_class, num_classes=NUM_CLASSES): + super(AudioMNISTModel, self).__init__() + + self.reshape_dims = (1, 128, 125) + dummy_batch = torch.zeros(1, *self.reshape_dims) + self.backbone = backbone_class(one_batch=dummy_batch, num_classes=num_classes) + + def forward(self, x): + x = x.view(x.size(0), *self.reshape_dims) + return self.backbone(x) + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad(set_to_none=True) + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(NUM_CLASSES)[y_true] + + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="AudioMNIST Augmented Training") + parser.add_argument("--data", type=str, default="./data/AudioMNIST") + parser.add_argument("--output", type=str, default="audiomnist_dynamiclstm_aug.pt") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=10) + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + model = AudioMNISTModel(backbone_class=DynamicLSTM, num_classes=NUM_CLASSES) + + train_loader, test_loader = load_data(args.data, args.batch_size, augment_train=True, split_tsv="split_indices_aug.tsv") + + train(model, train_loader, device, epochs=args.epochs) + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + print("Model statistics on clean test dataset") + evaluate_model(model, test_loader, device) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/AUDIOMNIST/DynamicLSTM/model_fgsm_AUDIOMNIST_DynamicLSTM.py b/manuscripts/Poison26/bin/train/AUDIOMNIST/DynamicLSTM/model_fgsm_AUDIOMNIST_DynamicLSTM.py new file mode 100644 index 0000000..73a1fe2 --- /dev/null +++ b/manuscripts/Poison26/bin/train/AUDIOMNIST/DynamicLSTM/model_fgsm_AUDIOMNIST_DynamicLSTM.py @@ -0,0 +1,282 @@ +import argparse +import os +import glob +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader, Dataset, Subset + +import librosa +import torchaudio.transforms as T +from sklearn.metrics import roc_auc_score, average_precision_score + +import hashlib +import csv +import random + +from DynamicLSTM import DynamicLSTM + +# ---------------------------- +# Constants +# ---------------------------- +SAMPLING_RATE = 16000 +NUM_CLASSES = 10 +MAX_AUDIO_LENGTH = 16000 + +# ---------------------------- +# Audio Preprocessing +# ---------------------------- +def normalize_audio(x): + max_val = np.max(np.abs(x)) + return x / max_val if max_val > 0 else x + +def pad_audio(audio, max_len=MAX_AUDIO_LENGTH): + return audio[:max_len] if len(audio) > max_len else np.pad(audio, (0, max_len - len(audio)), 'constant') + +# ---------------------------- +# Dataset & Wrapper +# ---------------------------- +class AudioMNISTBaseDataset(Dataset): + def __init__(self, data_path): + self.data = [] + self.labels = [] + + wav_files = glob.glob(os.path.join(data_path, '*', '*.wav')) + wav_files = sorted(wav_files, key=lambda x: hashlib.md5(x.encode()).hexdigest()) + self.wav_files = wav_files.copy() + + for audio_path in tqdm(wav_files, desc="Loading audio files"): + audio, _ = librosa.load(audio_path, sr=SAMPLING_RATE) + audio = normalize_audio(audio) + audio = pad_audio(audio) + label = int(os.path.basename(audio_path)[0]) + self.data.append(audio) + self.labels.append(label) + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return self.data[idx], self.labels[idx] + +class AudioSubsetWrapper(Dataset): + def __init__(self, subset, augment=False): + self.subset = subset + self.augment = augment + + def __len__(self): + return len(self.subset) + + def apply_augmentation(self, x): + if random.random() < 0.5: + x = np.clip(x + np.random.randn(len(x)) * 0.005, -1.0, 1.0) + if random.random() < 0.5: + x = np.roll(x, np.random.randint(-200, 200)) + if random.random() < 0.5: + x = np.clip(x * np.random.uniform(0.8, 1.2), -1.0, 1.0) + return x + + def __getitem__(self, idx): + x, y = self.subset[idx] + if self.augment: + x = self.apply_augmentation(x) + x = torch.tensor(x, dtype=torch.float32).unsqueeze(0) + return x, y + +def load_data(data_path, batch_size, augment_train=False, split_tsv="split_indices_standard.tsv"): + base_dataset = AudioMNISTBaseDataset(data_path) + + train_size = int(0.8 * len(base_dataset)) + train_indices = list(range(0, train_size)) + test_indices = list(range(train_size, len(base_dataset))) + + train_dataset = AudioSubsetWrapper(Subset(base_dataset, train_indices), augment=augment_train) + test_dataset = AudioSubsetWrapper(Subset(base_dataset, test_indices), augment=False) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) + + with open(split_tsv, "w", newline="") as f: + writer = csv.writer(f, delimiter="\t") + writer.writerow(["index", "split", "label", "file_path"]) + for idx in train_indices: + writer.writerow([idx, "train", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + for idx in test_indices: + writer.writerow([idx, "test", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + + return train_loader, test_loader + +# ---------------------------- +# Model Definition +# ---------------------------- +class AudioMNISTModel(nn.Module): + def __init__(self, backbone_class, num_classes=NUM_CLASSES): + super(AudioMNISTModel, self).__init__() + self.spectrogram = T.MelSpectrogram( + sample_rate=SAMPLING_RATE, + n_fft=1024, + hop_length=512, + n_mels=64 + ) + dummy_batch = torch.zeros(1, 1, 64, 32) + self.backbone = backbone_class(one_batch=dummy_batch, num_classes=num_classes) + + def forward(self, x): + x = self.spectrogram(x) + x = torch.log(x + 1e-9) + return self.backbone(x) + +# ---------------------------- +# FGSM attack on 1D Audio [-1, 1] +# ---------------------------- +@torch.enable_grad() +def fgsm_attack(model, x, y, eps=0.05, clamp_min=-1.0, clamp_max=1.0): + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + x.requires_grad_(True) + logits = model(x) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x, retain_graph=False, create_graph=False)[0] + x_adv = x + eps * torch.sign(grad.detach()) + + x_adv = x_adv.clamp(clamp_min, clamp_max) + return x_adv.detach() + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, adv_lambda=1.0): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for inputs, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + inputs, labels = inputs.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + if adv: + model.eval() + inputs_adv = fgsm_attack(model, inputs, labels, eps=eps, clamp_min=-1.0, clamp_max=1.0) + model.train() + else: + inputs_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(inputs_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(inputs) + out_adv = model(inputs_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv + else: + outputs = model(inputs) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + test_loss = 0.0 + correct = 0 + total = 0 + + for inputs, labels in test_loader: + inputs, labels = inputs.to(device), labels.to(device) + + if robust: + inputs_in = fgsm_attack(model, inputs, labels, eps=eps, clamp_min=-1.0, clamp_max=1.0) + else: + inputs_in = inputs + + with torch.no_grad(): + outputs = model(inputs_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(NUM_CLASSES)[y_true] + + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (FGSM)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f} | Acc: {accuracy:.4f} | auROC: {auroc:.4f} | auPRC: {auprc:.4f}") + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="AudioMNIST FGSM Training") + parser.add_argument("--data", type=str, default="./data/AudioMNIST") + parser.add_argument("--output", type=str, default="audiomnist_dynamiclstm_fgsm.pt") + parser.add_argument("--epochs", type=int, default=10) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--adv-train", action="store_true", help="Enable FGSM adversarial training") + parser.add_argument("--eps", type=float, default=0.05) + parser.add_argument("--adv-lambda", type=float, default=1.0) + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + model = AudioMNISTModel(backbone_class=DynamicLSTM, num_classes=NUM_CLASSES) + + train_loader, test_loader = load_data(args.data, args.batch_size, augment_train=False, split_tsv="split_indices_fgsm.tsv") + + train(model, train_loader, device, epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, adv_lambda=args.adv_lambda) + + torch.save(model.state_dict(), args.output) + + print("\nEvaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/AUDIOMNIST/DynamicLSTM/model_pgd_AUDIOMNIST_DynamicLSTM.py b/manuscripts/Poison26/bin/train/AUDIOMNIST/DynamicLSTM/model_pgd_AUDIOMNIST_DynamicLSTM.py new file mode 100644 index 0000000..aced1a5 --- /dev/null +++ b/manuscripts/Poison26/bin/train/AUDIOMNIST/DynamicLSTM/model_pgd_AUDIOMNIST_DynamicLSTM.py @@ -0,0 +1,298 @@ +import argparse +import os +import glob +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader, Dataset, Subset + +import librosa +import torchaudio.transforms as T +from sklearn.metrics import roc_auc_score, average_precision_score + +import hashlib +import csv +import random + +from DynamicLSTM import DynamicLSTM + +# ---------------------------- +# Constants +# ---------------------------- +SAMPLING_RATE = 16000 +NUM_CLASSES = 10 +MAX_AUDIO_LENGTH = 16000 + +# ---------------------------- +# Audio Preprocessing +# ---------------------------- +def normalize_audio(x): + max_val = np.max(np.abs(x)) + return x / max_val if max_val > 0 else x + +def pad_audio(audio, max_len=MAX_AUDIO_LENGTH): + return audio[:max_len] if len(audio) > max_len else np.pad(audio, (0, max_len - len(audio)), 'constant') + +# ---------------------------- +# Dataset & Wrapper +# ---------------------------- +class AudioMNISTBaseDataset(Dataset): + def __init__(self, data_path): + self.data = [] + self.labels = [] + + wav_files = glob.glob(os.path.join(data_path, '*', '*.wav')) + wav_files = sorted(wav_files, key=lambda x: hashlib.md5(x.encode()).hexdigest()) + self.wav_files = wav_files.copy() + + for audio_path in tqdm(wav_files, desc="Loading audio files"): + audio, _ = librosa.load(audio_path, sr=SAMPLING_RATE) + audio = normalize_audio(audio) + audio = pad_audio(audio) + label = int(os.path.basename(audio_path)[0]) + self.data.append(audio) + self.labels.append(label) + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return self.data[idx], self.labels[idx] + +class AudioSubsetWrapper(Dataset): + def __init__(self, subset, augment=False): + self.subset = subset + self.augment = augment + + def __len__(self): + return len(self.subset) + + def apply_augmentation(self, x): + if random.random() < 0.5: + x = np.clip(x + np.random.randn(len(x)) * 0.005, -1.0, 1.0) + if random.random() < 0.5: + x = np.roll(x, np.random.randint(-200, 200)) + if random.random() < 0.5: + x = np.clip(x * np.random.uniform(0.8, 1.2), -1.0, 1.0) + return x + + def __getitem__(self, idx): + x, y = self.subset[idx] + if self.augment: + x = self.apply_augmentation(x) + x = torch.tensor(x, dtype=torch.float32).unsqueeze(0) + return x, y + +def load_data(data_path, batch_size, augment_train=False, split_tsv="split_indices_standard.tsv"): + base_dataset = AudioMNISTBaseDataset(data_path) + + train_size = int(0.8 * len(base_dataset)) + train_indices = list(range(0, train_size)) + test_indices = list(range(train_size, len(base_dataset))) + + train_dataset = AudioSubsetWrapper(Subset(base_dataset, train_indices), augment=augment_train) + test_dataset = AudioSubsetWrapper(Subset(base_dataset, test_indices), augment=False) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) + + with open(split_tsv, "w", newline="") as f: + writer = csv.writer(f, delimiter="\t") + writer.writerow(["index", "split", "label", "file_path"]) + for idx in train_indices: + writer.writerow([idx, "train", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + for idx in test_indices: + writer.writerow([idx, "test", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + + return train_loader, test_loader + +# ---------------------------- +# Model Definition +# ---------------------------- +class AudioMNISTModel(nn.Module): + def __init__(self, backbone_class, num_classes=NUM_CLASSES): + super(AudioMNISTModel, self).__init__() + + self.spectrogram = T.MelSpectrogram( + sample_rate=SAMPLING_RATE, + n_fft=1024, + hop_length=512, + n_mels=64 + ) + + dummy_batch = torch.zeros(1, 1, 64, 32) + self.backbone = backbone_class(one_batch=dummy_batch, num_classes=num_classes) + + def forward(self, x): + x = self.spectrogram(x) + x = torch.log(x + 1e-9) + return self.backbone(x) + +# ---------------------------- +# PGD attack on 1D Audio [-1, 1] +# ---------------------------- +@torch.enable_grad() +def pgd_attack(model, x, y, eps=0.05, alpha=0.01, steps=10, random_start=True, clamp_min=-1.0, clamp_max=1.0): + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + if random_start: + x_adv = x + torch.empty_like(x).uniform_(-eps, eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + else: + x_adv = x.clone() + + for _ in range(steps): + x_adv.requires_grad_(True) + logits = model(x_adv) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x_adv, retain_graph=False, create_graph=False)[0] + x_adv = x_adv.detach() + alpha * torch.sign(grad.detach()) + + x_adv = torch.max(torch.min(x_adv, x + eps), x - eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + + return x_adv.detach() + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, alpha=0.01, pgd_steps=10, random_start=True, adv_lambda=1.0): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for inputs, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + inputs, labels = inputs.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + if adv: + model.eval() + inputs_adv = pgd_attack(model, inputs, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=random_start, clamp_min=-1.0, clamp_max=1.0) + model.train() + else: + inputs_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(inputs_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(inputs) + out_adv = model(inputs_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv + else: + outputs = model(inputs) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05, alpha=0.01, pgd_steps=10): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + test_loss = 0.0 + correct = 0 + total = 0 + + for inputs, labels in test_loader: + inputs, labels = inputs.to(device), labels.to(device) + + if robust: + inputs_in = pgd_attack(model, inputs, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=True, clamp_min=-1.0, clamp_max=1.0) + else: + inputs_in = inputs + + with torch.no_grad(): + outputs = model(inputs_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(NUM_CLASSES)[y_true] + + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (PGD)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f} | Acc: {accuracy:.4f} | auROC: {auroc:.4f} | auPRC: {auprc:.4f}") + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="AudioMNIST PGD Adversarial Training") + parser.add_argument("--data", type=str, default="./data/AudioMNIST") + parser.add_argument("--output", type=str, default="audiomnist_dynamiclstm_pgd.pt") + parser.add_argument("--epochs", type=int, default=10) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + parser.add_argument("--adv-train", action="store_true", help="Enable PGD adversarial training") + parser.add_argument("--eps", type=float, default=0.05) + parser.add_argument("--alpha", type=float, default=0.01) + parser.add_argument("--pgd-steps", type=int, default=10) + parser.add_argument("--no-random-start", action="store_true") + parser.add_argument("--adv-lambda", type=float, default=1.0) + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + model = AudioMNISTModel(backbone_class=DynamicLSTM, num_classes=NUM_CLASSES) + + train_loader, test_loader = load_data(args.data, args.batch_size, augment_train=False, split_tsv="split_indices_pgd.tsv") + + train(model, train_loader, device, epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, alpha=args.alpha, + pgd_steps=args.pgd_steps, random_start=not args.no_random_start, adv_lambda=args.adv_lambda) + + torch.save(model.state_dict(), args.output) + + print("\nEvaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps, alpha=args.alpha, pgd_steps=args.pgd_steps) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/AUDIOMNIST/DynamicLSTM/model_standard_AUDIOMNIST_DynamicLSTM.py b/manuscripts/Poison26/bin/train/AUDIOMNIST/DynamicLSTM/model_standard_AUDIOMNIST_DynamicLSTM.py new file mode 100644 index 0000000..31aa294 --- /dev/null +++ b/manuscripts/Poison26/bin/train/AUDIOMNIST/DynamicLSTM/model_standard_AUDIOMNIST_DynamicLSTM.py @@ -0,0 +1,237 @@ +import argparse +import os +import glob +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader, Dataset, Subset + +import librosa +import torchaudio.transforms as T +from sklearn.metrics import roc_auc_score, average_precision_score + +import hashlib +import csv +import random + +from DynamicLSTM import DynamicLSTM + +# ---------------------------- +# Constants +# ---------------------------- +SAMPLING_RATE = 16000 +NUM_CLASSES = 10 +MAX_AUDIO_LENGTH = 16000 + +# ---------------------------- +# Audio Preprocessing +# ---------------------------- +def normalize_audio(x): + max_val = np.max(np.abs(x)) + return x / max_val if max_val > 0 else x + +def pad_audio(audio, max_len=MAX_AUDIO_LENGTH): + return audio[:max_len] if len(audio) > max_len else np.pad(audio, (0, max_len - len(audio)), 'constant') + +# ---------------------------- +# Dataset & Wrapper +# ---------------------------- +class AudioMNISTBaseDataset(Dataset): + def __init__(self, data_path): + self.data = [] + self.labels = [] + + wav_files = glob.glob(os.path.join(data_path, '*', '*.wav')) + wav_files = sorted(wav_files, key=lambda x: hashlib.md5(x.encode()).hexdigest()) + self.wav_files = wav_files.copy() + + for audio_path in tqdm(wav_files, desc="Loading audio files"): + audio, _ = librosa.load(audio_path, sr=SAMPLING_RATE) + audio = normalize_audio(audio) + audio = pad_audio(audio) + label = int(os.path.basename(audio_path)[0]) + self.data.append(audio) + self.labels.append(label) + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return self.data[idx], self.labels[idx] + +class AudioSubsetWrapper(Dataset): + def __init__(self, subset, augment=False): + self.subset = subset + self.augment = augment + + def __len__(self): + return len(self.subset) + + def apply_augmentation(self, x): + if random.random() < 0.5: + x = np.clip(x + np.random.randn(len(x)) * 0.005, -1.0, 1.0) + if random.random() < 0.5: + x = np.roll(x, np.random.randint(-200, 200)) + if random.random() < 0.5: + x = np.clip(x * np.random.uniform(0.8, 1.2), -1.0, 1.0) + return x + + def __getitem__(self, idx): + x, y = self.subset[idx] + if self.augment: + x = self.apply_augmentation(x) + x = torch.tensor(x, dtype=torch.float32).unsqueeze(0) + return x, y + +def load_data(data_path, batch_size, augment_train=False, split_tsv="split_indices_standard.tsv"): + base_dataset = AudioMNISTBaseDataset(data_path) + + train_size = int(0.8 * len(base_dataset)) + train_indices = list(range(0, train_size)) + test_indices = list(range(train_size, len(base_dataset))) + + train_dataset = AudioSubsetWrapper(Subset(base_dataset, train_indices), augment=augment_train) + test_dataset = AudioSubsetWrapper(Subset(base_dataset, test_indices), augment=False) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) + + with open(split_tsv, "w", newline="") as f: + writer = csv.writer(f, delimiter="\t") + writer.writerow(["index", "split", "label", "file_path"]) + for idx in train_indices: + writer.writerow([idx, "train", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + for idx in test_indices: + writer.writerow([idx, "test", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + + return train_loader, test_loader + +# ---------------------------- +# Model Definition +# ---------------------------- +class AudioMNISTModel(nn.Module): + def __init__(self, backbone_class, num_classes=NUM_CLASSES): + super(AudioMNISTModel, self).__init__() + + self.spectrogram = T.MelSpectrogram( + sample_rate=SAMPLING_RATE, + n_fft=1024, + hop_length=512, + n_mels=64 + ) + + dummy_batch = torch.zeros(1, 1, 64, 32) + self.backbone = backbone_class(one_batch=dummy_batch, num_classes=num_classes) + + def forward(self, x): + x = self.spectrogram(x) + x = torch.log(x + 1e-9) + return self.backbone(x) + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for inputs, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + inputs, labels = inputs.to(device), labels.to(device) + + optimizer.zero_grad(set_to_none=True) + outputs = model(inputs) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for inputs, labels in test_loader: + inputs, labels = inputs.to(device), labels.to(device) + outputs = model(inputs) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(NUM_CLASSES)[y_true] + + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f} | Acc: {accuracy:.4f} | auROC: {auroc:.4f} | auPRC: {auprc:.4f}") + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="AudioMNIST Standard Training") + parser.add_argument("--data", type=str, default="./data/AudioMNIST") + parser.add_argument("--output", type=str, default="audiomnist_dynamiclstm_standard.pt") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=10) + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + model = AudioMNISTModel(backbone_class=DynamicLSTM, num_classes=NUM_CLASSES) + + train_loader, test_loader = load_data(args.data, args.batch_size, augment_train=False) + + train(model, train_loader, device, epochs=args.epochs) + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/AUDIOMNIST/HybridViT/HybridViT.py b/manuscripts/Poison26/bin/train/AUDIOMNIST/HybridViT/HybridViT.py new file mode 100644 index 0000000..829ab66 --- /dev/null +++ b/manuscripts/Poison26/bin/train/AUDIOMNIST/HybridViT/HybridViT.py @@ -0,0 +1,115 @@ +import torch +import torch.nn as nn + +class HybridViT(nn.Module): + """ + A Convolutional-Transformer Hybrid Architecture. + Dynamically sizes its positional embeddings and patch count based on the input batch, + making it easily adaptable to different image or spectrogram resolutions. + """ + def __init__(self, one_batch=None, num_classes=10, embed_dim=64, num_heads=4, depth=2): + super(HybridViT, self).__init__() + + # ------------------------- + # Dynamic Input Handling + # ------------------------- + if one_batch is not None: + _, in_channels, H, W = one_batch.shape + self.input_channels = in_channels + self.input_size = (in_channels, H, W) + else: + self.input_channels = 1 + self.input_size = (1, 28, 28) # Default to MNIST size + + self.embed_dim = embed_dim + + # ------------------------- + # 1. CNN Stem (Local Feature Extraction) + # ------------------------- + self.stem = nn.Sequential( + nn.Conv2d(self.input_channels, 32, kernel_size=3, stride=2, padding=1, bias=False), + nn.BatchNorm2d(32), + nn.ReLU(inplace=True), + nn.Conv2d(32, self.embed_dim, kernel_size=3, stride=2, padding=1, bias=False), + nn.BatchNorm2d(self.embed_dim), + nn.ReLU(inplace=True) + ) + + # ------------------------- + # 2. Dynamic Patch Calculation + # ------------------------- + self.num_patches = self._get_num_patches(one_batch) + + # Learnable positional embedding to retain spatial awareness + self.pos_embed = nn.Parameter(torch.randn(1, self.num_patches, self.embed_dim)) + + # ------------------------- + # 3. Transformer Encoder (Global Context) + # ------------------------- + encoder_layer = nn.TransformerEncoderLayer( + d_model=self.embed_dim, + nhead=num_heads, + dim_feedforward=self.embed_dim * 4, + activation='gelu', + batch_first=True + ) + self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=depth) + + # ------------------------- + # 4. Classifier Head + # ------------------------- + self.norm = nn.LayerNorm(self.embed_dim) + self.fc = nn.Linear(self.embed_dim, num_classes) + + # ------------------------- + # Compute Patch Size Dynamically + # ------------------------- + def _get_num_patches(self, one_batch): + # Explicitly set eval() to prevent tracking BatchNorm stats with the dummy batch + was_training = self.training + self.eval() + + with torch.no_grad(): + if one_batch is None: + dummy_input = torch.zeros(1, *self.input_size) + else: + _, C, H, W = one_batch.shape + dummy_input = torch.zeros(1, C, H, W) + + # Pass through the stem + x = self.stem(dummy_input) + + # Extract resulting spatial dimensions + _, _, h_out, w_out = x.shape + num_patches = h_out * w_out + + if was_training: + self.train() + + return num_patches + + # ------------------------- + # Forward + # ------------------------- + def forward(self, x): + # 1. Pass through CNN stem + x = self.stem(x) + + # 2. Flatten spatial dimensions to create a sequence + B, C, H, W = x.shape + x = x.flatten(2).transpose(1, 2) # Shape: (Batch, num_patches, embed_dim) + + # 3. Add positional embeddings + x = x + self.pos_embed + + # 4. Pass through Transformer + x = self.transformer(x) + + # 5. Global Average Pooling over the sequence length + x = x.mean(dim=1) # Shape: (Batch, embed_dim) + + # 6. Classify + x = self.norm(x) + x = self.fc(x) + + return x \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/AUDIOMNIST/HybridViT/model_aug_AUDIOMNIST_HybridViT.py b/manuscripts/Poison26/bin/train/AUDIOMNIST/HybridViT/model_aug_AUDIOMNIST_HybridViT.py new file mode 100644 index 0000000..61b19b1 --- /dev/null +++ b/manuscripts/Poison26/bin/train/AUDIOMNIST/HybridViT/model_aug_AUDIOMNIST_HybridViT.py @@ -0,0 +1,238 @@ +import argparse +import os +import glob +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader, Dataset, Subset + +import librosa +from sklearn.metrics import roc_auc_score, average_precision_score + +import hashlib +import csv +import random + +# --- UPDATED IMPORT --- +from HybridViT import HybridViT + +# ---------------------------- +# Constants +# ---------------------------- +SAMPLING_RATE = 16000 +NUM_CLASSES = 10 +MAX_AUDIO_LENGTH = 16000 + +# ---------------------------- +# Audio Preprocessing +# ---------------------------- +def normalize_audio(x): + max_val = np.max(np.abs(x)) + return x / max_val if max_val > 0 else x + +def pad_audio(audio, max_len=MAX_AUDIO_LENGTH): + return audio[:max_len] if len(audio) > max_len else np.pad(audio, (0, max_len - len(audio)), 'constant') + +# ---------------------------- +# Dataset & Wrapper +# ---------------------------- +class AudioMNISTBaseDataset(Dataset): + """Loads all audio into memory ONCE. Returns raw numpy arrays.""" + def __init__(self, data_path): + self.data = [] + self.labels = [] + + wav_files = glob.glob(os.path.join(data_path, '*', '*.wav')) + wav_files = sorted(wav_files, key=lambda x: hashlib.md5(x.encode()).hexdigest()) + self.wav_files = wav_files.copy() + + for audio_path in tqdm(wav_files, desc="Loading audio files"): + audio, _ = librosa.load(audio_path, sr=SAMPLING_RATE) + audio = normalize_audio(audio) + audio = pad_audio(audio) + label = int(os.path.basename(audio_path)[0]) + self.data.append(audio) + self.labels.append(label) + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return self.data[idx], self.labels[idx] + +class AudioSubsetWrapper(Dataset): + """Wraps a subset to apply dynamic augmentation and convert to Tensors.""" + def __init__(self, subset, augment=False): + self.subset = subset + self.augment = augment + + def __len__(self): + return len(self.subset) + + def apply_augmentation(self, x): + if random.random() < 0.5: + x = np.clip(x + np.random.randn(len(x)) * 0.005, -1.0, 1.0) # noise + if random.random() < 0.5: + x = np.roll(x, np.random.randint(-200, 200)) # time shift + if random.random() < 0.5: + x = np.clip(x * np.random.uniform(0.8, 1.2), -1.0, 1.0) # random gain + return x + + def __getitem__(self, idx): + x, y = self.subset[idx] + if self.augment: + x = self.apply_augmentation(x) + x = torch.tensor(x, dtype=torch.float32).unsqueeze(0) # (1, length) + return x, y + +def load_data(data_path, batch_size, augment_train=False, split_tsv="split_indices_standard.tsv"): + base_dataset = AudioMNISTBaseDataset(data_path) + + train_size = int(0.8 * len(base_dataset)) + train_indices = list(range(0, train_size)) + test_indices = list(range(train_size, len(base_dataset))) + + # Isolate augmentation using the wrapper + train_dataset = AudioSubsetWrapper(Subset(base_dataset, train_indices), augment=augment_train) + test_dataset = AudioSubsetWrapper(Subset(base_dataset, test_indices), augment=False) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) + + with open(split_tsv, "w", newline="") as f: + writer = csv.writer(f, delimiter="\t") + writer.writerow(["index", "split", "label", "file_path"]) + for idx in train_indices: + writer.writerow([idx, "train", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + for idx in test_indices: + writer.writerow([idx, "test", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + print(f"Saved split information to {split_tsv}") + + return train_loader, test_loader + +# ---------------------------- +# Model Definition +# ---------------------------- +class AudioMNISTModel(nn.Module): + def __init__(self, backbone_class, num_classes=NUM_CLASSES): + super(AudioMNISTModel, self).__init__() + + self.reshape_dims = (1, 128, 125) + dummy_batch = torch.zeros(1, *self.reshape_dims) + self.backbone = backbone_class(one_batch=dummy_batch, num_classes=num_classes) + + def forward(self, x): + x = x.view(x.size(0), *self.reshape_dims) + return self.backbone(x) + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad(set_to_none=True) + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(NUM_CLASSES)[y_true] + + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="AudioMNIST Augmented Training") + parser.add_argument("--data", type=str, default="./data/AudioMNIST", help="Path to dataset") + parser.add_argument("--output", type=str, default="audiomnist_hybridvit_aug.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=10) + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # --- UPDATED MODEL CALL --- + model = AudioMNISTModel(backbone_class=HybridViT, num_classes=NUM_CLASSES) + + train_loader, test_loader = load_data(args.data, args.batch_size, augment_train=True, split_tsv="split_indices_aug.tsv") + + train(model, train_loader, device, epochs=args.epochs) + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + print("Model statistics on clean test dataset") + evaluate_model(model, test_loader, device) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/AUDIOMNIST/HybridViT/model_fgsm_AUDIOMNIST_HybridViT.py b/manuscripts/Poison26/bin/train/AUDIOMNIST/HybridViT/model_fgsm_AUDIOMNIST_HybridViT.py new file mode 100644 index 0000000..5b7424d --- /dev/null +++ b/manuscripts/Poison26/bin/train/AUDIOMNIST/HybridViT/model_fgsm_AUDIOMNIST_HybridViT.py @@ -0,0 +1,284 @@ +import argparse +import os +import glob +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader, Dataset, Subset + +import librosa +import torchaudio.transforms as T +from sklearn.metrics import roc_auc_score, average_precision_score + +import hashlib +import csv +import random + +# --- UPDATED IMPORT --- +from HybridViT import HybridViT + +# ---------------------------- +# Constants +# ---------------------------- +SAMPLING_RATE = 16000 +NUM_CLASSES = 10 +MAX_AUDIO_LENGTH = 16000 + +# ---------------------------- +# Audio Preprocessing +# ---------------------------- +def normalize_audio(x): + max_val = np.max(np.abs(x)) + return x / max_val if max_val > 0 else x + +def pad_audio(audio, max_len=MAX_AUDIO_LENGTH): + return audio[:max_len] if len(audio) > max_len else np.pad(audio, (0, max_len - len(audio)), 'constant') + +# ---------------------------- +# Dataset & Wrapper +# ---------------------------- +class AudioMNISTBaseDataset(Dataset): + def __init__(self, data_path): + self.data = [] + self.labels = [] + + wav_files = glob.glob(os.path.join(data_path, '*', '*.wav')) + wav_files = sorted(wav_files, key=lambda x: hashlib.md5(x.encode()).hexdigest()) + self.wav_files = wav_files.copy() + + for audio_path in tqdm(wav_files, desc="Loading audio files"): + audio, _ = librosa.load(audio_path, sr=SAMPLING_RATE) + audio = normalize_audio(audio) + audio = pad_audio(audio) + label = int(os.path.basename(audio_path)[0]) + self.data.append(audio) + self.labels.append(label) + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return self.data[idx], self.labels[idx] + +class AudioSubsetWrapper(Dataset): + def __init__(self, subset, augment=False): + self.subset = subset + self.augment = augment + + def __len__(self): + return len(self.subset) + + def apply_augmentation(self, x): + if random.random() < 0.5: + x = np.clip(x + np.random.randn(len(x)) * 0.005, -1.0, 1.0) + if random.random() < 0.5: + x = np.roll(x, np.random.randint(-200, 200)) + if random.random() < 0.5: + x = np.clip(x * np.random.uniform(0.8, 1.2), -1.0, 1.0) + return x + + def __getitem__(self, idx): + x, y = self.subset[idx] + if self.augment: + x = self.apply_augmentation(x) + x = torch.tensor(x, dtype=torch.float32).unsqueeze(0) + return x, y + +def load_data(data_path, batch_size, augment_train=False, split_tsv="split_indices_standard.tsv"): + base_dataset = AudioMNISTBaseDataset(data_path) + + train_size = int(0.8 * len(base_dataset)) + train_indices = list(range(0, train_size)) + test_indices = list(range(train_size, len(base_dataset))) + + train_dataset = AudioSubsetWrapper(Subset(base_dataset, train_indices), augment=augment_train) + test_dataset = AudioSubsetWrapper(Subset(base_dataset, test_indices), augment=False) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) + + with open(split_tsv, "w", newline="") as f: + writer = csv.writer(f, delimiter="\t") + writer.writerow(["index", "split", "label", "file_path"]) + for idx in train_indices: + writer.writerow([idx, "train", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + for idx in test_indices: + writer.writerow([idx, "test", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + + return train_loader, test_loader + +# ---------------------------- +# Model Definition +# ---------------------------- +class AudioMNISTModel(nn.Module): + def __init__(self, backbone_class, num_classes=NUM_CLASSES): + super(AudioMNISTModel, self).__init__() + self.spectrogram = T.MelSpectrogram( + sample_rate=SAMPLING_RATE, + n_fft=1024, + hop_length=512, + n_mels=64 + ) + dummy_batch = torch.zeros(1, 1, 64, 32) + self.backbone = backbone_class(one_batch=dummy_batch, num_classes=num_classes) + + def forward(self, x): + x = self.spectrogram(x) + x = torch.log(x + 1e-9) + return self.backbone(x) + +# ---------------------------- +# FGSM attack on 1D Audio [-1, 1] +# ---------------------------- +@torch.enable_grad() +def fgsm_attack(model, x, y, eps=0.05, clamp_min=-1.0, clamp_max=1.0): + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + x.requires_grad_(True) + logits = model(x) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x, retain_graph=False, create_graph=False)[0] + x_adv = x + eps * torch.sign(grad.detach()) + + x_adv = x_adv.clamp(clamp_min, clamp_max) + return x_adv.detach() + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, adv_lambda=1.0): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for inputs, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + inputs, labels = inputs.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + if adv: + model.eval() + inputs_adv = fgsm_attack(model, inputs, labels, eps=eps, clamp_min=-1.0, clamp_max=1.0) + model.train() + else: + inputs_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(inputs_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(inputs) + out_adv = model(inputs_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv + else: + outputs = model(inputs) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + test_loss = 0.0 + correct = 0 + total = 0 + + for inputs, labels in test_loader: + inputs, labels = inputs.to(device), labels.to(device) + + if robust: + inputs_in = fgsm_attack(model, inputs, labels, eps=eps, clamp_min=-1.0, clamp_max=1.0) + else: + inputs_in = inputs + + with torch.no_grad(): + outputs = model(inputs_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(NUM_CLASSES)[y_true] + + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (FGSM)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f} | Acc: {accuracy:.4f} | auROC: {auroc:.4f} | auPRC: {auprc:.4f}") + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="AudioMNIST FGSM Training") + parser.add_argument("--data", type=str, default="./data/AudioMNIST") + parser.add_argument("--output", type=str, default="audiomnist_hybridvit_fgsm.pt") + parser.add_argument("--epochs", type=int, default=10) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--adv-train", action="store_true", help="Enable FGSM adversarial training") + parser.add_argument("--eps", type=float, default=0.05) + parser.add_argument("--adv-lambda", type=float, default=1.0) + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # --- UPDATED MODEL CALL --- + model = AudioMNISTModel(backbone_class=HybridViT, num_classes=NUM_CLASSES) + + train_loader, test_loader = load_data(args.data, args.batch_size, augment_train=False, split_tsv="split_indices_fgsm.tsv") + + train(model, train_loader, device, epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, adv_lambda=args.adv_lambda) + + torch.save(model.state_dict(), args.output) + + print("\nEvaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/AUDIOMNIST/HybridViT/model_pgd_AUDIOMNIST_HybridViT.py b/manuscripts/Poison26/bin/train/AUDIOMNIST/HybridViT/model_pgd_AUDIOMNIST_HybridViT.py new file mode 100644 index 0000000..b10f678 --- /dev/null +++ b/manuscripts/Poison26/bin/train/AUDIOMNIST/HybridViT/model_pgd_AUDIOMNIST_HybridViT.py @@ -0,0 +1,265 @@ +import argparse +import os +import glob +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader, Dataset, Subset + +import librosa +import torchaudio.transforms as T +from sklearn.metrics import roc_auc_score, average_precision_score + +import hashlib +import csv +import random + +# --- UPDATED IMPORT --- +from HybridViT import HybridViT + +# ---------------------------- +# Constants +# ---------------------------- +SAMPLING_RATE = 16000 +NUM_CLASSES = 10 +MAX_AUDIO_LENGTH = 16000 + +# ---------------------------- +# Audio Preprocessing +# ---------------------------- +def normalize_audio(x): + max_val = np.max(np.abs(x)) + return x / max_val if max_val > 0 else x + +def pad_audio(audio, max_len=MAX_AUDIO_LENGTH): + return audio[:max_len] if len(audio) > max_len else np.pad(audio, (0, max_len - len(audio)), 'constant') + +# ---------------------------- +# Dataset & Wrapper +# ---------------------------- +class AudioMNISTBaseDataset(Dataset): + def __init__(self, data_path): + self.data = [] + self.labels = [] + + wav_files = glob.glob(os.path.join(data_path, '*', '*.wav')) + wav_files = sorted(wav_files, key=lambda x: hashlib.md5(x.encode()).hexdigest()) + self.wav_files = wav_files.copy() + + for audio_path in tqdm(wav_files, desc="Loading audio files"): + audio, _ = librosa.load(audio_path, sr=SAMPLING_RATE) + audio = normalize_audio(audio) + audio = pad_audio(audio) + label = int(os.path.basename(audio_path)[0]) + self.data.append(audio) + self.labels.append(label) + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return self.data[idx], self.labels[idx] + +class AudioSubsetWrapper(Dataset): + def __init__(self, subset, augment=False): + self.subset = subset + self.augment = augment + + def __len__(self): + return len(self.subset) + + def apply_augmentation(self, x): + if random.random() < 0.5: + x = np.clip(x + np.random.randn(len(x)) * 0.005, -1.0, 1.0) + if random.random() < 0.5: + x = np.roll(x, np.random.randint(-200, 200)) + if random.random() < 0.5: + x = np.clip(x * np.random.uniform(0.8, 1.2), -1.0, 1.0) + return x + + def __getitem__(self, idx): + x, y = self.subset[idx] + if self.augment: + x = self.apply_augmentation(x) + x = torch.tensor(x, dtype=torch.float32).unsqueeze(0) + return x, y + +def load_data(data_path, batch_size, augment_train=False, split_tsv="split_indices_standard.tsv"): + base_dataset = AudioMNISTBaseDataset(data_path) + + train_size = int(0.8 * len(base_dataset)) + train_indices = list(range(0, train_size)) + test_indices = list(range(train_size, len(base_dataset))) + + train_dataset = AudioSubsetWrapper(Subset(base_dataset, train_indices), augment=augment_train) + test_dataset = AudioSubsetWrapper(Subset(base_dataset, test_indices), augment=False) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) + + with open(split_tsv, "w", newline="") as f: + writer = csv.writer(f, delimiter="\t") + writer.writerow(["index", "split", "label", "file_path"]) + for idx in train_indices: + writer.writerow([idx, "train", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + for idx in test_indices: + writer.writerow([idx, "test", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + + return train_loader, test_loader + +# ---------------------------- +# Model Definition +# ---------------------------- +class AudioMNISTModel(nn.Module): + def __init__(self, backbone_class, num_classes=NUM_CLASSES): + super(AudioMNISTModel, self).__init__() + + self.spectrogram = T.MelSpectrogram( + sample_rate=SAMPLING_RATE, + n_fft=1024, + hop_length=512, + n_mels=64 + ) + + dummy_batch = torch.zeros(1, 1, 64, 32) + self.backbone = backbone_class(one_batch=dummy_batch, num_classes=num_classes) + + def forward(self, x): + x = self.spectrogram(x) + x = torch.log(x + 1e-9) + return self.backbone(x) + +# ---------------------------- +# PGD attack on 1D Audio [-1, 1] +# ---------------------------- +@torch.enable_grad() +def pgd_attack(model, x, y, eps=0.05, alpha=0.01, steps=10, random_start=True, clamp_min=-1.0, clamp_max=1.0): + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + if random_start: + x_adv = x + torch.empty_like(x).uniform_(-eps, eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + else: + x_adv = x.clone() + + for _ in range(steps): + x_adv.requires_grad_(True) + logits = model(x_adv) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x_adv, retain_graph=False, create_graph=False)[0] + x_adv = x_adv.detach() + alpha * torch.sign(grad.detach()) + + x_adv = torch.max(torch.min(x_adv, x + eps), x - eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + + return x_adv.detach() + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, alpha=0.01, pgd_steps=10, random_start=True, adv_lambda=1.0): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for inputs, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + inputs, labels = inputs.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + if adv: + model.eval() + inputs_adv = pgd_attack(model, inputs, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=random_start, clamp_min=-1.0, clamp_max=1.0) + model.train() + else: + inputs_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(inputs_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(inputs) + out_adv = model(inputs_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv + else: + outputs = model(inputs) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05, alpha=0.01, pgd_steps=10): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + test_loss = 0.0 + correct = 0 + total = 0 + + for inputs, labels in test_loader: + inputs, labels = inputs.to(device), labels.to(device) + + if robust: + inputs_in = pgd_attack(model, inputs, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=True, clamp_min=-1.0, clamp_max=1.0) + else: + inputs_in = inputs + + with torch.no_grad(): + outputs = model(inputs_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(NUM_CLASSES)[y_true] + + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (PGD)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f} | Acc: {accuracy:.4f} | auROC: {auroc:.4f} | auPRC: {auprc:.4f}") + +# ---------------------------- +# Main +# ---------------------------- \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/AUDIOMNIST/HybridViT/model_standard_AUDIOMNIST_HybridViT.py b/manuscripts/Poison26/bin/train/AUDIOMNIST/HybridViT/model_standard_AUDIOMNIST_HybridViT.py new file mode 100644 index 0000000..1412a7b --- /dev/null +++ b/manuscripts/Poison26/bin/train/AUDIOMNIST/HybridViT/model_standard_AUDIOMNIST_HybridViT.py @@ -0,0 +1,242 @@ +import argparse +import os +import glob +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader, Dataset, Subset + +import librosa +import torchaudio.transforms as T +from sklearn.metrics import roc_auc_score, average_precision_score + +import hashlib +import csv +import random + +# --- UPDATED IMPORT --- +from HybridViT import HybridViT + +# ---------------------------- +# Constants +# ---------------------------- +SAMPLING_RATE = 16000 +NUM_CLASSES = 10 +MAX_AUDIO_LENGTH = 16000 + +# ---------------------------- +# Audio Preprocessing +# ---------------------------- +def normalize_audio(x): + max_val = np.max(np.abs(x)) + return x / max_val if max_val > 0 else x + +def pad_audio(audio, max_len=MAX_AUDIO_LENGTH): + return audio[:max_len] if len(audio) > max_len else np.pad(audio, (0, max_len - len(audio)), 'constant') + +# ---------------------------- +# Dataset & Wrapper +# ---------------------------- +class AudioMNISTBaseDataset(Dataset): + def __init__(self, data_path): + self.data = [] + self.labels = [] + + wav_files = glob.glob(os.path.join(data_path, '*', '*.wav')) + wav_files = sorted(wav_files, key=lambda x: hashlib.md5(x.encode()).hexdigest()) + self.wav_files = wav_files.copy() + + for audio_path in tqdm(wav_files, desc="Loading audio files"): + audio, _ = librosa.load(audio_path, sr=SAMPLING_RATE) + audio = normalize_audio(audio) + audio = pad_audio(audio) + label = int(os.path.basename(audio_path)[0]) + self.data.append(audio) + self.labels.append(label) + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + return self.data[idx], self.labels[idx] + +class AudioSubsetWrapper(Dataset): + def __init__(self, subset, augment=False): + self.subset = subset + self.augment = augment + + def __len__(self): + return len(self.subset) + + def apply_augmentation(self, x): + if random.random() < 0.5: + x = np.clip(x + np.random.randn(len(x)) * 0.005, -1.0, 1.0) + if random.random() < 0.5: + x = np.roll(x, np.random.randint(-200, 200)) + if random.random() < 0.5: + x = np.clip(x * np.random.uniform(0.8, 1.2), -1.0, 1.0) + return x + + def __getitem__(self, idx): + x, y = self.subset[idx] + if self.augment: + x = self.apply_augmentation(x) + x = torch.tensor(x, dtype=torch.float32).unsqueeze(0) + return x, y + +def load_data(data_path, batch_size, augment_train=False, split_tsv="split_indices_standard.tsv"): + base_dataset = AudioMNISTBaseDataset(data_path) + + train_size = int(0.8 * len(base_dataset)) + train_indices = list(range(0, train_size)) + test_indices = list(range(train_size, len(base_dataset))) + + train_dataset = AudioSubsetWrapper(Subset(base_dataset, train_indices), augment=augment_train) + test_dataset = AudioSubsetWrapper(Subset(base_dataset, test_indices), augment=False) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) + + with open(split_tsv, "w", newline="") as f: + writer = csv.writer(f, delimiter="\t") + writer.writerow(["index", "split", "label", "file_path"]) + for idx in train_indices: + writer.writerow([idx, "train", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + for idx in test_indices: + writer.writerow([idx, "test", base_dataset.labels[idx], base_dataset.wav_files[idx]]) + + return train_loader, test_loader + +# ---------------------------- +# Model Definition +# ---------------------------- +class AudioMNISTModel(nn.Module): + def __init__(self, backbone_class, num_classes=NUM_CLASSES): + super(AudioMNISTModel, self).__init__() + + # 1D Audio to 2D Spectrogram + self.spectrogram = T.MelSpectrogram( + sample_rate=SAMPLING_RATE, + n_fft=1024, + hop_length=512, + n_mels=64 + ) + + # 16000 samples / 512 hop_length ~= 32 time frames. Shape: (1, 64, 32) + dummy_batch = torch.zeros(1, 1, 64, 32) + self.backbone = backbone_class(one_batch=dummy_batch, num_classes=num_classes) + + def forward(self, x): + # x is raw audio: (Batch, 1, 16000) + x = self.spectrogram(x) + x = torch.log(x + 1e-9) # Log scaling for numerical stability + return self.backbone(x) + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for inputs, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + inputs, labels = inputs.to(device), labels.to(device) + + optimizer.zero_grad(set_to_none=True) + outputs = model(inputs) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for inputs, labels in test_loader: + inputs, labels = inputs.to(device), labels.to(device) + outputs = model(inputs) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(NUM_CLASSES)[y_true] + + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f} | Acc: {accuracy:.4f} | auROC: {auroc:.4f} | auPRC: {auprc:.4f}") + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="AudioMNIST Standard Training") + parser.add_argument("--data", type=str, default="./data/AudioMNIST", help="Path to dataset") + parser.add_argument("--output", type=str, default="audiomnist_hybridvit_standard.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=10) + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # --- UPDATED MODEL CALL --- + model = AudioMNISTModel(backbone_class=HybridViT, num_classes=NUM_CLASSES) + + train_loader, test_loader = load_data(args.data, args.batch_size, augment_train=False) + + train(model, train_loader, device, epochs=args.epochs) + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/CIFAR10/MobileNet/MobileNet.py b/manuscripts/Poison26/bin/train/CIFAR10/MobileNet/MobileNet.py new file mode 100644 index 0000000..21cc2ab --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/MobileNet/MobileNet.py @@ -0,0 +1,121 @@ +import torch +import torch.nn as nn +from collections import OrderedDict + +# ------------------------- +# Depthwise Separable Convolution Block for MobileNetV1 +# ------------------------- +class DepthwiseSeparableConv(nn.Module): + def __init__(self, in_channels, out_channels, stride=1): + super(DepthwiseSeparableConv, self).__init__() + + # Depthwise layer with BN and ReLU6 + self.depthwise = nn.Sequential( + nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=stride, padding=1, groups=in_channels, bias=False), + nn.BatchNorm2d(in_channels), + nn.ReLU6(inplace=True) + ) + + # Pointwise layer with BN and ReLU6 + self.pointwise = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU6(inplace=True) + ) + + def forward(self, x): + x = self.depthwise(x) + x = self.pointwise(x) + return x + +# ------------------------- +# MobileNetV1 (Model Definition) +# ------------------------- +class MobileNet(nn.Module): + def __init__(self, one_batch=None, num_classes=1000): + super(MobileNet, self).__init__() + + # Handle dynamic input sizes + if one_batch is not None: + _, in_channels, H, W = one_batch.shape + self.input_channels = in_channels + self.input_size = (in_channels, H, W) + else: + self.input_channels = 3 + self.input_size = (3, 224, 224) + + # ------------------------- + # Stem + # ------------------------- + self.stem = nn.Sequential(OrderedDict([ + ('conv1', nn.Conv2d(self.input_channels, 32, kernel_size=3, stride=2, padding=1, bias=False)), + ('bn1', nn.BatchNorm2d(32)), + ('relu1', nn.ReLU6(inplace=True)), + ])) + + # ------------------------- + # Full MobileNetV1 Architecture + # ------------------------- + layers = [ + DepthwiseSeparableConv(32, 64, stride=1), + DepthwiseSeparableConv(64, 128, stride=2), + DepthwiseSeparableConv(128, 128, stride=1), + DepthwiseSeparableConv(128, 256, stride=2), + DepthwiseSeparableConv(256, 256, stride=1), + DepthwiseSeparableConv(256, 512, stride=2) + ] + + # 5x repeating blocks of 512 channels + for _ in range(5): + layers.append(DepthwiseSeparableConv(512, 512, stride=1)) + + # Final expansion to 1024 channels + layers.extend([ + DepthwiseSeparableConv(512, 1024, stride=2), + DepthwiseSeparableConv(1024, 1024, stride=1) + ]) + + # Pack the layers into an nn.Sequential for cleaner forward pass + self.features = nn.Sequential(*layers) + + # ------------------------- + # Classifier Setup + # ------------------------- + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc_input_features = self._get_flattened_feature_size(one_batch) + self.fc = nn.Linear(self.fc_input_features, num_classes) + + # ------------------------- + # Compute FC feature size dynamically + # ------------------------- + def _get_flattened_feature_size(self, one_batch): + was_training = self.training + self.eval() + + with torch.no_grad(): + if one_batch is None: + dummy = torch.zeros(1, *self.input_size) + else: + _, C, H, W = one_batch.shape + dummy = torch.zeros(1, C, H, W) + + x = self.stem(dummy) + x = self.features(x) + x = self.pool(x) + out_features = x.view(1, -1).size(1) + + if was_training: + self.train() + + return out_features + + # ------------------------- + # Forward + # ------------------------- + def forward(self, x): + x = self.stem(x) + x = self.features(x) + x = self.pool(x) + x = torch.flatten(x, 1) + x = self.fc(x) + return x diff --git a/manuscripts/Poison26/bin/train/CIFAR10/MobileNet/model_aug_CIFAR10_MobileNet.py b/manuscripts/Poison26/bin/train/CIFAR10/MobileNet/model_aug_CIFAR10_MobileNet.py new file mode 100644 index 0000000..3ba696b --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/MobileNet/model_aug_CIFAR10_MobileNet.py @@ -0,0 +1,158 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + +# Import the MobileNet model +from MobileNet import MobileNet + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + # Data augmentation for training + train_transform = transforms.Compose([ + transforms.RandomCrop(32, padding=4), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), # Converts to [0,1] + ]) + + # No augmentation for testing + test_transform = transforms.Compose([ + transforms.ToTensor(), + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=train_transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=test_transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="cifar10 training code (PyTorch) with Augmentation") + parser.add_argument("--output", type=str, default="cifar10_model_aug.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + dummy_batch = train_loader.dataset[0][0].unsqueeze(0) # Get a single sample and add batch dimension + + # Initialize MobileNet dynamically for 1-channel, 28x28 inputs + model = MobileNet(one_batch=dummy_batch, num_classes=10) + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/CIFAR10/MobileNet/model_fgsm_CIFAR10_MobileNet.py b/manuscripts/Poison26/bin/train/CIFAR10/MobileNet/model_fgsm_CIFAR10_MobileNet.py new file mode 100644 index 0000000..59f431d --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/MobileNet/model_fgsm_CIFAR10_MobileNet.py @@ -0,0 +1,213 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + +# Import the MobileNet model +from MobileNet import MobileNet + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # Converts to [0,1] + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader + + +# ---------------------------- +# FGSM attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def fgsm_attack(model, x, y, eps=0.05, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using FGSM (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps is in the same scale as x (e.g., cifar10 ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + x.requires_grad_(True) + logits = model(x) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x, retain_graph=False, create_graph=False)[0] + x_adv = x + eps * torch.sign(grad.detach()) + + # project back to [clamp_min, clamp_max] + x_adv = x_adv.clamp(clamp_min, clamp_max) + return x_adv.detach() + +# ---------------------------- +# Training loop (with optional FGSM adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under FGSM) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (FGSM)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="cifar10 training code (PyTorch) with FGSM adversarial training") + parser.add_argument("--output", type=str, default="cifar10_model_fgsm.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable FGSM adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="FGSM epsilon (in [0,1] pixel scale)") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + dummy_batch = train_loader.dataset[0][0].unsqueeze(0) # Get a single sample and add batch dimension + + # Initialize MobileNet dynamically for 1-channel, 28x28 inputs + model = MobileNet(one_batch=dummy_batch, num_classes=10) + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/CIFAR10/MobileNet/model_pgd_CIFAR10_MobileNet.py b/manuscripts/Poison26/bin/train/CIFAR10/MobileNet/model_pgd_CIFAR10_MobileNet.py new file mode 100644 index 0000000..6dd0e4d --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/MobileNet/model_pgd_CIFAR10_MobileNet.py @@ -0,0 +1,229 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + +# Import the MobileNet model +from MobileNet import MobileNet + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # Converts to [0,1] + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader + + +# ---------------------------- +# PGD attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def pgd_attack(model, x, y, eps=0.05, alpha=0.01, steps=40, random_start=True, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using PGD (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps/alpha are in the same scale as x (e.g., cifar10 ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + # start from a random point in the epsilon-ball if desired + if random_start: + x_adv = x + torch.empty_like(x).uniform_(-eps, eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + else: + x_adv = x.clone() + + for _ in range(steps): + x_adv.requires_grad_(True) + logits = model(x_adv) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x_adv, retain_graph=False, create_graph=False)[0] + x_adv = x_adv.detach() + alpha * torch.sign(grad.detach()) + + # project back to the epsilon l_inf ball around x, then clip to image bounds + x_adv = torch.max(torch.min(x_adv, x + eps), x - eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + + return x_adv.detach() + + +# ---------------------------- +# Training loop (with optional PGD adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, alpha=0.01, pgd_steps=40, random_start=True, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=random_start, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under PGD) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05, alpha=0.01, pgd_steps=40): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=True, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (PGD)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="cifar10 training code (PyTorch) with PGD adversarial training") + parser.add_argument("--output", type=str, default="cifar10_model_pgd.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable PGD adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="PGD epsilon (in [0,1] pixel scale)") + parser.add_argument("--alpha", type=float, default=0.01, help="PGD step size (in [0,1] scale)") + parser.add_argument("--pgd-steps", type=int, default=40, help="Number of PGD steps") + parser.add_argument("--no-random-start", action="store_true", help="Disable random PGD start") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + dummy_batch = train_loader.dataset[0][0].unsqueeze(0) # Get a single sample and add batch dimension + + # Initialize MobileNet dynamically for 1-channel, 28x28 inputs + model = MobileNet(one_batch=dummy_batch, num_classes=10) + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, alpha=args.alpha, + pgd_steps=args.pgd_steps, random_start=not args.no_random_start, + adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps, alpha=args.alpha, pgd_steps=args.pgd_steps) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/CIFAR10/MobileNet/model_standard_CIFAR10_MobileNet.py b/manuscripts/Poison26/bin/train/CIFAR10/MobileNet/model_standard_CIFAR10_MobileNet.py new file mode 100644 index 0000000..8d0613b --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/MobileNet/model_standard_CIFAR10_MobileNet.py @@ -0,0 +1,152 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + +# Import the MobileNet model +from MobileNet import MobileNet + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # Converts to [0,1] + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader + + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="cifar10 training code (PyTorch)") + parser.add_argument("--output", type=str, default="cifar10_model1.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + dummy_batch = train_loader.dataset[0][0].unsqueeze(0) # Get a single sample and add batch dimension + + # Initialize MobileNet dynamically for 1-channel, 28x28 inputs + model = MobileNet(one_batch=dummy_batch, num_classes=10) + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/CIFAR10/RegNetX/RegNetX.py b/manuscripts/Poison26/bin/train/CIFAR10/RegNetX/RegNetX.py new file mode 100644 index 0000000..29ce32b --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/RegNetX/RegNetX.py @@ -0,0 +1,140 @@ +import torch +import torch.nn as nn +from collections import OrderedDict + +# ------------------------- +# XBlock for RegNetX +# ------------------------- +class XBlock(nn.Module): + def __init__(self, in_channels, out_channels, stride=1, group_width=16): + super(XBlock, self).__init__() + + # Calculate the number of groups for the ResNeXt-style 3x3 convolution. + # RegNetX fixes the bottleneck ratio to 1, so inner channels = out_channels. + groups = out_channels // group_width + + self.block = nn.Sequential(OrderedDict([ + # 1x1 projection + ('conv1', nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)), + ('bn1', nn.BatchNorm2d(out_channels)), + ('relu1', nn.ReLU(inplace=True)), + + # 3x3 Group Convolution (CRITICAL FIX) + ('conv2', nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, + padding=1, groups=groups, bias=False)), + ('bn2', nn.BatchNorm2d(out_channels)), + ('relu2', nn.ReLU(inplace=True)), + + # 1x1 expansion + ('conv3', nn.Conv2d(out_channels, out_channels, kernel_size=1, bias=False)), + ('bn3', nn.BatchNorm2d(out_channels)) + ])) + + self.shortcut = nn.Sequential() + if stride != 1 or in_channels != out_channels: + self.shortcut = nn.Sequential(OrderedDict([ + ('shortcut_conv', nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False)), + ('shortcut_bn', nn.BatchNorm2d(out_channels)) + ])) + + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + return self.relu(self.block(x) + self.shortcut(x)) + + +# ------------------------- +# RegNetX-400MF +# ------------------------- +class RegNetX_400MF(nn.Module): + def __init__(self, one_batch=None, num_classes=200): + super(RegNetX_400MF, self).__init__() + + # Handle dynamic input sizes + if one_batch is not None: + _, in_channels, H, W = one_batch.shape + self.input_channels = in_channels + self.input_size = (in_channels, H, W) + else: + self.input_channels = 3 + self.input_size = (3, 224, 224) + + # ------------------------- + # Stem + # ------------------------- + self.stem = nn.Sequential(OrderedDict([ + ('stem_conv', nn.Conv2d(self.input_channels, 32, kernel_size=3, stride=2, padding=1, bias=False)), + ('stem_bn', nn.BatchNorm2d(32)), + ('stem_relu', nn.ReLU(inplace=True)) + ])) + + # ------------------------- + # Stages (Fixed to exact 400MF specs) + # Depths: [1, 2, 7, 12] + # Widths: [32, 64, 160, 384] + # ------------------------- + self.stage1 = self._make_stage('stage1', in_channels=32, out_channels=32, num_blocks=1, stride=2) + self.stage2 = self._make_stage('stage2', in_channels=32, out_channels=64, num_blocks=2, stride=2) + self.stage3 = self._make_stage('stage3', in_channels=64, out_channels=160, num_blocks=7, stride=2) + self.stage4 = self._make_stage('stage4', in_channels=160, out_channels=384, num_blocks=12, stride=2) + + # ------------------------- + # Classifier Setup + # ------------------------- + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc_input_features = self._get_flattened_feature_size(one_batch) + self.fc = nn.Linear(self.fc_input_features, num_classes) + + def _make_stage(self, stage_name, in_channels, out_channels, num_blocks, stride): + blocks = [] + for i in range(num_blocks): + # Only the first block in the stage handles the downsampling + s = stride if i == 0 else 1 + block_name = f"{stage_name}_block{i}" + # group_width=16 is constant across all stages for RegNetX-400MF + blocks.append((block_name, XBlock(in_channels, out_channels, stride=s, group_width=16))) + in_channels = out_channels + return nn.Sequential(OrderedDict(blocks)) + + # ------------------------- + # Compute FC feature size dynamically + # ------------------------- + def _get_flattened_feature_size(self, one_batch): + # We explicitly set eval() to prevent tracking BatchNorm stats with the dummy batch! + was_training = self.training + self.eval() + + with torch.no_grad(): + if one_batch is None: + dummy_input = torch.zeros(1, *self.input_size) + else: + _, C, H, W = one_batch.shape + dummy_input = torch.zeros(1, C, H, W) + + x = self.stem(dummy_input) + x = self.stage1(x) + x = self.stage2(x) + x = self.stage3(x) + x = self.stage4(x) + x = self.pool(x) + out_features = x.view(1, -1).size(1) + + if was_training: + self.train() + + return out_features + + # ------------------------- + # Forward + # ------------------------- + def forward(self, x): + x = self.stem(x) + x = self.stage1(x) + x = self.stage2(x) + x = self.stage3(x) + x = self.stage4(x) + + x = self.pool(x) + x = torch.flatten(x, 1) + x = self.fc(x) + return x diff --git a/manuscripts/Poison26/bin/train/CIFAR10/RegNetX/model_aug_CIFAR10_RegNetX.py b/manuscripts/Poison26/bin/train/CIFAR10/RegNetX/model_aug_CIFAR10_RegNetX.py new file mode 100644 index 0000000..6a61688 --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/RegNetX/model_aug_CIFAR10_RegNetX.py @@ -0,0 +1,159 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + +# Import the RegNetX model +from RegNetX import RegNetX_400MF + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + # Data augmentation for training + train_transform = transforms.Compose([ + transforms.RandomCrop(32, padding=4), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), # Converts to [0,1] + ]) + + # No augmentation for testing + test_transform = transforms.Compose([ + transforms.ToTensor(), + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=train_transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=test_transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="cifar10 training code (PyTorch) with Augmentation") + parser.add_argument("--output", type=str, default="cifar10_regnet_aug.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + dummy_batch = train_loader.dataset[0][0].unsqueeze(0) # Get a single sample and add batch dimension + # Initialize RegNetX dynamically for 1-channel, 28x28 inputs and 10 classes + model = RegNetX_400MF(one_batch=dummy_batch, num_classes=10) + + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/CIFAR10/RegNetX/model_fgsm_CIFAR10_RegNetX.py b/manuscripts/Poison26/bin/train/CIFAR10/RegNetX/model_fgsm_CIFAR10_RegNetX.py new file mode 100644 index 0000000..c844a14 --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/RegNetX/model_fgsm_CIFAR10_RegNetX.py @@ -0,0 +1,214 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + +# Import the RegNetX model +from RegNetX import RegNetX_400MF + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # Converts to [0,1] + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader + + +# ---------------------------- +# FGSM attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def fgsm_attack(model, x, y, eps=0.05, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using FGSM (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps is in the same scale as x (e.g., cifar10 ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + x.requires_grad_(True) + logits = model(x) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x, retain_graph=False, create_graph=False)[0] + x_adv = x + eps * torch.sign(grad.detach()) + + # project back to [clamp_min, clamp_max] + x_adv = x_adv.clamp(clamp_min, clamp_max) + return x_adv.detach() + +# ---------------------------- +# Training loop (with optional FGSM adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under FGSM) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (FGSM)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="cifar10 training code (PyTorch) with FGSM adversarial training") + parser.add_argument("--output", type=str, default="cifar10_regnet_fgsm.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable FGSM adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="FGSM epsilon (in [0,1] pixel scale)") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + dummy_batch = train_loader.dataset[0][0].unsqueeze(0) # Get a single sample and add batch dimension + # Initialize RegNetX dynamically for 1-channel, 28x28 inputs and 10 classes + model = RegNetX_400MF(one_batch=dummy_batch, num_classes=10) + + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/CIFAR10/RegNetX/model_pgd_CIFAR10_RegNetX.py b/manuscripts/Poison26/bin/train/CIFAR10/RegNetX/model_pgd_CIFAR10_RegNetX.py new file mode 100644 index 0000000..741fe22 --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/RegNetX/model_pgd_CIFAR10_RegNetX.py @@ -0,0 +1,230 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + +# Import the RegNetX model +from RegNetX import RegNetX_400MF + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # Converts to [0,1] + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader + + +# ---------------------------- +# PGD attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def pgd_attack(model, x, y, eps=0.05, alpha=0.01, steps=40, random_start=True, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using PGD (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps/alpha are in the same scale as x (e.g., cifar10 ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + # start from a random point in the epsilon-ball if desired + if random_start: + x_adv = x + torch.empty_like(x).uniform_(-eps, eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + else: + x_adv = x.clone() + + for _ in range(steps): + x_adv.requires_grad_(True) + logits = model(x_adv) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x_adv, retain_graph=False, create_graph=False)[0] + x_adv = x_adv.detach() + alpha * torch.sign(grad.detach()) + + # project back to the epsilon l_inf ball around x, then clip to image bounds + x_adv = torch.max(torch.min(x_adv, x + eps), x - eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + + return x_adv.detach() + + +# ---------------------------- +# Training loop (with optional PGD adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, alpha=0.01, pgd_steps=40, random_start=True, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=random_start, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under PGD) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05, alpha=0.01, pgd_steps=40): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=True, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (PGD)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="cifar10 training code (PyTorch) with PGD adversarial training") + parser.add_argument("--output", type=str, default="cifar10_regnet_pgd.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable PGD adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="PGD epsilon (in [0,1] pixel scale)") + parser.add_argument("--alpha", type=float, default=0.01, help="PGD step size (in [0,1] scale)") + parser.add_argument("--pgd-steps", type=int, default=40, help="Number of PGD steps") + parser.add_argument("--no-random-start", action="store_true", help="Disable random PGD start") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + dummy_batch = train_loader.dataset[0][0].unsqueeze(0) # Get a single sample and add batch dimension + # Initialize RegNetX dynamically for 1-channel, 28x28 inputs and 10 classes + model = RegNetX_400MF(one_batch=dummy_batch, num_classes=10) + + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, alpha=args.alpha, + pgd_steps=args.pgd_steps, random_start=not args.no_random_start, + adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps, alpha=args.alpha, pgd_steps=args.pgd_steps) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/CIFAR10/RegNetX/model_standard_CIFAR10_RegNetX.py b/manuscripts/Poison26/bin/train/CIFAR10/RegNetX/model_standard_CIFAR10_RegNetX.py new file mode 100644 index 0000000..b4b2b06 --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/RegNetX/model_standard_CIFAR10_RegNetX.py @@ -0,0 +1,153 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + +# Import the RegNetX model +from RegNetX import RegNetX_400MF + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # Converts to [0,1] + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader + + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="cifar10 training code (PyTorch)") + parser.add_argument("--output", type=str, default="cifar10_regnet_standard.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + dummy_batch = train_loader.dataset[0][0].unsqueeze(0) # Get a single sample and add batch dimension + # Initialize RegNetX dynamically for 1-channel, 28x28 inputs and 10 classes + model = RegNetX_400MF(one_batch=dummy_batch, num_classes=10) + + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/CIFAR10/adv/model_aug_CIFAR10_adv.py b/manuscripts/Poison26/bin/train/CIFAR10/adv/model_aug_CIFAR10_adv.py new file mode 100644 index 0000000..cf41c40 --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/adv/model_aug_CIFAR10_adv.py @@ -0,0 +1,222 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class CIFARModel(nn.Module): + def __init__(self, num_classes=10): + super(CIFARModel, self).__init__() + + # Block 1 + self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=0) # (32x32 → 30x30) + self.bn1 = nn.BatchNorm2d(32) + self.conv2 = nn.Conv2d(32, 32, kernel_size=3, padding=0) # (30x30 → 28x28) + self.bn2 = nn.BatchNorm2d(32) + self.pool1 = nn.MaxPool2d(2, 2) # (28x28 → 14x14) + self.drop1 = nn.Dropout(0.25) + + # Block 2 + self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1) # same padding + self.bn3 = nn.BatchNorm2d(64) + self.conv4 = nn.Conv2d(64, 64, kernel_size=3, padding=0) # (14x14 → 12x12) + self.bn4 = nn.BatchNorm2d(64) + self.pool2 = nn.MaxPool2d(2, 2) # (12x12 → 6x6) + self.drop2 = nn.Dropout(0.25) + + # Block 3 + self.conv5 = nn.Conv2d(64, 128, kernel_size=3, padding=1) # same padding + self.bn5 = nn.BatchNorm2d(128) + self.conv6 = nn.Conv2d(128, 128, kernel_size=3, padding=0) # (6x6 → 4x4) + self.bn6 = nn.BatchNorm2d(128) + self.pool3 = nn.MaxPool2d(2, 2) # (4x4 → 2x2) + self.drop3 = nn.Dropout(0.25) + + # Fully connected layers + self.fc1 = nn.Linear(128 * 2 * 2, 128) # flatten from 128 channels × 2×2 + self.drop4 = nn.Dropout(0.25) + self.fc2 = nn.Linear(128, num_classes) + + def forward(self, x): + # Block 1 + x = F.relu(self.bn1(self.conv1(x))) + x = F.relu(self.bn2(self.conv2(x))) + x = self.pool1(x) + x = self.drop1(x) + + # Block 2 + x = F.relu(self.bn3(self.conv3(x))) + x = F.relu(self.bn4(self.conv4(x))) + x = self.pool2(x) + x = self.drop2(x) + + # Block 3 + x = F.relu(self.bn5(self.conv5(x))) + x = F.relu(self.bn6(self.conv6(x))) + x = self.pool3(x) + x = self.drop3(x) + + # Flatten + x = x.view(x.size(0), -1) + + # Fully connected + x = F.relu(self.fc1(x)) + x = self.drop4(x) + x = self.fc2(x) + + return x # logits (use CrossEntropyLoss) + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + # Data augmentation for training + train_transform = transforms.Compose([ + transforms.RandomCrop(32, padding=4), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), # Converts to [0,1] + ]) + + # No augmentation for testing + test_transform = transforms.Compose([ + transforms.ToTensor(), + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=train_transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=test_transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="CIFAR10 training code (PyTorch)") + parser.add_argument("--output", type=str, default="cifar10_model3.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = CIFARModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/bin/train/CIFAR10/adv/model_fgsm_CIFAR10_adv.py b/manuscripts/Poison26/bin/train/CIFAR10/adv/model_fgsm_CIFAR10_adv.py new file mode 100644 index 0000000..da421b7 --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/adv/model_fgsm_CIFAR10_adv.py @@ -0,0 +1,276 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class CIFARModel(nn.Module): + def __init__(self, num_classes=10): + super(CIFARModel, self).__init__() + + # Block 1 + self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=0) # (32x32 → 30x30) + self.bn1 = nn.BatchNorm2d(32) + self.conv2 = nn.Conv2d(32, 32, kernel_size=3, padding=0) # (30x30 → 28x28) + self.bn2 = nn.BatchNorm2d(32) + self.pool1 = nn.MaxPool2d(2, 2) # (28x28 → 14x14) + self.drop1 = nn.Dropout(0.25) + + # Block 2 + self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1) # same padding + self.bn3 = nn.BatchNorm2d(64) + self.conv4 = nn.Conv2d(64, 64, kernel_size=3, padding=0) # (14x14 → 12x12) + self.bn4 = nn.BatchNorm2d(64) + self.pool2 = nn.MaxPool2d(2, 2) # (12x12 → 6x6) + self.drop2 = nn.Dropout(0.25) + + # Block 3 + self.conv5 = nn.Conv2d(64, 128, kernel_size=3, padding=1) # same padding + self.bn5 = nn.BatchNorm2d(128) + self.conv6 = nn.Conv2d(128, 128, kernel_size=3, padding=0) # (6x6 → 4x4) + self.bn6 = nn.BatchNorm2d(128) + self.pool3 = nn.MaxPool2d(2, 2) # (4x4 → 2x2) + self.drop3 = nn.Dropout(0.25) + + # Fully connected layers + self.fc1 = nn.Linear(128 * 2 * 2, 128) # flatten from 128 channels × 2×2 + self.drop4 = nn.Dropout(0.25) + self.fc2 = nn.Linear(128, num_classes) + + def forward(self, x): + # Block 1 + x = F.relu(self.bn1(self.conv1(x))) + x = F.relu(self.bn2(self.conv2(x))) + x = self.pool1(x) + x = self.drop1(x) + + # Block 2 + x = F.relu(self.bn3(self.conv3(x))) + x = F.relu(self.bn4(self.conv4(x))) + x = self.pool2(x) + x = self.drop2(x) + + # Block 3 + x = F.relu(self.bn5(self.conv5(x))) + x = F.relu(self.bn6(self.conv6(x))) + x = self.pool3(x) + x = self.drop3(x) + + # Flatten + x = x.view(x.size(0), -1) + + # Fully connected + x = F.relu(self.fc1(x)) + x = self.drop4(x) + x = self.fc2(x) + + return x # logits (use CrossEntropyLoss) + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # Converts to [0,1] + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader + + +# ---------------------------- +# FGSM attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def fgsm_attack(model, x, y, eps=0.05, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using FGSM (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps is in the same scale as x (e.g., cifar10 ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + x.requires_grad_(True) + logits = model(x) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x, retain_graph=False, create_graph=False)[0] + x_adv = x + eps * torch.sign(grad.detach()) + + # project back to [clamp_min, clamp_max] + x_adv = x_adv.clamp(clamp_min, clamp_max) + return x_adv.detach() + +# ---------------------------- +# Training loop (with optional FGSM adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under FGSM) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (FGSM)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="CIFAR10 training code (PyTorch) with FGSM adversarial training") + parser.add_argument("--output", type=str, default="cifar10_model4.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable FGSM adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="FGSM epsilon (in [0,1] pixel scale)") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = CIFARModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps) + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/bin/train/CIFAR10/adv/model_pgd_CIFAR10_adv.py b/manuscripts/Poison26/bin/train/CIFAR10/adv/model_pgd_CIFAR10_adv.py new file mode 100644 index 0000000..ed1b56b --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/adv/model_pgd_CIFAR10_adv.py @@ -0,0 +1,293 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class CIFARModel(nn.Module): + def __init__(self, num_classes=10): + super(CIFARModel, self).__init__() + + # Block 1 + self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=0) # (32x32 → 30x30) + self.bn1 = nn.BatchNorm2d(32) + self.conv2 = nn.Conv2d(32, 32, kernel_size=3, padding=0) # (30x30 → 28x28) + self.bn2 = nn.BatchNorm2d(32) + self.pool1 = nn.MaxPool2d(2, 2) # (28x28 → 14x14) + self.drop1 = nn.Dropout(0.25) + + # Block 2 + self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1) # same padding + self.bn3 = nn.BatchNorm2d(64) + self.conv4 = nn.Conv2d(64, 64, kernel_size=3, padding=0) # (14x14 → 12x12) + self.bn4 = nn.BatchNorm2d(64) + self.pool2 = nn.MaxPool2d(2, 2) # (12x12 → 6x6) + self.drop2 = nn.Dropout(0.25) + + # Block 3 + self.conv5 = nn.Conv2d(64, 128, kernel_size=3, padding=1) # same padding + self.bn5 = nn.BatchNorm2d(128) + self.conv6 = nn.Conv2d(128, 128, kernel_size=3, padding=0) # (6x6 → 4x4) + self.bn6 = nn.BatchNorm2d(128) + self.pool3 = nn.MaxPool2d(2, 2) # (4x4 → 2x2) + self.drop3 = nn.Dropout(0.25) + + # Fully connected layers + self.fc1 = nn.Linear(128 * 2 * 2, 128) # flatten from 128 channels × 2×2 + self.drop4 = nn.Dropout(0.25) + self.fc2 = nn.Linear(128, num_classes) + + def forward(self, x): + # Block 1 + x = F.relu(self.bn1(self.conv1(x))) + x = F.relu(self.bn2(self.conv2(x))) + x = self.pool1(x) + x = self.drop1(x) + + # Block 2 + x = F.relu(self.bn3(self.conv3(x))) + x = F.relu(self.bn4(self.conv4(x))) + x = self.pool2(x) + x = self.drop2(x) + + # Block 3 + x = F.relu(self.bn5(self.conv5(x))) + x = F.relu(self.bn6(self.conv6(x))) + x = self.pool3(x) + x = self.drop3(x) + + # Flatten + x = x.view(x.size(0), -1) + + # Fully connected + x = F.relu(self.fc1(x)) + x = self.drop4(x) + x = self.fc2(x) + + return x # logits (use CrossEntropyLoss) + + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # Converts to [0,1] + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader + + +# ---------------------------- +# PGD attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def pgd_attack(model, x, y, eps=0.05, alpha=0.01, steps=40, random_start=True, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using PGD (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps/alpha are in the same scale as x (e.g., cifar10 ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + # start from a random point in the epsilon-ball if desired + if random_start: + x_adv = x + torch.empty_like(x).uniform_(-eps, eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + else: + x_adv = x.clone() + + for _ in range(steps): + x_adv.requires_grad_(True) + logits = model(x_adv) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x_adv, retain_graph=False, create_graph=False)[0] + x_adv = x_adv.detach() + alpha * torch.sign(grad.detach()) + + # project back to the epsilon l_inf ball around x, then clip to image bounds + x_adv = torch.max(torch.min(x_adv, x + eps), x - eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + + return x_adv.detach() + + +# ---------------------------- +# Training loop (with optional PGD adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, alpha=0.01, pgd_steps=40, random_start=True, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=random_start, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under PGD) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05, alpha=0.01, pgd_steps=40): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=True, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (PGD)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="cifar10 training code (PyTorch) with PGD adversarial training") + parser.add_argument("--output", type=str, default="cifar10_model5.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable PGD adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="PGD epsilon (in [0,1] pixel scale)") + parser.add_argument("--alpha", type=float, default=0.01, help="PGD step size (in [0,1] scale)") + parser.add_argument("--pgd-steps", type=int, default=40, help="Number of PGD steps") + parser.add_argument("--no-random-start", action="store_true", help="Disable random PGD start") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = CIFARModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, alpha=args.alpha, + pgd_steps=args.pgd_steps, random_start=not args.no_random_start, + adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps, alpha=args.alpha, pgd_steps=args.pgd_steps) + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/bin/train/CIFAR10/adv/model_standard_CIFAR10_adv.py b/manuscripts/Poison26/bin/train/CIFAR10/adv/model_standard_CIFAR10_adv.py new file mode 100644 index 0000000..87d05b0 --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/adv/model_standard_CIFAR10_adv.py @@ -0,0 +1,215 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class CIFARModel(nn.Module): + def __init__(self, num_classes=10): + super(CIFARModel, self).__init__() + + # Block 1 + self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=0) # (32x32 → 30x30) + self.bn1 = nn.BatchNorm2d(32) + self.conv2 = nn.Conv2d(32, 32, kernel_size=3, padding=0) # (30x30 → 28x28) + self.bn2 = nn.BatchNorm2d(32) + self.pool1 = nn.MaxPool2d(2, 2) # (28x28 → 14x14) + self.drop1 = nn.Dropout(0.25) + + # Block 2 + self.conv3 = nn.Conv2d(32, 64, kernel_size=3, padding=1) # same padding + self.bn3 = nn.BatchNorm2d(64) + self.conv4 = nn.Conv2d(64, 64, kernel_size=3, padding=0) # (14x14 → 12x12) + self.bn4 = nn.BatchNorm2d(64) + self.pool2 = nn.MaxPool2d(2, 2) # (12x12 → 6x6) + self.drop2 = nn.Dropout(0.25) + + # Block 3 + self.conv5 = nn.Conv2d(64, 128, kernel_size=3, padding=1) # same padding + self.bn5 = nn.BatchNorm2d(128) + self.conv6 = nn.Conv2d(128, 128, kernel_size=3, padding=0) # (6x6 → 4x4) + self.bn6 = nn.BatchNorm2d(128) + self.pool3 = nn.MaxPool2d(2, 2) # (4x4 → 2x2) + self.drop3 = nn.Dropout(0.25) + + # Fully connected layers + self.fc1 = nn.Linear(128 * 2 * 2, 128) # flatten from 128 channels × 2×2 + self.drop4 = nn.Dropout(0.25) + self.fc2 = nn.Linear(128, num_classes) + + def forward(self, x): + # Block 1 + x = F.relu(self.bn1(self.conv1(x))) + x = F.relu(self.bn2(self.conv2(x))) + x = self.pool1(x) + x = self.drop1(x) + + # Block 2 + x = F.relu(self.bn3(self.conv3(x))) + x = F.relu(self.bn4(self.conv4(x))) + x = self.pool2(x) + x = self.drop2(x) + + # Block 3 + x = F.relu(self.bn5(self.conv5(x))) + x = F.relu(self.bn6(self.conv6(x))) + x = self.pool3(x) + x = self.drop3(x) + + # Flatten + x = x.view(x.size(0), -1) + + # Fully connected + x = F.relu(self.fc1(x)) + x = self.drop4(x) + x = self.fc2(x) + + return x # logits (use CrossEntropyLoss) + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # Converts to [0,1] + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader + + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="cifar10 training code (PyTorch)") + parser.add_argument("--output", type=str, default="cifar10_model2.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = CIFARModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/bin/train/CIFAR10/basic/model_aug_CIFAR10_basic.py b/manuscripts/Poison26/bin/train/CIFAR10/basic/model_aug_CIFAR10_basic.py new file mode 100644 index 0000000..b8946b1 --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/basic/model_aug_CIFAR10_basic.py @@ -0,0 +1,180 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class CIFARModel(nn.Module): + def __init__(self): + super(CIFARModel, self).__init__() + self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) # 3-channel input + self.pool1 = nn.MaxPool2d(2, 2) # 32→16 + self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) + self.pool2 = nn.MaxPool2d(2, 2) # 16→8 + self.fc1 = nn.Linear(64 * 8 * 8, 128) # flatten size = 64*8*8 + self.fc2 = nn.Linear(128, 10) + self.relu = nn.ReLU() + + def forward(self, x): + x = self.relu(self.conv1(x)) + x = self.pool1(x) + x = self.relu(self.conv2(x)) + x = self.pool2(x) + x = x.view(x.size(0), -1) # flatten + x = self.relu(self.fc1(x)) + x = self.fc2(x) # raw logits + return x + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + # Data augmentation for training + train_transform = transforms.Compose([ + transforms.RandomCrop(32, padding=4), + transforms.RandomHorizontalFlip(), + transforms.ToTensor(), # Converts to [0,1] + ]) + + # No augmentation for testing + test_transform = transforms.Compose([ + transforms.ToTensor(), + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=train_transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=test_transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="cifar10 training code (PyTorch)") + parser.add_argument("--output", type=str, default="cifar10_model1.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = CIFARModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/bin/train/CIFAR10/basic/model_fgsm_CIFAR10_basic.py b/manuscripts/Poison26/bin/train/CIFAR10/basic/model_fgsm_CIFAR10_basic.py new file mode 100644 index 0000000..7acf432 --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/basic/model_fgsm_CIFAR10_basic.py @@ -0,0 +1,234 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class CIFARModel(nn.Module): + def __init__(self): + super(CIFARModel, self).__init__() + self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) # 3-channel input + self.pool1 = nn.MaxPool2d(2, 2) # 32→16 + self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) + self.pool2 = nn.MaxPool2d(2, 2) # 16→8 + self.fc1 = nn.Linear(64 * 8 * 8, 128) # flatten size = 64*8*8 + self.fc2 = nn.Linear(128, 10) + self.relu = nn.ReLU() + + def forward(self, x): + x = self.relu(self.conv1(x)) + x = self.pool1(x) + x = self.relu(self.conv2(x)) + x = self.pool2(x) + x = x.view(x.size(0), -1) # flatten + x = self.relu(self.fc1(x)) + x = self.fc2(x) # raw logits + return x + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # Converts to [0,1] + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader + + +# ---------------------------- +# FGSM attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def fgsm_attack(model, x, y, eps=0.05, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using FGSM (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps is in the same scale as x (e.g., cifar10 ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + x.requires_grad_(True) + logits = model(x) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x, retain_graph=False, create_graph=False)[0] + x_adv = x + eps * torch.sign(grad.detach()) + + # project back to [clamp_min, clamp_max] + x_adv = x_adv.clamp(clamp_min, clamp_max) + return x_adv.detach() + +# ---------------------------- +# Training loop (with optional FGSM adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under FGSM) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (FGSM)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="cifar10 training code (PyTorch) with FGSM adversarial training") + parser.add_argument("--output", type=str, default="cifar10_model4.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable FGSM adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="FGSM epsilon (in [0,1] pixel scale)") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = CIFARModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps) + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/bin/train/CIFAR10/basic/model_pgd_CIFAR10_basic.py b/manuscripts/Poison26/bin/train/CIFAR10/basic/model_pgd_CIFAR10_basic.py new file mode 100644 index 0000000..dcc07bf --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/basic/model_pgd_CIFAR10_basic.py @@ -0,0 +1,250 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class CIFARModel(nn.Module): + def __init__(self): + super(CIFARModel, self).__init__() + self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) # 3-channel input + self.pool1 = nn.MaxPool2d(2, 2) # 32→16 + self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) + self.pool2 = nn.MaxPool2d(2, 2) # 16→8 + self.fc1 = nn.Linear(64 * 8 * 8, 128) # flatten size = 64*8*8 + self.fc2 = nn.Linear(128, 10) + self.relu = nn.ReLU() + + def forward(self, x): + x = self.relu(self.conv1(x)) + x = self.pool1(x) + x = self.relu(self.conv2(x)) + x = self.pool2(x) + x = x.view(x.size(0), -1) # flatten + x = self.relu(self.fc1(x)) + x = self.fc2(x) # raw logits + return x + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # Converts to [0,1] + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader + + +# ---------------------------- +# PGD attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def pgd_attack(model, x, y, eps=0.05, alpha=0.01, steps=40, random_start=True, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using PGD (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps/alpha are in the same scale as x (e.g., cifar10 ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + # start from a random point in the epsilon-ball if desired + if random_start: + x_adv = x + torch.empty_like(x).uniform_(-eps, eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + else: + x_adv = x.clone() + + for _ in range(steps): + x_adv.requires_grad_(True) + logits = model(x_adv) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x_adv, retain_graph=False, create_graph=False)[0] + x_adv = x_adv.detach() + alpha * torch.sign(grad.detach()) + + # project back to the epsilon l_inf ball around x, then clip to image bounds + x_adv = torch.max(torch.min(x_adv, x + eps), x - eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + + return x_adv.detach() + + +# ---------------------------- +# Training loop (with optional PGD adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, alpha=0.01, pgd_steps=40, random_start=True, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=random_start, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under PGD) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05, alpha=0.01, pgd_steps=40): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=True, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (PGD)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="cifar10 training code (PyTorch) with PGD adversarial training") + parser.add_argument("--output", type=str, default="cifar10_model5.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable PGD adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="PGD epsilon (in [0,1] pixel scale)") + parser.add_argument("--alpha", type=float, default=0.01, help="PGD step size (in [0,1] scale)") + parser.add_argument("--pgd-steps", type=int, default=40, help="Number of PGD steps") + parser.add_argument("--no-random-start", action="store_true", help="Disable random PGD start") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = CIFARModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, alpha=args.alpha, + pgd_steps=args.pgd_steps, random_start=not args.no_random_start, + adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps, alpha=args.alpha, pgd_steps=args.pgd_steps) + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/bin/train/CIFAR10/basic/model_standard_CIFAR10_basic.py b/manuscripts/Poison26/bin/train/CIFAR10/basic/model_standard_CIFAR10_basic.py new file mode 100644 index 0000000..ce839d0 --- /dev/null +++ b/manuscripts/Poison26/bin/train/CIFAR10/basic/model_standard_CIFAR10_basic.py @@ -0,0 +1,174 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class CIFARModel(nn.Module): + def __init__(self): + super(CIFARModel, self).__init__() + self.conv1 = nn.Conv2d(3, 32, kernel_size=3, padding=1) # 3-channel input + self.pool1 = nn.MaxPool2d(2, 2) # 32→16 + self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) + self.pool2 = nn.MaxPool2d(2, 2) # 16→8 + self.fc1 = nn.Linear(64 * 8 * 8, 128) # flatten size = 64*8*8 + self.fc2 = nn.Linear(128, 10) + self.relu = nn.ReLU() + + def forward(self, x): + x = self.relu(self.conv1(x)) + x = self.pool1(x) + x = self.relu(self.conv2(x)) + x = self.pool2(x) + x = x.view(x.size(0), -1) # flatten + x = self.relu(self.fc1(x)) + x = self.fc2(x) # raw logits + return x + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # Converts to [0,1] + ]) + + train_dataset = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2) + + return train_loader, test_loader + + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="cifar10 training code (PyTorch)") + parser.add_argument("--output", type=str, default="cifar10_model1.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = CIFARModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/bin/train/MNIST/MobileNet/MobileNet.py b/manuscripts/Poison26/bin/train/MNIST/MobileNet/MobileNet.py new file mode 100644 index 0000000..21cc2ab --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/MobileNet/MobileNet.py @@ -0,0 +1,121 @@ +import torch +import torch.nn as nn +from collections import OrderedDict + +# ------------------------- +# Depthwise Separable Convolution Block for MobileNetV1 +# ------------------------- +class DepthwiseSeparableConv(nn.Module): + def __init__(self, in_channels, out_channels, stride=1): + super(DepthwiseSeparableConv, self).__init__() + + # Depthwise layer with BN and ReLU6 + self.depthwise = nn.Sequential( + nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=stride, padding=1, groups=in_channels, bias=False), + nn.BatchNorm2d(in_channels), + nn.ReLU6(inplace=True) + ) + + # Pointwise layer with BN and ReLU6 + self.pointwise = nn.Sequential( + nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False), + nn.BatchNorm2d(out_channels), + nn.ReLU6(inplace=True) + ) + + def forward(self, x): + x = self.depthwise(x) + x = self.pointwise(x) + return x + +# ------------------------- +# MobileNetV1 (Model Definition) +# ------------------------- +class MobileNet(nn.Module): + def __init__(self, one_batch=None, num_classes=1000): + super(MobileNet, self).__init__() + + # Handle dynamic input sizes + if one_batch is not None: + _, in_channels, H, W = one_batch.shape + self.input_channels = in_channels + self.input_size = (in_channels, H, W) + else: + self.input_channels = 3 + self.input_size = (3, 224, 224) + + # ------------------------- + # Stem + # ------------------------- + self.stem = nn.Sequential(OrderedDict([ + ('conv1', nn.Conv2d(self.input_channels, 32, kernel_size=3, stride=2, padding=1, bias=False)), + ('bn1', nn.BatchNorm2d(32)), + ('relu1', nn.ReLU6(inplace=True)), + ])) + + # ------------------------- + # Full MobileNetV1 Architecture + # ------------------------- + layers = [ + DepthwiseSeparableConv(32, 64, stride=1), + DepthwiseSeparableConv(64, 128, stride=2), + DepthwiseSeparableConv(128, 128, stride=1), + DepthwiseSeparableConv(128, 256, stride=2), + DepthwiseSeparableConv(256, 256, stride=1), + DepthwiseSeparableConv(256, 512, stride=2) + ] + + # 5x repeating blocks of 512 channels + for _ in range(5): + layers.append(DepthwiseSeparableConv(512, 512, stride=1)) + + # Final expansion to 1024 channels + layers.extend([ + DepthwiseSeparableConv(512, 1024, stride=2), + DepthwiseSeparableConv(1024, 1024, stride=1) + ]) + + # Pack the layers into an nn.Sequential for cleaner forward pass + self.features = nn.Sequential(*layers) + + # ------------------------- + # Classifier Setup + # ------------------------- + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc_input_features = self._get_flattened_feature_size(one_batch) + self.fc = nn.Linear(self.fc_input_features, num_classes) + + # ------------------------- + # Compute FC feature size dynamically + # ------------------------- + def _get_flattened_feature_size(self, one_batch): + was_training = self.training + self.eval() + + with torch.no_grad(): + if one_batch is None: + dummy = torch.zeros(1, *self.input_size) + else: + _, C, H, W = one_batch.shape + dummy = torch.zeros(1, C, H, W) + + x = self.stem(dummy) + x = self.features(x) + x = self.pool(x) + out_features = x.view(1, -1).size(1) + + if was_training: + self.train() + + return out_features + + # ------------------------- + # Forward + # ------------------------- + def forward(self, x): + x = self.stem(x) + x = self.features(x) + x = self.pool(x) + x = torch.flatten(x, 1) + x = self.fc(x) + return x diff --git a/manuscripts/Poison26/bin/train/MNIST/MobileNet/model_aug_MNIST_MobileNet.py b/manuscripts/Poison26/bin/train/MNIST/MobileNet/model_aug_MNIST_MobileNet.py new file mode 100644 index 0000000..bb02368 --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/MobileNet/model_aug_MNIST_MobileNet.py @@ -0,0 +1,165 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + +# Import the MobileNet model +from MobileNet import MobileNet + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + # Define transforms for training (with augmentation) and test (no augmentation) + train_transform = transforms.Compose([ + transforms.RandomRotation(10), # rotation_range=10 + transforms.RandomAffine(0, translate=(0.1, 0.1)), # width/height shift + transforms.RandomResizedCrop(28, scale=(0.9, 1.1)), # zoom_range=0.10 + transforms.ToTensor(), # scale to [0,1] + ]) + + test_transform = transforms.Compose([ + transforms.ToTensor(), # only rescale + ]) + + # Download MNIST dataset + train_dataset = datasets.MNIST( + root="./data", train=True, download=True, transform=train_transform + ) + test_dataset = datasets.MNIST( + root="./data", train=False, download=True, transform=test_transform + ) + + # Create DataLoaders + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) + + return train_loader, test_loader + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch) with Augmentation") + parser.add_argument("--output", type=str, default="mnist_model_aug.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize MobileNet dynamically for 1-channel, 28x28 inputs + dummy_batch = torch.zeros(1, 1, 28, 28) + model = MobileNet(one_batch=dummy_batch, num_classes=10) + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/MNIST/MobileNet/model_fgsm_MNIST_MobileNet.py b/manuscripts/Poison26/bin/train/MNIST/MobileNet/model_fgsm_MNIST_MobileNet.py new file mode 100644 index 0000000..421ca10 --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/MobileNet/model_fgsm_MNIST_MobileNet.py @@ -0,0 +1,210 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + +# Import the MobileNet model +from MobileNet import MobileNet + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # converts to [0,1] + ]) + train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.MNIST(root="./data", train=False, download=True, transform=transform) + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) + return train_loader, test_loader + + +# ---------------------------- +# FGSM attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def fgsm_attack(model, x, y, eps=0.05, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using FGSM (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps is in the same scale as x (e.g., MNIST ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + x.requires_grad_(True) + logits = model(x) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x, retain_graph=False, create_graph=False)[0] + x_adv = x + eps * torch.sign(grad.detach()) + + # project back to [clamp_min, clamp_max] + x_adv = x_adv.clamp(clamp_min, clamp_max) + return x_adv.detach() + +# ---------------------------- +# Training loop (with optional FGSM adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under FGSM) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (FGSM)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch) with FGSM adversarial training") + parser.add_argument("--output", type=str, default="mnist_model_fgsm.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable FGSM adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="FGSM epsilon (in [0,1] pixel scale)") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize MobileNet dynamically for 1-channel, 28x28 inputs + dummy_batch = torch.zeros(1, 1, 28, 28) + model = MobileNet(one_batch=dummy_batch, num_classes=10) + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/MNIST/MobileNet/model_pgd_MNIST_MobileNet.py b/manuscripts/Poison26/bin/train/MNIST/MobileNet/model_pgd_MNIST_MobileNet.py new file mode 100644 index 0000000..43a439e --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/MobileNet/model_pgd_MNIST_MobileNet.py @@ -0,0 +1,226 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + +# Import the MobileNet model +from MobileNet import MobileNet + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # converts to [0,1] + ]) + train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.MNIST(root="./data", train=False, download=True, transform=transform) + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) + return train_loader, test_loader + + +# ---------------------------- +# PGD attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def pgd_attack(model, x, y, eps=0.05, alpha=0.01, steps=40, random_start=True, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using PGD (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps/alpha are in the same scale as x (e.g., MNIST ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + # start from a random point in the epsilon-ball if desired + if random_start: + x_adv = x + torch.empty_like(x).uniform_(-eps, eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + else: + x_adv = x.clone() + + for _ in range(steps): + x_adv.requires_grad_(True) + logits = model(x_adv) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x_adv, retain_graph=False, create_graph=False)[0] + x_adv = x_adv.detach() + alpha * torch.sign(grad.detach()) + + # project back to the epsilon l_inf ball around x, then clip to image bounds + x_adv = torch.max(torch.min(x_adv, x + eps), x - eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + + return x_adv.detach() + + +# ---------------------------- +# Training loop (with optional PGD adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, alpha=0.01, pgd_steps=40, random_start=True, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=random_start, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under PGD) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05, alpha=0.01, pgd_steps=40): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=True, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (PGD)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch) with PGD adversarial training") + parser.add_argument("--output", type=str, default="mnist_model_pgd.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable PGD adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="PGD epsilon (in [0,1] pixel scale)") + parser.add_argument("--alpha", type=float, default=0.01, help="PGD step size (in [0,1] scale)") + parser.add_argument("--pgd-steps", type=int, default=40, help="Number of PGD steps") + parser.add_argument("--no-random-start", action="store_true", help="Disable random PGD start") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize MobileNet dynamically for 1-channel, 28x28 inputs + dummy_batch = torch.zeros(1, 1, 28, 28) + model = MobileNet(one_batch=dummy_batch, num_classes=10) + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, alpha=args.alpha, + pgd_steps=args.pgd_steps, random_start=not args.no_random_start, + adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps, alpha=args.alpha, pgd_steps=args.pgd_steps) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/MNIST/MobileNet/model_standard_MNIST_MobileNet.py b/manuscripts/Poison26/bin/train/MNIST/MobileNet/model_standard_MNIST_MobileNet.py new file mode 100644 index 0000000..a643d91 --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/MobileNet/model_standard_MNIST_MobileNet.py @@ -0,0 +1,154 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + +# Import the MobileNet model +from MobileNet import MobileNet + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # converts to [0,1] + ]) + + train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.MNIST(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) + + return train_loader, test_loader + + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch)") + parser.add_argument("--output", type=str, default="mnist_model1.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + # Pass a dummy batch to configure the MobileNet stem for 1-channel MNIST images + # and properly calculate the fully-connected layer inputs for 28x28 resolution. + dummy_batch = torch.zeros(1, 1, 28, 28) + model = MobileNet(one_batch=dummy_batch, num_classes=10) + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/MNIST/RegNetX/RegNetX.py b/manuscripts/Poison26/bin/train/MNIST/RegNetX/RegNetX.py new file mode 100644 index 0000000..29ce32b --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/RegNetX/RegNetX.py @@ -0,0 +1,140 @@ +import torch +import torch.nn as nn +from collections import OrderedDict + +# ------------------------- +# XBlock for RegNetX +# ------------------------- +class XBlock(nn.Module): + def __init__(self, in_channels, out_channels, stride=1, group_width=16): + super(XBlock, self).__init__() + + # Calculate the number of groups for the ResNeXt-style 3x3 convolution. + # RegNetX fixes the bottleneck ratio to 1, so inner channels = out_channels. + groups = out_channels // group_width + + self.block = nn.Sequential(OrderedDict([ + # 1x1 projection + ('conv1', nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)), + ('bn1', nn.BatchNorm2d(out_channels)), + ('relu1', nn.ReLU(inplace=True)), + + # 3x3 Group Convolution (CRITICAL FIX) + ('conv2', nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=stride, + padding=1, groups=groups, bias=False)), + ('bn2', nn.BatchNorm2d(out_channels)), + ('relu2', nn.ReLU(inplace=True)), + + # 1x1 expansion + ('conv3', nn.Conv2d(out_channels, out_channels, kernel_size=1, bias=False)), + ('bn3', nn.BatchNorm2d(out_channels)) + ])) + + self.shortcut = nn.Sequential() + if stride != 1 or in_channels != out_channels: + self.shortcut = nn.Sequential(OrderedDict([ + ('shortcut_conv', nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False)), + ('shortcut_bn', nn.BatchNorm2d(out_channels)) + ])) + + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + return self.relu(self.block(x) + self.shortcut(x)) + + +# ------------------------- +# RegNetX-400MF +# ------------------------- +class RegNetX_400MF(nn.Module): + def __init__(self, one_batch=None, num_classes=200): + super(RegNetX_400MF, self).__init__() + + # Handle dynamic input sizes + if one_batch is not None: + _, in_channels, H, W = one_batch.shape + self.input_channels = in_channels + self.input_size = (in_channels, H, W) + else: + self.input_channels = 3 + self.input_size = (3, 224, 224) + + # ------------------------- + # Stem + # ------------------------- + self.stem = nn.Sequential(OrderedDict([ + ('stem_conv', nn.Conv2d(self.input_channels, 32, kernel_size=3, stride=2, padding=1, bias=False)), + ('stem_bn', nn.BatchNorm2d(32)), + ('stem_relu', nn.ReLU(inplace=True)) + ])) + + # ------------------------- + # Stages (Fixed to exact 400MF specs) + # Depths: [1, 2, 7, 12] + # Widths: [32, 64, 160, 384] + # ------------------------- + self.stage1 = self._make_stage('stage1', in_channels=32, out_channels=32, num_blocks=1, stride=2) + self.stage2 = self._make_stage('stage2', in_channels=32, out_channels=64, num_blocks=2, stride=2) + self.stage3 = self._make_stage('stage3', in_channels=64, out_channels=160, num_blocks=7, stride=2) + self.stage4 = self._make_stage('stage4', in_channels=160, out_channels=384, num_blocks=12, stride=2) + + # ------------------------- + # Classifier Setup + # ------------------------- + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc_input_features = self._get_flattened_feature_size(one_batch) + self.fc = nn.Linear(self.fc_input_features, num_classes) + + def _make_stage(self, stage_name, in_channels, out_channels, num_blocks, stride): + blocks = [] + for i in range(num_blocks): + # Only the first block in the stage handles the downsampling + s = stride if i == 0 else 1 + block_name = f"{stage_name}_block{i}" + # group_width=16 is constant across all stages for RegNetX-400MF + blocks.append((block_name, XBlock(in_channels, out_channels, stride=s, group_width=16))) + in_channels = out_channels + return nn.Sequential(OrderedDict(blocks)) + + # ------------------------- + # Compute FC feature size dynamically + # ------------------------- + def _get_flattened_feature_size(self, one_batch): + # We explicitly set eval() to prevent tracking BatchNorm stats with the dummy batch! + was_training = self.training + self.eval() + + with torch.no_grad(): + if one_batch is None: + dummy_input = torch.zeros(1, *self.input_size) + else: + _, C, H, W = one_batch.shape + dummy_input = torch.zeros(1, C, H, W) + + x = self.stem(dummy_input) + x = self.stage1(x) + x = self.stage2(x) + x = self.stage3(x) + x = self.stage4(x) + x = self.pool(x) + out_features = x.view(1, -1).size(1) + + if was_training: + self.train() + + return out_features + + # ------------------------- + # Forward + # ------------------------- + def forward(self, x): + x = self.stem(x) + x = self.stage1(x) + x = self.stage2(x) + x = self.stage3(x) + x = self.stage4(x) + + x = self.pool(x) + x = torch.flatten(x, 1) + x = self.fc(x) + return x diff --git a/manuscripts/Poison26/bin/train/MNIST/RegNetX/model_aug_MNIST_RegNetX.py b/manuscripts/Poison26/bin/train/MNIST/RegNetX/model_aug_MNIST_RegNetX.py new file mode 100644 index 0000000..6883d43 --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/RegNetX/model_aug_MNIST_RegNetX.py @@ -0,0 +1,165 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + +# Import the RegNetX model +from RegNetX import RegNetX_400MF + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + # Define transforms for training (with augmentation) and test (no augmentation) + train_transform = transforms.Compose([ + transforms.RandomRotation(10), # rotation_range=10 + transforms.RandomAffine(0, translate=(0.1, 0.1)), # width/height shift + transforms.RandomResizedCrop(28, scale=(0.9, 1.1)), # zoom_range=0.10 + transforms.ToTensor(), # scale to [0,1] + ]) + + test_transform = transforms.Compose([ + transforms.ToTensor(), # only rescale + ]) + + # Download MNIST dataset + train_dataset = datasets.MNIST( + root="./data", train=True, download=True, transform=train_transform + ) + test_dataset = datasets.MNIST( + root="./data", train=False, download=True, transform=test_transform + ) + + # Create DataLoaders + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) + + return train_loader, test_loader + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch) with Augmentation") + parser.add_argument("--output", type=str, default="mnist_regnet_aug.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize RegNetX dynamically for 1-channel, 28x28 inputs and 10 classes + dummy_batch = torch.zeros(1, 1, 28, 28) + model = RegNetX_400MF(one_batch=dummy_batch, num_classes=10) + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/MNIST/RegNetX/model_fgsm_MNIST_RegNetX.py b/manuscripts/Poison26/bin/train/MNIST/RegNetX/model_fgsm_MNIST_RegNetX.py new file mode 100644 index 0000000..9cf4d18 --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/RegNetX/model_fgsm_MNIST_RegNetX.py @@ -0,0 +1,210 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + +# Import the RegNetX model +from RegNetX import RegNetX_400MF + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # converts to [0,1] + ]) + train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.MNIST(root="./data", train=False, download=True, transform=transform) + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) + return train_loader, test_loader + + +# ---------------------------- +# FGSM attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def fgsm_attack(model, x, y, eps=0.05, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using FGSM (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps is in the same scale as x (e.g., MNIST ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + x.requires_grad_(True) + logits = model(x) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x, retain_graph=False, create_graph=False)[0] + x_adv = x + eps * torch.sign(grad.detach()) + + # project back to [clamp_min, clamp_max] + x_adv = x_adv.clamp(clamp_min, clamp_max) + return x_adv.detach() + +# ---------------------------- +# Training loop (with optional FGSM adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under FGSM) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (FGSM)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch) with FGSM adversarial training") + parser.add_argument("--output", type=str, default="mnist_regnet_fgsm.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable FGSM adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="FGSM epsilon (in [0,1] pixel scale)") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize RegNetX dynamically for 1-channel, 28x28 inputs and 10 classes + dummy_batch = torch.zeros(1, 1, 28, 28) + model = RegNetX_400MF(one_batch=dummy_batch, num_classes=10) + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/MNIST/RegNetX/model_pgd_MNIST_RegNetX.py b/manuscripts/Poison26/bin/train/MNIST/RegNetX/model_pgd_MNIST_RegNetX.py new file mode 100644 index 0000000..8ab68eb --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/RegNetX/model_pgd_MNIST_RegNetX.py @@ -0,0 +1,226 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + +# Import the RegNetX model +from RegNetX import RegNetX_400MF + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # converts to [0,1] + ]) + train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.MNIST(root="./data", train=False, download=True, transform=transform) + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) + return train_loader, test_loader + + +# ---------------------------- +# PGD attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def pgd_attack(model, x, y, eps=0.05, alpha=0.01, steps=40, random_start=True, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using PGD (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps/alpha are in the same scale as x (e.g., MNIST ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + # start from a random point in the epsilon-ball if desired + if random_start: + x_adv = x + torch.empty_like(x).uniform_(-eps, eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + else: + x_adv = x.clone() + + for _ in range(steps): + x_adv.requires_grad_(True) + logits = model(x_adv) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x_adv, retain_graph=False, create_graph=False)[0] + x_adv = x_adv.detach() + alpha * torch.sign(grad.detach()) + + # project back to the epsilon l_inf ball around x, then clip to image bounds + x_adv = torch.max(torch.min(x_adv, x + eps), x - eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + + return x_adv.detach() + + +# ---------------------------- +# Training loop (with optional PGD adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, alpha=0.01, pgd_steps=40, random_start=True, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=random_start, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under PGD) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05, alpha=0.01, pgd_steps=40): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=True, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (PGD)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch) with PGD adversarial training") + parser.add_argument("--output", type=str, default="mnist_regnet_pgd.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable PGD adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="PGD epsilon (in [0,1] pixel scale)") + parser.add_argument("--alpha", type=float, default=0.01, help="PGD step size (in [0,1] scale)") + parser.add_argument("--pgd-steps", type=int, default=40, help="Number of PGD steps") + parser.add_argument("--no-random-start", action="store_true", help="Disable random PGD start") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize RegNetX dynamically for 1-channel, 28x28 inputs and 10 classes + dummy_batch = torch.zeros(1, 1, 28, 28) + model = RegNetX_400MF(one_batch=dummy_batch, num_classes=10) + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, alpha=args.alpha, + pgd_steps=args.pgd_steps, random_start=not args.no_random_start, + adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps, alpha=args.alpha, pgd_steps=args.pgd_steps) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/MNIST/RegNetX/model_standard_MNIST_RegNetX.py b/manuscripts/Poison26/bin/train/MNIST/RegNetX/model_standard_MNIST_RegNetX.py new file mode 100644 index 0000000..d2e6a86 --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/RegNetX/model_standard_MNIST_RegNetX.py @@ -0,0 +1,152 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + +# Import the RegNetX model +from RegNetX import RegNetX_400MF + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # converts to [0,1] + ]) + + train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.MNIST(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) + + return train_loader, test_loader + + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch)") + parser.add_argument("--output", type=str, default="mnist_regnet_standard.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize RegNetX dynamically for 1-channel, 28x28 inputs and 10 classes + dummy_batch = torch.zeros(1, 1, 28, 28) + model = RegNetX_400MF(one_batch=dummy_batch, num_classes=10) + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/manuscripts/Poison26/bin/train/MNIST/adv/model_aug_MNIST_adv.py b/manuscripts/Poison26/bin/train/MNIST/adv/model_aug_MNIST_adv.py new file mode 100644 index 0000000..53e9b6b --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/adv/model_aug_MNIST_adv.py @@ -0,0 +1,209 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class MNISTModel(nn.Module): + def __init__(self): + super(MNISTModel, self).__init__() + self.conv1 = nn.Conv2d(1, 32, kernel_size=3) # 28x28 -> 26x26 + self.bn1 = nn.BatchNorm2d(32) + self.conv2 = nn.Conv2d(32, 32, kernel_size=3) # 26 -> 24 + self.bn2 = nn.BatchNorm2d(32) + self.conv3 = nn.Conv2d(32, 32, kernel_size=5, stride=2, padding=2) # 24 -> 12 + self.bn3 = nn.BatchNorm2d(32) + self.dropout1 = nn.Dropout(0.4) + + self.conv4 = nn.Conv2d(32, 64, kernel_size=3) # 12 -> 10 + self.bn4 = nn.BatchNorm2d(64) + self.conv5 = nn.Conv2d(64, 64, kernel_size=3) # 10 -> 8 + self.bn5 = nn.BatchNorm2d(64) + self.conv6 = nn.Conv2d(64, 64, kernel_size=5, stride=2, padding=2) # 8 -> 4 + self.bn6 = nn.BatchNorm2d(64) + self.dropout2 = nn.Dropout(0.4) + + self.conv7 = nn.Conv2d(64, 128, kernel_size=4) # 4 -> 1 + self.bn7 = nn.BatchNorm2d(128) + + self.dropout3 = nn.Dropout(0.4) + self.fc = nn.Linear(128, 10) + + def forward(self, x): + x = F.relu(self.bn1(self.conv1(x))) + x = F.relu(self.bn2(self.conv2(x))) + x = F.relu(self.bn3(self.conv3(x))) + x = self.dropout1(x) + + x = F.relu(self.bn4(self.conv4(x))) + x = F.relu(self.bn5(self.conv5(x))) + x = F.relu(self.bn6(self.conv6(x))) + x = self.dropout2(x) + + x = F.relu(self.bn7(self.conv7(x))) + x = torch.flatten(x, 1) + x = self.dropout3(x) + x = self.fc(x) + return x + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + # Define transforms for training (with augmentation) and test (no augmentation) + train_transform = transforms.Compose([ + transforms.RandomRotation(10), # rotation_range=10 + transforms.RandomAffine(0, translate=(0.1, 0.1)), # width/height shift + transforms.RandomResizedCrop(28, scale=(0.9, 1.1)), # zoom_range=0.10 + transforms.ToTensor(), # scale to [0,1] + ]) + + test_transform = transforms.Compose([ + transforms.ToTensor(), # only rescale + ]) + + # Download MNIST dataset + train_dataset = datasets.MNIST( + root="./data", train=True, download=True, transform=train_transform + ) + test_dataset = datasets.MNIST( + root="./data", train=False, download=True, transform=test_transform + ) + + # Create DataLoaders + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) + + return train_loader, test_loader + + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch)") + parser.add_argument("--output", type=str, default="mnist_model3.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = MNISTModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/bin/train/MNIST/adv/model_fgsm_MNIST_adv.py b/manuscripts/Poison26/bin/train/MNIST/adv/model_fgsm_MNIST_adv.py new file mode 100644 index 0000000..b5215c7 --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/adv/model_fgsm_MNIST_adv.py @@ -0,0 +1,253 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class MNISTModel(nn.Module): + def __init__(self): + super(MNISTModel, self).__init__() + self.conv1 = nn.Conv2d(1, 32, kernel_size=3) # 28x28 -> 26x26 + self.bn1 = nn.BatchNorm2d(32) + self.conv2 = nn.Conv2d(32, 32, kernel_size=3) # 26 -> 24 + self.bn2 = nn.BatchNorm2d(32) + self.conv3 = nn.Conv2d(32, 32, kernel_size=5, stride=2, padding=2) # 24 -> 12 + self.bn3 = nn.BatchNorm2d(32) + self.dropout1 = nn.Dropout(0.4) + + self.conv4 = nn.Conv2d(32, 64, kernel_size=3) # 12 -> 10 + self.bn4 = nn.BatchNorm2d(64) + self.conv5 = nn.Conv2d(64, 64, kernel_size=3) # 10 -> 8 + self.bn5 = nn.BatchNorm2d(64) + self.conv6 = nn.Conv2d(64, 64, kernel_size=5, stride=2, padding=2) # 8 -> 4 + self.bn6 = nn.BatchNorm2d(64) + self.dropout2 = nn.Dropout(0.4) + + self.conv7 = nn.Conv2d(64, 128, kernel_size=4) # 4 -> 1 + self.bn7 = nn.BatchNorm2d(128) + + self.dropout3 = nn.Dropout(0.4) + self.fc = nn.Linear(128, 10) + + def forward(self, x): + x = F.relu(self.bn1(self.conv1(x))) + x = F.relu(self.bn2(self.conv2(x))) + x = F.relu(self.bn3(self.conv3(x))) + x = self.dropout1(x) + + x = F.relu(self.bn4(self.conv4(x))) + x = F.relu(self.bn5(self.conv5(x))) + x = F.relu(self.bn6(self.conv6(x))) + x = self.dropout2(x) + + x = F.relu(self.bn7(self.conv7(x))) + x = torch.flatten(x, 1) + x = self.dropout3(x) + x = self.fc(x) + return x + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # converts to [0,1] + ]) + train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.MNIST(root="./data", train=False, download=True, transform=transform) + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) + return train_loader, test_loader + + +# ---------------------------- +# FGSM attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def fgsm_attack(model, x, y, eps=0.05, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using FGSM (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps is in the same scale as x (e.g., MNIST ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + x.requires_grad_(True) + logits = model(x) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x, retain_graph=False, create_graph=False)[0] + x_adv = x + eps * torch.sign(grad.detach()) + + # project back to [clamp_min, clamp_max] + x_adv = x_adv.clamp(clamp_min, clamp_max) + return x_adv.detach() + +# ---------------------------- +# Training loop (with optional FGSM adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under FGSM) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (FGSM)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch) with FGSM adversarial training") + parser.add_argument("--output", type=str, default="mnist_model4.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable FGSM adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="FGSM epsilon (in [0,1] pixel scale)") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = MNISTModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps) + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/bin/train/MNIST/adv/model_pgd_MNIST_adv.py b/manuscripts/Poison26/bin/train/MNIST/adv/model_pgd_MNIST_adv.py new file mode 100644 index 0000000..95fdc9e --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/adv/model_pgd_MNIST_adv.py @@ -0,0 +1,269 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class MNISTModel(nn.Module): + def __init__(self): + super(MNISTModel, self).__init__() + self.conv1 = nn.Conv2d(1, 32, kernel_size=3) # 28x28 -> 26x26 + self.bn1 = nn.BatchNorm2d(32) + self.conv2 = nn.Conv2d(32, 32, kernel_size=3) # 26 -> 24 + self.bn2 = nn.BatchNorm2d(32) + self.conv3 = nn.Conv2d(32, 32, kernel_size=5, stride=2, padding=2) # 24 -> 12 + self.bn3 = nn.BatchNorm2d(32) + self.dropout1 = nn.Dropout(0.4) + + self.conv4 = nn.Conv2d(32, 64, kernel_size=3) # 12 -> 10 + self.bn4 = nn.BatchNorm2d(64) + self.conv5 = nn.Conv2d(64, 64, kernel_size=3) # 10 -> 8 + self.bn5 = nn.BatchNorm2d(64) + self.conv6 = nn.Conv2d(64, 64, kernel_size=5, stride=2, padding=2) # 8 -> 4 + self.bn6 = nn.BatchNorm2d(64) + self.dropout2 = nn.Dropout(0.4) + + self.conv7 = nn.Conv2d(64, 128, kernel_size=4) # 4 -> 1 + self.bn7 = nn.BatchNorm2d(128) + + self.dropout3 = nn.Dropout(0.4) + self.fc = nn.Linear(128, 10) + + def forward(self, x): + x = F.relu(self.bn1(self.conv1(x))) + x = F.relu(self.bn2(self.conv2(x))) + x = F.relu(self.bn3(self.conv3(x))) + x = self.dropout1(x) + + x = F.relu(self.bn4(self.conv4(x))) + x = F.relu(self.bn5(self.conv5(x))) + x = F.relu(self.bn6(self.conv6(x))) + x = self.dropout2(x) + + x = F.relu(self.bn7(self.conv7(x))) + x = torch.flatten(x, 1) + x = self.dropout3(x) + x = self.fc(x) + return x + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # converts to [0,1] + ]) + train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.MNIST(root="./data", train=False, download=True, transform=transform) + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) + return train_loader, test_loader + + +# ---------------------------- +# PGD attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def pgd_attack(model, x, y, eps=0.05, alpha=0.01, steps=40, random_start=True, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using PGD (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps/alpha are in the same scale as x (e.g., MNIST ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + # start from a random point in the epsilon-ball if desired + if random_start: + x_adv = x + torch.empty_like(x).uniform_(-eps, eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + else: + x_adv = x.clone() + + for _ in range(steps): + x_adv.requires_grad_(True) + logits = model(x_adv) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x_adv, retain_graph=False, create_graph=False)[0] + x_adv = x_adv.detach() + alpha * torch.sign(grad.detach()) + + # project back to the epsilon l_inf ball around x, then clip to image bounds + x_adv = torch.max(torch.min(x_adv, x + eps), x - eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + + return x_adv.detach() + + +# ---------------------------- +# Training loop (with optional PGD adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, alpha=0.01, pgd_steps=40, random_start=True, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=random_start, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under PGD) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05, alpha=0.01, pgd_steps=40): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=True, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (PGD)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch) with PGD adversarial training") + parser.add_argument("--output", type=str, default="mnist_model5.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable PGD adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="PGD epsilon (in [0,1] pixel scale)") + parser.add_argument("--alpha", type=float, default=0.01, help="PGD step size (in [0,1] scale)") + parser.add_argument("--pgd-steps", type=int, default=40, help="Number of PGD steps") + parser.add_argument("--no-random-start", action="store_true", help="Disable random PGD start") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = MNISTModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, alpha=args.alpha, + pgd_steps=args.pgd_steps, random_start=not args.no_random_start, + adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps, alpha=args.alpha, pgd_steps=args.pgd_steps) + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/bin/train/MNIST/adv/model_standard_MNIST_adv.py b/manuscripts/Poison26/bin/train/MNIST/adv/model_standard_MNIST_adv.py new file mode 100644 index 0000000..09ad9a4 --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/adv/model_standard_MNIST_adv.py @@ -0,0 +1,195 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class MNISTModel(nn.Module): + def __init__(self): + super(MNISTModel, self).__init__() + self.conv1 = nn.Conv2d(1, 32, kernel_size=3) # 28x28 -> 26x26 + self.bn1 = nn.BatchNorm2d(32) + self.conv2 = nn.Conv2d(32, 32, kernel_size=3) # 26 -> 24 + self.bn2 = nn.BatchNorm2d(32) + self.conv3 = nn.Conv2d(32, 32, kernel_size=5, stride=2, padding=2) # 24 -> 12 + self.bn3 = nn.BatchNorm2d(32) + self.dropout1 = nn.Dropout(0.4) + + self.conv4 = nn.Conv2d(32, 64, kernel_size=3) # 12 -> 10 + self.bn4 = nn.BatchNorm2d(64) + self.conv5 = nn.Conv2d(64, 64, kernel_size=3) # 10 -> 8 + self.bn5 = nn.BatchNorm2d(64) + self.conv6 = nn.Conv2d(64, 64, kernel_size=5, stride=2, padding=2) # 8 -> 4 + self.bn6 = nn.BatchNorm2d(64) + self.dropout2 = nn.Dropout(0.4) + + self.conv7 = nn.Conv2d(64, 128, kernel_size=4) # 4 -> 1 + self.bn7 = nn.BatchNorm2d(128) + + self.dropout3 = nn.Dropout(0.4) + self.fc = nn.Linear(128, 10) + + def forward(self, x): + x = F.relu(self.bn1(self.conv1(x))) + x = F.relu(self.bn2(self.conv2(x))) + x = F.relu(self.bn3(self.conv3(x))) + x = self.dropout1(x) + + x = F.relu(self.bn4(self.conv4(x))) + x = F.relu(self.bn5(self.conv5(x))) + x = F.relu(self.bn6(self.conv6(x))) + x = self.dropout2(x) + + x = F.relu(self.bn7(self.conv7(x))) + x = torch.flatten(x, 1) + x = self.dropout3(x) + x = self.fc(x) + return x + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # converts to [0,1] + ]) + + train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.MNIST(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) + + return train_loader, test_loader + + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch)") + parser.add_argument("--output", type=str, default="mnist_model2.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = MNISTModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/bin/train/MNIST/basic/model_aug_MNIST_basic.py b/manuscripts/Poison26/bin/train/MNIST/basic/model_aug_MNIST_basic.py new file mode 100644 index 0000000..a4e1363 --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/basic/model_aug_MNIST_basic.py @@ -0,0 +1,188 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class MNISTModel(nn.Module): + def __init__(self): + super(MNISTModel, self).__init__() + self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1) + self.pool1 = nn.MaxPool2d(2, 2) + self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) + self.pool2 = nn.MaxPool2d(2, 2) + self.fc1 = nn.Linear(64 * 7 * 7, 128) + self.fc2 = nn.Linear(128, 10) + self.relu = nn.ReLU() + self.softmax = nn.Softmax(dim=1) + + def forward(self, x): + x = self.relu(self.conv1(x)) + x = self.pool1(x) + x = self.relu(self.conv2(x)) + x = self.pool2(x) + x = x.view(x.size(0), -1) # flatten + x = self.relu(self.fc1(x)) + x = self.fc2(x) # raw logits + return x + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + # Define transforms for training (with augmentation) and test (no augmentation) + train_transform = transforms.Compose([ + transforms.RandomRotation(10), # rotation_range=10 + transforms.RandomAffine(0, translate=(0.1, 0.1)), # width/height shift + transforms.RandomResizedCrop(28, scale=(0.9, 1.1)), # zoom_range=0.10 + transforms.ToTensor(), # scale to [0,1] + ]) + + test_transform = transforms.Compose([ + transforms.ToTensor(), # only rescale + ]) + + # Download MNIST dataset + train_dataset = datasets.MNIST( + root="./data", train=True, download=True, transform=train_transform + ) + test_dataset = datasets.MNIST( + root="./data", train=False, download=True, transform=test_transform + ) + + # Create DataLoaders + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) + + return train_loader, test_loader + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch)") + parser.add_argument("--output", type=str, default="mnist_model1.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = MNISTModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/bin/train/MNIST/basic/model_fgsm_MNIST_basic.py b/manuscripts/Poison26/bin/train/MNIST/basic/model_fgsm_MNIST_basic.py new file mode 100644 index 0000000..1c069ea --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/basic/model_fgsm_MNIST_basic.py @@ -0,0 +1,232 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class MNISTModel(nn.Module): + def __init__(self): + super(MNISTModel, self).__init__() + self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1) + self.pool1 = nn.MaxPool2d(2, 2) + self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) + self.pool2 = nn.MaxPool2d(2, 2) + self.fc1 = nn.Linear(64 * 7 * 7, 128) + self.fc2 = nn.Linear(128, 10) + self.relu = nn.ReLU() + self.softmax = nn.Softmax(dim=1) + + def forward(self, x): + x = self.relu(self.conv1(x)) + x = self.pool1(x) + x = self.relu(self.conv2(x)) + x = self.pool2(x) + x = x.view(x.size(0), -1) # flatten + x = self.relu(self.fc1(x)) + x = self.fc2(x) # raw logits + return x + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # converts to [0,1] + ]) + train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.MNIST(root="./data", train=False, download=True, transform=transform) + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) + return train_loader, test_loader + + +# ---------------------------- +# FGSM attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def fgsm_attack(model, x, y, eps=0.05, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using FGSM (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps is in the same scale as x (e.g., MNIST ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + x.requires_grad_(True) + logits = model(x) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x, retain_graph=False, create_graph=False)[0] + x_adv = x + eps * torch.sign(grad.detach()) + + # project back to [clamp_min, clamp_max] + x_adv = x_adv.clamp(clamp_min, clamp_max) + return x_adv.detach() + +# ---------------------------- +# Training loop (with optional FGSM adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under FGSM) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = fgsm_attack(model, images, labels, eps=eps, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (FGSM)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch) with FGSM adversarial training") + parser.add_argument("--output", type=str, default="mnist_model4.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable FGSM adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="FGSM epsilon (in [0,1] pixel scale)") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = MNISTModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps) + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/bin/train/MNIST/basic/model_pgd_MNIST_basic.py b/manuscripts/Poison26/bin/train/MNIST/basic/model_pgd_MNIST_basic.py new file mode 100644 index 0000000..ded5eac --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/basic/model_pgd_MNIST_basic.py @@ -0,0 +1,248 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class MNISTModel(nn.Module): + def __init__(self): + super(MNISTModel, self).__init__() + self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1) + self.pool1 = nn.MaxPool2d(2, 2) + self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) + self.pool2 = nn.MaxPool2d(2, 2) + self.fc1 = nn.Linear(64 * 7 * 7, 128) + self.fc2 = nn.Linear(128, 10) + self.relu = nn.ReLU() + self.softmax = nn.Softmax(dim=1) + + def forward(self, x): + x = self.relu(self.conv1(x)) + x = self.pool1(x) + x = self.relu(self.conv2(x)) + x = self.pool2(x) + x = x.view(x.size(0), -1) # flatten + x = self.relu(self.fc1(x)) + x = self.fc2(x) # raw logits + return x + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # converts to [0,1] + ]) + train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.MNIST(root="./data", train=False, download=True, transform=transform) + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True) + return train_loader, test_loader + + +# ---------------------------- +# PGD attack (untargeted) on [0,1] inputs +# ---------------------------- +@torch.enable_grad() +def pgd_attack(model, x, y, eps=0.05, alpha=0.01, steps=40, random_start=True, clamp_min=0.0, clamp_max=1.0): + """ + Generates adversarial examples for x using PGD (l_infty). + Assumes inputs are in [clamp_min, clamp_max]. + eps/alpha are in the same scale as x (e.g., MNIST ToTensor -> [0,1]). + """ + model_device = next(model.parameters()).device + x = x.detach().to(model_device) + y = y.detach().to(model_device) + + # start from a random point in the epsilon-ball if desired + if random_start: + x_adv = x + torch.empty_like(x).uniform_(-eps, eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + else: + x_adv = x.clone() + + for _ in range(steps): + x_adv.requires_grad_(True) + logits = model(x_adv) + loss = F.cross_entropy(logits, y) + grad = torch.autograd.grad(loss, x_adv, retain_graph=False, create_graph=False)[0] + x_adv = x_adv.detach() + alpha * torch.sign(grad.detach()) + + # project back to the epsilon l_inf ball around x, then clip to image bounds + x_adv = torch.max(torch.min(x_adv, x + eps), x - eps) + x_adv = x_adv.clamp(clamp_min, clamp_max) + + return x_adv.detach() + + +# ---------------------------- +# Training loop (with optional PGD adversarial training) +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001, adv=True, eps=0.05, alpha=0.01, pgd_steps=40, random_start=True, adv_lambda=1.0): + """ + adv=True: use adversarial examples in training. + adv_lambda in [0,1]: loss = (1-λ)*CE(clean) + λ*CE(adv). Set λ=1.0 for pure adversarial training. + """ + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + + for epoch in range(epochs): + start_time = time.time() + model.train() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device, non_blocking=True), labels.to(device, non_blocking=True) + + # optionally craft adversarial batch with BN/Dropout frozen for stability + if adv: + model.eval() + images_adv = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=random_start, clamp_min=0.0, clamp_max=1.0) + model.train() + else: + images_adv = None + + optimizer.zero_grad(set_to_none=True) + + if adv and adv_lambda >= 1.0 - 1e-8: + outputs = model(images_adv) + loss = criterion(outputs, labels) + elif adv and 0.0 < adv_lambda < 1.0: + out_clean = model(images) + out_adv = model(images_adv) + loss = (1.0 - adv_lambda) * criterion(out_clean, labels) + adv_lambda * criterion(out_adv, labels) + outputs = out_adv # for accuracy, count the adv preds + else: + outputs = model(images) + loss = criterion(outputs, labels) + + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation (clean or robust under PGD) +# ---------------------------- +def evaluate_model(model, test_loader, device, robust=False, eps=0.05, alpha=0.01, pgd_steps=40): + model.to(device) + criterion = nn.CrossEntropyLoss() + model.eval() + + y_true = [] + y_pred = [] + + test_loss = 0.0 + correct = 0 + total = 0 + + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + + if robust: + # generate test-time adversarial examples with BN/Dropout frozen + images_in = pgd_attack(model, images, labels, eps=eps, alpha=alpha, steps=pgd_steps, random_start=True, clamp_min=0.0, clamp_max=1.0) + else: + images_in = images + + with torch.no_grad(): + outputs = model(images_in) + loss = criterion(outputs, labels) + + test_loss += loss.item() + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.detach().cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).detach().cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + tag = "Robust (PGD)" if robust else "Clean" + print(f"{tag} Test Loss: {avg_loss:.4f}") + print(f"{tag} Test Accuracy: {accuracy:.4f}") + print(f"{tag} Test auROC: {auroc:.4f}") + print(f"{tag} Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch) with PGD adversarial training") + parser.add_argument("--output", type=str, default="mnist_model5.pt", help="Model output name") + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--lr", type=float, default=1e-3) + + # Adversarial training flags + parser.add_argument("--adv-train", action="store_true", help="Enable PGD adversarial training") + parser.add_argument("--eps", type=float, default=0.05, help="PGD epsilon (in [0,1] pixel scale)") + parser.add_argument("--alpha", type=float, default=0.01, help="PGD step size (in [0,1] scale)") + parser.add_argument("--pgd-steps", type=int, default=40, help="Number of PGD steps") + parser.add_argument("--no-random-start", action="store_true", help="Disable random PGD start") + parser.add_argument("--adv-lambda", type=float, default=1.0, help="λ in [0,1]: mix clean/adv loss (1.0 = pure adv)") + + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = MNISTModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, + epochs=args.epochs, lr=args.lr, + adv=args.adv_train, eps=args.eps, alpha=args.alpha, + pgd_steps=args.pgd_steps, random_start=not args.no_random_start, + adv_lambda=args.adv_lambda) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate (clean) + print("Evaluate test dataset") + evaluate_model(model, test_loader, device, robust=False) + evaluate_model(model, test_loader, device, robust=True, eps=args.eps, alpha=args.alpha, pgd_steps=args.pgd_steps) + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/bin/train/MNIST/basic/model_standard_MNIST_basic.py b/manuscripts/Poison26/bin/train/MNIST/basic/model_standard_MNIST_basic.py new file mode 100644 index 0000000..ccb6ddc --- /dev/null +++ b/manuscripts/Poison26/bin/train/MNIST/basic/model_standard_MNIST_basic.py @@ -0,0 +1,175 @@ +import argparse +import time +import numpy as np +from tqdm import tqdm + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from sklearn.metrics import roc_auc_score, average_precision_score +import matplotlib.pyplot as plt +import os, sys, json + + +# ---------------------------- +# Model definition +# ---------------------------- +class MNISTModel(nn.Module): + def __init__(self): + super(MNISTModel, self).__init__() + self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1) + self.pool1 = nn.MaxPool2d(2, 2) + self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1) + self.pool2 = nn.MaxPool2d(2, 2) + self.fc1 = nn.Linear(64 * 7 * 7, 128) + self.fc2 = nn.Linear(128, 10) + self.relu = nn.ReLU() + self.softmax = nn.Softmax(dim=1) + + def forward(self, x): + x = self.relu(self.conv1(x)) + x = self.pool1(x) + x = self.relu(self.conv2(x)) + x = self.pool2(x) + x = x.view(x.size(0), -1) # flatten + x = self.relu(self.fc1(x)) + x = self.fc2(x) # raw logits + return x + + +# ---------------------------- +# Dataset loading +# ---------------------------- +def load_data(batch_size=32): + transform = transforms.Compose([ + transforms.ToTensor(), # converts to [0,1] + ]) + + train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform) + test_dataset = datasets.MNIST(root="./data", train=False, download=True, transform=transform) + + train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True) + test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False) + + return train_loader, test_loader + + +# ---------------------------- +# Training loop +# ---------------------------- +def train(model, train_loader, device, epochs=10, lr=0.001): + criterion = nn.CrossEntropyLoss() + optimizer = optim.Adam(model.parameters(), lr=lr) + + model.to(device) + model.train() + + for epoch in range(epochs): + start_time = time.time() + running_loss = 0.0 + running_correct = 0 + total = 0 + + for images, labels in tqdm(train_loader, desc=f"Epoch {epoch+1}/{epochs}", unit="batch"): + images, labels = images.to(device), labels.to(device) + + optimizer.zero_grad() + outputs = model(images) + loss = criterion(outputs, labels) + loss.backward() + optimizer.step() + + running_loss += loss.item() + _, predicted = torch.max(outputs, 1) + running_correct += (predicted == labels).sum().item() + total += labels.size(0) + + avg_loss = running_loss / len(train_loader) + avg_acc = running_correct / total + elapsed = time.time() - start_time + print(f"Epoch {epoch+1} finished in {elapsed:.2f}s - Loss: {avg_loss:.4f}, Accuracy: {avg_acc:.4f}") + + +# ---------------------------- +# Evaluation +# ---------------------------- +def evaluate_model(model, test_loader, device): + model.to(device) + model.eval() + + y_true = [] + y_pred = [] + criterion = nn.CrossEntropyLoss() + + test_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(device), labels.to(device) + outputs = model(images) + + loss = criterion(outputs, labels) + test_loss += loss.item() + + _, predicted = torch.max(outputs, 1) + correct += (predicted == labels).sum().item() + total += labels.size(0) + + y_true.extend(labels.cpu().numpy()) + y_pred.extend(torch.softmax(outputs, dim=1).cpu().numpy()) + + avg_loss = test_loss / len(test_loader) + accuracy = correct / total + + y_true = np.array(y_true) + y_pred = np.array(y_pred) + + # compute AUROC and AUPRC + y_true_onehot = np.eye(10)[y_true] + auroc = roc_auc_score(y_true_onehot, y_pred, multi_class="ovr") + auprc = average_precision_score(y_true_onehot, y_pred) + + print(f"Test Loss: {avg_loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test auROC: {auroc:.4f}") + print(f"Test auPRC: {auprc:.4f}") + + +# ---------------------------- +# Main +# ---------------------------- +def main(): + parser = argparse.ArgumentParser(description="MNIST training code (PyTorch)") + parser.add_argument("--output", type=str, default="mnist_model1.pt", help="Model output name") + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--epochs", type=int, default=5, help="Number of training epochs") + args = parser.parse_args() + + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + # Initialize model + model = MNISTModel() + + # Load data + train_loader, test_loader = load_data(batch_size=args.batch_size) + + # Train + train(model, train_loader, device, epochs=args.epochs) + + # Save model + torch.save(model.state_dict(), args.output) + print(f"Model saved to {args.output}") + + # Evaluate + print("Model statistics on test dataset") + evaluate_model(model, test_loader, device) + + +if __name__ == "__main__": + main() + diff --git a/manuscripts/Poison26/singularity/README.md b/manuscripts/Poison26/singularity/README.md new file mode 100644 index 0000000..6ac80f9 --- /dev/null +++ b/manuscripts/Poison26/singularity/README.md @@ -0,0 +1,5 @@ +# Build SIF file +--- + ``` + singularity build apso_poison.sif apso_poison.def + ``` diff --git a/manuscripts/Poison26/singularity/apso_poison.def b/manuscripts/Poison26/singularity/apso_poison.def new file mode 100644 index 0000000..b3151bd --- /dev/null +++ b/manuscripts/Poison26/singularity/apso_poison.def @@ -0,0 +1,55 @@ +Bootstrap: docker +From: continuumio/miniconda3:latest + +%files + Adversarial_Observation /opt/Adversarial_Observation + +%post + apt-get update && apt-get install -y git && apt-get clean + + # Enable conda command in non-interactive build shell + . /opt/conda/etc/profile.d/conda.sh + + # Write your environment.yml inside the container + cat <<'EOF' > /tmp/environment.yml +name: pytorch-captum +channels: + - pytorch + - conda-forge + - defaults +dependencies: + - python=3.10 + - pytorch=2.4.1 + - torchvision=0.19.1 + - torchaudio=2.4.1 + - numpy + - scipy + - librosa + - matplotlib + - scikit-learn + - pandas + - imageio + - pip + - pip: + - captum +EOF + + # Create the conda environment + conda env create -f /tmp/environment.yml + + # Clean conda and pip caches + conda clean -a -y + rm -rf /opt/conda/pkgs/* + rm -rf /root/.cache/pip + + # Remove only the environment.yml we created + rm -f /tmp/environment.yml + +%environment + export PATH=/opt/conda/bin:$PATH + export PYTHONPATH=/opt:$PYTHONPATH + . /opt/conda/etc/profile.d/conda.sh + conda activate pytorch-captum + +%runscript + exec "$@" diff --git a/manuscripts/Poison26/singularity/pytorch-captum.def b/manuscripts/Poison26/singularity/pytorch-captum.def new file mode 100644 index 0000000..33d5f9b --- /dev/null +++ b/manuscripts/Poison26/singularity/pytorch-captum.def @@ -0,0 +1,40 @@ +Bootstrap: docker +From: continuumio/miniconda3:latest + +%post + apt-get update && apt-get install -y git && apt-get clean + + # Enable conda command in non-interactive build shell + . /opt/conda/etc/profile.d/conda.sh + + # Write your environment.yml inside the container + cat <<'EOF' > /tmp/environment.yml +name: pytorch-captum +channels: + - pytorch + - defaults +dependencies: + - python=3.10 + - pytorch=2.4.1 + - torchvision=0.19.1 + - torchaudio=2.4.1 + - numpy + - scipy + - matplotlib + - scikit-learn + - pandas + - imageio + - pip + - pip: + - captum +EOF + + conda env create -f /tmp/environment.yml + +%environment + export PATH=/opt/conda/bin:$PATH + . /opt/conda/etc/profile.d/conda.sh + conda activate pytorch-captum + +%runscript + exec "$@"