| |
| |
|
|
| import torch |
| import torch.nn as nn |
| import torchvision |
| import torchvision.transforms as transforms |
| from torchvision.models import vgg16, vgg19, googlenet, resnet18 |
| import timm |
| import numpy as np |
| import matplotlib.pyplot as plt |
| from torchattacks import FGSM, PGD, APGD |
| import os |
| import time |
| from datetime import datetime |
| import gradio as gr |
|
|
| class LeNet(nn.Module): |
| def __init__(self): |
| super(LeNet, self).__init__() |
| self.conv1 = nn.Conv2d(1, 6, 5) |
| self.conv2 = nn.Conv2d(6, 16, 5) |
| self.fc1 = nn.Linear(16 * 4 * 4, 120) |
| self.fc2 = nn.Linear(120, 84) |
| self.fc3 = nn.Linear(84, 10) |
| self.relu = nn.ReLU() |
| self.pool = nn.MaxPool2d(2, 2) |
| |
| def forward(self, x, return_all=False): |
| outputs = [] |
| x1 = self.pool(self.relu(self.conv1(x))) |
| outputs.append(x1) |
| x2 = self.pool(self.relu(self.conv2(x1))) |
| outputs.append(x2) |
| x2_flat = x2.view(-1, 16 * 4 * 4) |
| x3 = self.relu(self.fc1(x2_flat)) |
| outputs.append(x3) |
| x4 = self.relu(self.fc2(x3)) |
| outputs.append(x4) |
| x5 = self.fc3(x4) |
| outputs.append(x5) |
| if return_all: |
| return outputs |
| else: |
| return x5 |
|
|
| def salt_pepper_noise(images, prob=0.01, device='cuda'): |
| batch_smap = torch.rand_like(images) < prob / 2 |
| pepper = torch.rand_like(images) < prob / 2 |
| noisy = images.clone() |
| noisy[batch_smap] = 1.0 |
| noisy[pepper] = 0.0 |
| return torch.clamp(noisy, 0, 1) |
|
|
| def pepper_statistical_noise(images, prob=0.01, device='cuda'): |
| pepper = torch.rand_like(images) < prob |
| noisy = images.clone() |
| noisy[pepper] = 0.0 |
| return torch.clamp(noisy, 0, 1) |
|
|
| def get_layer_outputs(model, input_tensor): |
| outputs = [] |
| def hook(module, input, output): |
| outputs.append(output) |
| hooks = [] |
| for layer in model.modules(): |
| if isinstance(layer, (nn.Conv2d, nn.Linear)): |
| hooks.append(layer.register_forward_hook(hook)) |
| model.eval() |
| with torch.no_grad(): |
| model(input_tensor) |
| for hook in hooks: |
| hook.remove() |
| return outputs |
|
|
| def compute_mvl(model, clean_images, adv_images, device='cuda'): |
| model.eval() |
| with torch.no_grad(): |
| try: |
| clean_outputs = model(clean_images, return_all=True) |
| adv_outputs = model(adv_images, return_all=True) |
| except TypeError: |
| clean_outputs = get_layer_outputs(model, clean_images) |
| adv_outputs = get_layer_outputs(model, adv_images) |
|
|
| mvl_list = [] |
| for clean_out, adv_out in zip(clean_outputs, adv_outputs): |
| if clean_out.ndim == 4: |
| diff = torch.norm(clean_out - adv_out, p=2, dim=(1,2,3)) |
| clean_norm = torch.norm(clean_out, p=2, dim=(1,2,3)) |
| else: |
| diff = torch.norm(clean_out - adv_out, p=2, dim=1) |
| clean_norm = torch.norm(clean_out, p=2, dim=1) |
| mvl = diff / (clean_norm + 1e-8) |
| mvl_list.append(mvl.mean().item()) |
| return mvl_list |
|
|
| def get_model_stats(model): |
| param_count = sum(p.numel() for p in model.parameters() if p.requires_grad) |
| layer_count = len([m for m in model.modules() if isinstance(m, (nn.Conv2d, nn.Linear))]) |
| return param_count, layer_count |
|
|
| def modify_model(model, model_name): |
| if model_name.startswith('VGG'): |
| model.classifier[6] = nn.Linear(4096, 10) |
| elif model_name == 'GoogLeNet': |
| model.fc = nn.Linear(1024, 10) |
| elif model_name == 'ResNet18': |
| model.fc = nn.Linear(512, 10) |
| elif model_name == 'WideResNet': |
| model.fc = nn.Linear(2048, 10) |
| elif model_name == 'DenseNet121': |
| model.classifier = nn.Linear(model.classifier.in_features, 10) |
| elif model_name == 'MobileNetV2': |
| if isinstance(model.classifier, nn.Sequential): |
| model.classifier[1] = nn.Linear(model.classifier[1].in_features, 10) |
| else: |
| model.classifier = nn.Linear(model.classifier.in_features, 10) |
| elif model_name == 'EfficientNet-B0': |
| model.classifier = nn.Linear(model.classifier.in_features, 10) |
| return model |
|
|
| def get_models_for_dataset(dataset_name): |
| if dataset_name == 'MNIST': |
| return ['LeNet'] |
| elif dataset_name == 'CIFAR-10': |
| return [ |
| 'VGG16', 'VGG19', 'GoogLeNet', 'ResNet18', 'WideResNet', |
| 'DenseNet121', 'MobileNetV2', 'EfficientNet-B0' |
| ] |
| else: |
| return [] |
|
|
| def get_dataset_and_transform(dataset_name): |
| if dataset_name == 'MNIST': |
| transform = transforms.Compose([ |
| transforms.Resize((28, 28)), |
| transforms.Grayscale(num_output_channels=1), |
| transforms.ToTensor(), |
| transforms.Normalize((0.1307,), (0.3081,)) |
| ]) |
| dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform) |
| else: |
| transform = transforms.Compose([ |
| transforms.Resize((224, 224)), |
| transforms.ToTensor(), |
| transforms.Normalize((0.485, 0.456, 0.406), |
| (0.229, 0.224, 0.225)) |
| ]) |
| dataset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform) |
| return dataset, transform |
|
|
| def initialize_model(model_name, device): |
| if model_name == 'LeNet': |
| model = LeNet() |
| elif model_name == 'VGG16': |
| model = modify_model(vgg16(weights='IMAGENET1K_V1'), model_name) |
| elif model_name == 'VGG19': |
| model = modify_model(vgg19(weights='IMAGENET1K_V1'), model_name) |
| elif model_name == 'GoogLeNet': |
| model = modify_model(googlenet(weights='IMAGENET1K_V1'), model_name) |
| elif model_name == 'ResNet18': |
| model = modify_model(resnet18(weights='IMAGENET1K_V1'), model_name) |
| elif model_name == 'WideResNet': |
| model = modify_model(timm.create_model('wide_resnet50_2', pretrained=True), model_name) |
| elif model_name == 'DenseNet121': |
| model = modify_model(timm.create_model('densenet121', pretrained=True), model_name) |
| elif model_name == 'MobileNetV2': |
| model = modify_model(timm.create_model('mobilenetv2_100', pretrained=True), model_name) |
| elif model_name == 'EfficientNet-B0': |
| model = modify_model(timm.create_model('efficientnet_b0', pretrained=True), model_name) |
| else: |
| raise ValueError(f"Unknown model {model_name}") |
| return model.to(device) |
|
|
| def layer_sustainability_analysis(dataset_name, model_name, selected_attacks, num_batches, output_dir_base='outputs'): |
| start_time = time.time() |
| logs = ["BSM:: experiment is being started ..."] |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
| logs.append(f"Loading {dataset_name} dataset...") |
| dataset, _ = get_dataset_and_transform(dataset_name) |
| testloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=False) |
| logs.append(f"{dataset_name} dataset loaded with {len(testloader)} batches.") |
|
|
| logs.append(f"Initializing model {model_name} on {device}...") |
| model = initialize_model(model_name, device) |
| logs.append(f"Model {model_name} initialized.") |
|
|
| param_count, layer_count = get_model_stats(model) |
| logs.append(f"Model stats: Parameters = {param_count}, Layers = {layer_count}") |
|
|
| all_attacks = { |
| 'FGSM': FGSM(model, eps=0.03), |
| 'PGD': PGD(model, eps=0.03, alpha=0.01, steps=40, random_start=True), |
| 'APGD': APGD(model, eps=0.03, steps=100, loss='ce'), |
| 'Salt & Pepper': lambda x, y: salt_pepper_noise(x, prob=0.01, device=device), |
| 'Pepper Statistical': lambda x, y: pepper_statistical_noise(x, prob=0.01, device=device) |
| } |
| attacks = {name: attack for name, attack in all_attacks.items() if name in selected_attacks} |
| if not attacks: |
| logs.append("Error: No valid attacks selected") |
| return ["No valid attacks selected", None] + [None]*6 + ["", '\n'.join(logs)] |
| logs.append(f"Selected attacks: {', '.join(attacks.keys())}") |
|
|
| timestamp = datetime.now().strftime('%Y%m%d_%H%M%S') |
| output_dir = os.path.join(output_dir_base, f"{model_name}_{timestamp}") |
| os.makedirs(output_dir, exist_ok=True) |
| logs.append(f"Output directory created: {output_dir}") |
|
|
| results = {atk: {'cm': [], 'mvl': []} for atk in attacks} |
|
|
| for i, (images, labels) in enumerate(testloader): |
| if i >= num_batches: |
| logs.append(f"Reached batch limit: {num_batches}") |
| break |
| images, labels = images.to(device), labels.to(device) |
| logs.append(f"Processing batch {i+1}/{num_batches}...") |
|
|
| for atk_name, atk in attacks.items(): |
| logs.append(f" Running attack: {atk_name} on batch {i+1}") |
| adv_images = atk(images, labels) |
| mvl_vals = compute_mvl(model, images, adv_images, device) |
| results[atk_name]['mvl'].append(mvl_vals) |
| batch_cm = np.mean(mvl_vals) |
| results[atk_name]['cm'].append(batch_cm) |
| logs.append(f" Attack {atk_name}: batch CM={batch_cm:.6f}") |
|
|
| logs.append("Finished processing batches, computing statistics...") |
|
|
| cm_means = {atk: np.mean(results[atk]['cm']) for atk in attacks} |
| cm_stds = {atk: np.std(results[atk]['cm']) for atk in attacks} |
|
|
| plt.figure(figsize=(8,6)) |
| attack_names = list(attacks.keys()) |
| means = [cm_means[a] for a in attack_names] |
| stds = [cm_stds[a] for a in attack_names] |
| x = np.arange(len(attack_names)) |
| plt.bar(x, means, yerr=stds, capsize=5) |
| plt.xticks(x, attack_names, rotation=45) |
| plt.ylabel("CM (Relative Error)") |
| plt.title(f"CM for {model_name} ({dataset_name})") |
| plt.tight_layout() |
| cm_plot_path = os.path.join(output_dir, "cm_plot.png") |
| plt.savefig(cm_plot_path) |
| plt.close() |
| logs.append(f"Saved CM plot: {cm_plot_path}") |
|
|
| mvl_plot_paths = [] |
| colors = ['skyblue', 'lightgreen', 'coral', 'lightgray', 'purple'] |
| for i, atk in enumerate(attack_names): |
| mvl_arr = np.array(results[atk]['mvl']) |
| mean_vals = np.mean(mvl_arr, axis=0) |
| std_vals = np.std(mvl_arr, axis=0) |
| layers = [f"Layer {j+1}" for j in range(len(mean_vals))] |
| plt.figure(figsize=(8,6)) |
| plt.plot(layers, mean_vals, marker='o', color=colors[i % len(colors)]) |
| plt.fill_between(layers, mean_vals - std_vals, mean_vals + std_vals, color=colors[i % len(colors)], alpha=0.3) |
| plt.title(f"MVL per Layer - {atk}") |
| plt.ylabel("MVL (Mean ± Std)") |
| plt.xticks(rotation=45) |
| plt.grid(True) |
| plt.tight_layout() |
| path = os.path.join(output_dir, f"mvl_{atk.lower().replace(' ', '_')}.png") |
| plt.savefig(path) |
| plt.close() |
| mvl_plot_paths.append(path) |
| logs.append(f"Saved MVL plot for {atk}: {path}") |
|
|
| plt.figure(figsize=(10,6)) |
| for i, atk in enumerate(attack_names): |
| mvl_arr = np.array(results[atk]['mvl']) |
| mean_vals = np.mean(mvl_arr, axis=0) |
| std_vals = np.std(mvl_arr, axis=0) |
| layers = [f"Layer {j+1}" for j in range(len(mean_vals))] |
| plt.plot(layers, mean_vals, marker='o', color=colors[i % len(colors)], label=atk) |
| plt.fill_between(layers, mean_vals - std_vals, mean_vals + std_vals, color=colors[i % len(colors)], alpha=0.3) |
| plt.title(f"Integrated MVL - {model_name}") |
| plt.ylabel("MVL (Mean ± Std)") |
| plt.xticks(rotation=45) |
| plt.legend() |
| plt.grid(True) |
| plt.tight_layout() |
| integrated_mvl_plot_path = os.path.join(output_dir, "integrated_mvl.png") |
| plt.savefig(integrated_mvl_plot_path) |
| plt.close() |
| logs.append(f"Saved integrated MVL plot: {integrated_mvl_plot_path}") |
|
|
| processing_time = time.time() - start_time |
| logs.append(f"Processing completed in {processing_time:.2f} seconds") |
|
|
| stats = { |
| 'Dataset': dataset_name, |
| 'Model': model_name, |
| 'Parameters': param_count, |
| 'Layers': layer_count, |
| 'Batches': num_batches, |
| 'Attacks': ', '.join(attack_names), |
| 'Time (s)': round(processing_time, 2) |
| } |
| stats_text = "## Model Statistics\n\n| Metric | Value |\n|---|---|\n" |
| for k,v in stats.items(): |
| stats_text += f"| {k} | {v} |\n" |
|
|
| while len(mvl_plot_paths) < 5: |
| mvl_plot_paths.append(None) |
|
|
| return [ |
| None, |
| cm_plot_path, |
| *mvl_plot_paths[:5], |
| integrated_mvl_plot_path, |
| stats_text, |
| '\n'.join(logs) |
| ] |
|
|
| paper_info_html = """ |
| <div style="border: 1px solid #ccc; padding: 15px; border-radius: 8px; margin-bottom: 15px;"> |
| <h2>Layer-wise Regularized Adversarial Training Using Layers Sustainability Analysis Framework</h2> |
| <h3>Authors</h3> |
| <p>Mohammad Khalooei, Mohammad Mehdi Homayounpour, Maryam Amirmazlaghani</p> |
| |
| <h3>Abstract</h3> |
| <ul> |
| <li>The layer sustainability analysis (LSA) framework is introduced to evaluate the behavior of layer-level representations of DNNs in dealing with network input perturbations using Lipschitz theoretical concepts.</li> |
| <li>A layer-wise regularized adversarial training (AT-LR) approach significantly improves the generalization and robustness of different deep neural network architectures for significant perturbations while reducing layer-level vulnerabilities.</li> |
| <li>AT-LR loss landscapes for each LSA MVL proposal can interpret layer importance for different layers, which is an intriguing aspect.</li> |
| </ul> |
| |
| <h3>Links</h3> |
| <ul> |
| <li><a href="https://arxiv.org/abs/2202.02626" target="_blank">ArXiv Paper</a></li> |
| <li><a href="https://github.com/khalooei/LSA" target="_blank">GitHub Repository</a></li> |
| <li><a href="https://www.sciencedirect.com/science/article/abs/pii/S0925231223002928" target="_blank">ScienceDirect Article</a></li> |
| </ul> |
| </div> |
| """ |
|
|
| def update_models(dataset_name): |
| if dataset_name == 'MNIST': |
| return gr.update(visible=False), "LeNet" |
| else: |
| models = get_models_for_dataset(dataset_name) |
| return gr.update(choices=models, value=models[0], visible=True), gr.update(visible=False) |
|
|
| def create_interface(): |
| datasets = ['MNIST', 'CIFAR-10'] |
| attacks = ['FGSM', 'PGD', 'APGD', 'Salt & Pepper', 'Pepper Statistical'] |
|
|
| with gr.Blocks() as interface: |
| gr.Markdown("# Layer-wise Sustainability Analysis") |
| gr.Markdown(paper_info_html) |
|
|
| initial_input="MNIST" |
| dataset_input = gr.Dropdown(datasets, label="Select Dataset", value=initial_input) |
| model_input = gr.Dropdown(get_models_for_dataset(initial_input), value=get_models_for_dataset(initial_input)[0], label="Select Model") |
| model_text = gr.Textbox(value="LeNet", visible=False, interactive=False, label="Model") |
|
|
| attack_input = gr.CheckboxGroup(choices=attacks, label="Select Attacks", value=attacks) |
| batch_input = gr.Slider(minimum=1, maximum=20, step=1, value=2, label="Number of Batches") |
| run_button = gr.Button("Run Analysis") |
|
|
| error_output = gr.Textbox(label="Error", visible=False) |
| cm_output = gr.Image(label="Comparative Measure (CM)") |
|
|
| with gr.Tabs(): |
| mvl_outputs = [] |
| for attack in attacks: |
| with gr.Tab(f"MVL: {attack}"): |
| mvl_output = gr.Image(label=f"MVL for {attack}") |
| mvl_outputs.append(mvl_output) |
| with gr.Tab("Integrated MVL"): |
| integrated_mvl_output = gr.Image(label="Integrated MVL for All Attacks") |
| with gr.Tab("Model Statistics"): |
| stats_output = gr.Markdown("## Model Statistics") |
| with gr.Tab("Logs"): |
| log_output = gr.Textbox(label="Processing Logs", lines=15, interactive=False) |
|
|
| dataset_input.change( |
| fn=update_models, |
| inputs=dataset_input, |
| outputs=[model_input, model_text] |
| ) |
|
|
| def get_model_for_mnist_or_dropdown(dataset_name, model_name): |
| return "LeNet" if dataset_name == 'MNIST' else model_name |
|
|
| def run_analysis(dataset_name, model_name, attacks, batches): |
| real_model = get_model_for_mnist_or_dropdown(dataset_name, model_name) |
| return layer_sustainability_analysis(dataset_name, real_model, attacks, batches) |
|
|
| run_button.click( |
| fn=run_analysis, |
| inputs=[dataset_input, model_input, attack_input, batch_input], |
| outputs=[error_output, cm_output] + mvl_outputs + [integrated_mvl_output, stats_output, log_output] |
| ) |
|
|
| return interface |
|
|
| if __name__ == '__main__': |
| interface = create_interface() |
| interface.launch() |
|
|
