initial commit

app.py

@@ -0,0 +1,312 @@
+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:  # CIFAR-10
+        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 = []
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+    dataset, transform = get_dataset_and_transform(dataset_name)
+    testloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=False)
+    logs.append(f"{dataset_name} dataset loaded")
+
+    model = initialize_model(model_name, device)
+    logs.append(f"Model {model_name} loaded on {device}")
+
+    param_count, layer_count = get_model_stats(model)
+    logs.append(f"Model stats: {param_count} parameters, {layer_count} layers")
+
+    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())}")
+
+    results = {attack_name: {'cm': [], 'mvl': []} for attack_name in attacks}
+
+    for i, (images, labels) in enumerate(testloader):
+        if i >= num_batches:
+            break
+        images, labels = images.to(device), labels.to(device)
+        logs.append(f"Processing batch {i+1}/{num_batches}...")
+
+        for attack_name, attack in attacks.items():
+            adv_images = attack(images, labels)
+            mvl_list = compute_mvl(model, images, adv_images, device)
+            results[attack_name]['mvl'].append(mvl_list)
+            cm = np.mean(mvl_list)
+            results[attack_name]['cm'].append(cm)
+
+    # Placeholders for plots (add your plot generation here)
+    cm_plot_path = None
+    mvl_plot_paths = [None]*5
+    integrated_mvl_plot_path = None
+
+    processing_time = time.time() - start_time
+
+    stats = {
+        'Dataset': dataset_name,
+        'Model': model_name,
+        'Parameter Count': param_count,
+        'Layer Count': layer_count,
+        'Processing Time (s)': round(processing_time, 2),
+        'Number of Batches': num_batches,
+        'Attacks Used': ', '.join(attacks.keys())
+    }
+    stats_text = "## Model Statistics\n\n| Metric | Value |\n|--------|-------|\n"
+    for k,v in stats.items():
+        stats_text += f"| {k} | {v} |\n"
+
+    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 Homaypour, 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):
+    models = get_models_for_dataset(dataset_name)
+    default_value = models[0] if models else None
+    return models, default_value  # Return choices and default value as a tuple
+
+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)
+
+        dataset_input = gr.Dropdown(datasets, label="Select Dataset", value='CIFAR-10')
+        model_input = gr.Dropdown(get_models_for_dataset('CIFAR-10'), label="Select Model")
+        attack_input = gr.CheckboxGroup(choices=attacks, label="Select Attacks", value=attacks)
+        batch_input = gr.Slider(minimum=1, maximum=20, step=1, value=5, 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")
+
+        # Return choices and value separately for older gradio versions
+        dataset_input.change(
+            fn=update_models,
+            inputs=dataset_input,
+            outputs=[model_input, model_input]
+        )
+
+        run_button.click(
+            fn=layer_sustainability_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()

requirements.txt

@@ -0,0 +1,8 @@
+torchattacks
+timm
+gradio
+datetime
+torch
+torchvision
+numpy
+matplotlib