hello, world

Author: hsirm

AROS/Main.py

@@ -0,0 +1,126 @@
+import torch
+import torch.nn as nn
+import torchvision
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader
+from sklearn.mixture import GaussianMixture
+import numpy as np
+from scipy.stats import multivariate_normal
+from sklearn.covariance import EmpiricalCovariance
+from robustbench.utils import load_model
+import torch.nn.functional as F
+from torch.utils.data import TensorDataset
+
+
+num_vclasses=100
+
+
+num_samples_needed=1
+
+fast=True
+epoch1=1
+epoch2=1
+epoch3=1
+
+
+model_name_='Wang2023Better_WRN-70-16'
+in_dataset='cifar100'
+threat_model_='Linf'
+ 
+
+
+# cifa10_models=['Ding2020MMA','Rebuffi2021Fixing_70_16_cutmix_extra'] 50000//num_classes
+
+# cifar100_models=['Wang2023Better_WRN-70-16','Rice2020Overfitting']
+
+
+
+trainloader,testloader,ID_OOD_loader=get_loaders(in_dataset=in_dataset)
+
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+
+
+robust_backbone = load_model(model_name=model_name_, dataset=in_dataset, threat_model=threat_model_).to(device)
+last_layer_name, last_layer = list(robust_backbone.named_children())[-1]
+setattr(robust_backbone, last_layer_name, nn.Identity())
+
+
+
+embeddings, labels = [], []
+
+with torch.no_grad():
+    for imgs, lbls in trainloader:
+        imgs = imgs.to(device, non_blocking=True)
+        embed = robust_backbone(imgs).cpu()  # move to CPU only once per batch
+        embeddings.append(embed)
+        labels.append(lbls)
+embeddings = torch.cat(embeddings).numpy()
+labels = torch.cat(labels).numpy()
+
+
+print("embedding")
+
+
+if fast==False:
+  gmm_dict = {}
+  for cls in np.unique(labels):
+      cls_embed = embeddings[labels == cls]
+      gmm = GaussianMixture(n_components=1, covariance_type='full').fit(cls_embed)
+      gmm_dict[cls] = gmm
+
+  print("fake start")
+
+  fake_data = []
+
+
+  for cls, gmm in gmm_dict.items():
+      samples, likelihoods = [], []
+      while len(samples) < num_samples_needed:
+          s = gmm.sample(100)[0]
+          likelihood = gmm.score_samples(s)
+          samples.append(s[likelihood < np.quantile(likelihood, 0.001)])
+          likelihoods.append(likelihood[likelihood < np.quantile(likelihood, 0.001)])
+          if sum(len(smp) for smp in samples) >= num_samples_needed:
+              break
+      samples = np.vstack(samples)[:num_samples_needed]
+      fake_data.append(samples)
+
+  fake_data = np.vstack(fake_data)
+  fake_data = torch.tensor(fake_data).float()
+  fake_data = F.normalize(fake_data, p=2, dim=1)
+
+  fake_labels = torch.full((fake_data.shape[0],), 10)
+  fake_loader = DataLoader(TensorDataset(fake_data, fake_labels), batch_size=128, shuffle=True)
+
+if fast==True:
+
+
+    noise_std = 0.1  # standard deviation of noise
+    noisy_embeddings = torch.tensor(embeddings) + noise_std * torch.randn_like(torch.tensor(embeddings))
+
+    # Normalize Noisy Embeddings
+    noisy_embeddings = F.normalize(noisy_embeddings, p=2, dim=1)[:len(trainloader.dataset)//num_classes]
+
+    # Convert to DataLoader if needed
+    fake_labels = torch.full((noisy_embeddings.shape[0],), num_classes)[:len(trainloader.dataset)//num_classes]
+    fake_loader = DataLoader(TensorDataset(noisy_embeddings, fake_labels), batch_size=128, shuffle=True)
+
+
+
+final_model=stability_loss_function_(trainloader,testloader,robust_backbone,num_classes,fake_loader,last_layer)
+
+
+
+attack_eps = 8/255
+attack_steps = 10
+attack_alpha = 2.5 * attack_eps / attack_steps
+test_attack = PGD_AUC(final_model, eps=attack_eps, steps=attack_steps, alpha=attack_alpha, num_classes=num_classes)
+
+
+
+get_clean_AUC(final_model, ID_OOD_loader , device, num_classes)
+
+adv_auc = get_auc_adversarial(model=final_model,  test_loader=ID_OOD_loader, test_attack=test_attack, device=device, num_classes=num_classes)

AROS/data_loader.py

@@ -0,0 +1,74 @@
+import torch
+import torchvision
+from torch.utils.data import DataLoader, Dataset,Subset, SubsetRandomSampler, TensorDataset, ConcatDataset
+from torchvision import datasets, transforms, models
+from torchvision.datasets import CIFAR10, CIFAR100, MNIST, SVHN, FashionMNIST
+
+
+
+class LabelChangedDataset(Dataset):
+    def __init__(self, original_dataset, new_label):
+        self.original_dataset = original_dataset
+        self.new_label = new_label
+
+    def __len__(self):
+        return len(self.original_dataset)
+
+    def __getitem__(self, idx):
+        image, _ = self.original_dataset[idx]
+        return image, self.new_label
+
+
+
+def get_subsampled_subset(dataset, subset_ratio=0.1):
+    subset_size = int(len(dataset) * subset_ratio)
+    remaining_size = len(dataset) - subset_size
+
+    subset_indices, _ = torch.utils.data.random_split(dataset, [subset_size, remaining_size])
+    subset_testset = Subset(dataset, subset_indices.indices)
+
+    return subset_testset
+
+
+
+transform_tensor = transforms.Compose([ transforms.ToTensor()])
+
+
+trainset_CIFAR10 = torchvision.datasets.CIFAR10(
+    root='./data', train=True, download=True, transform=transform_tensor)
+
+
+testset_CIFAR10 = torchvision.datasets.CIFAR10(
+    root='./data', train=False, download=True, transform=transform_tensor)
+
+
+trainloader_CIFAR10 = DataLoader(trainset_CIFAR10, batch_size=64, shuffle=True, num_workers=2)
+
+testloader_CIFAR10 = DataLoader(testset_CIFAR10, batch_size=16, shuffle=False, num_workers=2)
+
+
+
+
+trainset_CIFAR100 = torchvision.datasets.CIFAR100(
+    root='./data', train=True, download=True, transform=transform_tensor)
+
+testset_CIFAR100 = torchvision.datasets.CIFAR100(
+    root='./data', train=False, download=True, transform=transform_tensor)
+
+trainloader_CIFAR100 = DataLoader(trainset_CIFAR100, batch_size=64, shuffle=True, num_workers=2)
+
+testloader_CIFAR100 = DataLoader(testset_CIFAR100, batch_size=16, shuffle=False, num_workers=2)
+
+
+testset_CIFAR10_relabled = LabelChangedDataset(testset_CIFAR10, new_label=100)
+testset_CIFAR100_relabled = LabelChangedDataset(testset_CIFAR100, new_label=10)
+
+
+testloader_CIFAR10_vs_CIFAR100 = DataLoader(ConcatDataset([testset_CIFAR10, testset_CIFAR100_relabled]), shuffle=False, batch_size=16)
+testloader_CIFAR100_vs_CIFAR10 = DataLoader(ConcatDataset([testset_CIFAR100, testset_CIFAR10_relabled]), shuffle=False, batch_size=16)
+
+def get_loaders(in_dataset='CIFAR10'):
+    if in_dataset == 'cifar10':
+        return trainloader_CIFAR10, testloader_CIFAR10, testloader_CIFAR10_vs_CIFAR100
+    else:
+        raise ValueError(f"Dataset '{in_dataset}' is not supported.")