fix: saga attack is now correct

feat: attention is now cached during attack instead of calculating it afterwards

app.py

@@ -7,7 +7,7 @@ from utils.model_loader import load_model_and_labels
 from utils.preprocessing import get_default_transform, preprocess_image
 from utils.inference import predict_topk
 from utils.attacks import PGDIterations, FGSM, SAGA, MIFGSM
-from utils.visualization import extract_attention_maps, attention_rollout, create_attention_overlay, extract_attention_for_iterations, create_iteration_attention_overlays
+from utils.visualization import extract_attention_maps, attention_rollout, create_attention_overlay, extract_attention_for_iterations, create_iteration_attention_overlays, rollout_from_cached_attentions
 
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -56,7 +56,7 @@ def classify_image(model_file, image):
         img_tensor = preprocess_image(image, transform=transform).to(DEVICE)
 
         # Inferência
-        top_prob, top_idx, num_classes, probabilities = predict_topk(model, img_tensor, top_k=5, device=DEVICE)
+        top_prob, top_idx, num_classes, probabilities = predict_topk(model, img_tensor, top_k=5)
 
         top_k = len(top_prob)
         result = f"**Top {top_k} Predictions:**\n\n"
@@ -177,7 +177,7 @@ def run_attack(
     discard_ratio: float,
     head_fusion: str,
     alpha_overlay: float
-) -> Tuple[List[Image.Image], str, Image.Image, List[Image.Image]]:
+) -> Tuple[List[Image.Image], str, List[Image.Image]]:
     """
     Executa ataque adversarial (FGSM ou PGD) untargeted e extrai atenção.
     
@@ -193,13 +193,13 @@ def run_attack(
         alpha_overlay: transparência da sobreposição
     
     Returns:
-        (iteration_images, result_text, final_adv_image, attention_overlays)
+        (iteration_images, result_text, attention_overlays)
     """
     try:
         if model_file is None:
-            return [], "Please upload a model file (.pth)", None, []
+            return [], "Please upload a model file (.pth)", []
         if image is None:
-            return [], "Please upload an image", None, []
+            return [], "Please upload an image", []
         
         # Carregar modelo e labels
         model_path = _to_path(model_file)
@@ -230,18 +230,26 @@ def run_attack(
         adv_tensor, iteration_images = attack(img_tensor, original_label)
         
         # Extrair atenção para todas as iterações
-        # SAGA já calcula atenção internamente, podemos reutilizar!
-        if attack_type == "SAGA" and hasattr(attack, 'attention_masks_cache') and len(attack.attention_masks_cache) > 0:
-            # Usar atenção já calculada pelo SAGA
-            attention_masks = attack.attention_masks_cache
+        # Importante: para visualização, respeitar parâmetros do UI (discard_ratio, head_fusion)
+        # Mesmo quando o ataque é SAGA, recomputar os mapas para permitir controle do usuário.
+        # Preferir atenções em cache capturadas durante o ataque
+        attention_masks = []
+        cached_attns = getattr(attack, 'attentions_per_iter', None)
+        if cached_attns and len(cached_attns) == len(iteration_images):
+            attention_masks = rollout_from_cached_attentions(
+                cached_attns,
+                discard_ratio=discard_ratio,
+                head_fusion=head_fusion
+            )
         else:
-            # Para FGSM e PGD, calcular atenção normalmente
+            # Fallback: recomputar via hooks (depreciado)
             attention_masks = extract_attention_for_iterations(
                 model,
                 attack.iteration_tensors,
                 discard_ratio=discard_ratio,
                 head_fusion=head_fusion
             )
+            print("⚠️ Warning: recomputing attentions via hooks. Consider using an attack class that caches attentions.")
         
         # Criar overlays de atenção
         attention_overlays = create_iteration_attention_overlays(
@@ -255,10 +263,6 @@ def run_attack(
         adv_class = top_idx_adv[0].item()
         adv_prob = top_prob_adv[0].item()
         
-        # Converter imagem adversarial final para PIL
-        from utils.attacks import tensor_to_pil
-        final_adv_image = tensor_to_pil(adv_tensor[0])
-        
         # Format result
         result = f"## {attack_type} Attack Result (Untargeted)\n\n"
         result += f"**Configuration:**\n"
@@ -305,12 +309,12 @@ def run_attack(
         result += f"- Iteration 0 = Original image\n"
         result += f"- Iteration {steps} = Final adversarial image\n"
         
-        return iteration_images, result, final_adv_image, attention_overlays
+        return iteration_images, result, attention_overlays
         
     except Exception as e:
         import traceback
         error_msg = f"Error executing attack:\n```\n{str(e)}\n\n{traceback.format_exc()}\n```"
-        return [], error_msg, None, []
+        return [], error_msg, []
 
 
 def create_app():
@@ -318,7 +322,7 @@ def create_app():
     
     with gr.Blocks(title="ViTViz - Classifier & Attacks", theme=gr.themes.Citrus()) as app:
         gr.Markdown("""
-        # ViTViz: Vision Transformer Classifier & Adversarial Attacks
+        # ViTViz: Vision Transformer Adversarial Attack & Attention Visualization Tool
         """)
         
         with gr.Tabs():
@@ -406,12 +410,13 @@ def create_app():
                     outputs=[attention_output, output_text_attention]
                 )
 
-            # Tab 3: AAdversarial Attack + Attention
+            # Tab 3: Adversarial Attack + Attention
             with gr.Tab("Adversarial Attack + Attention"):
                 gr.Markdown("### Execute adversarial attacks and visualize how attention evolves")
                 
-                with gr.Row():
-                    with gr.Column(scale=1):
+                # with gr.Row():
+                with gr.Column(scale=1):
+                    with gr.Row():
                         model_upload_attack = gr.File(
                             label="Upload Model (.pth/.pt)",
                             file_types=[".pth", ".pt"]
@@ -420,45 +425,48 @@ def create_app():
                             label="Upload Image",
                             type="pil"
                         )
+                    gr.Markdown("---")
+
+                    with gr.Row():
+                        with gr.Column():
+                            gr.Markdown("#### Attack Configuration")
+                    
+                            attack_type = gr.Dropdown(
+                                choices=["PGD", "FGSM", "MIFGSM", "SAGA"],
+                                value="PGD",
+                                label="Attack Type",
+                                info="PGD: iterative | FGSM: single-step | MIFGSM: momentum | SAGA: gradient × attention"
+                            )
                         
-                        gr.Markdown("#### Attack Configuration")
-                        
-                        attack_type = gr.Dropdown(
-                            choices=["PGD", "FGSM", "MIFGSM", "SAGA"],
-                            value="PGD",
-                            label="Attack Type",
-                            info="PGD: iterative | FGSM: single-step | MIFGSM: momentum | SAGA: gradient × attention"
-                        )
-                        
-                        eps_input = gr.Slider(
-                            minimum=0.0,
-                            maximum=1.0,
-                            value=8/255,
-                            step=1/255,
-                            label="Epsilon (ε) - Maximum Perturbation"
-                        )
-                        
-                        # Parâmetros iterativos
-                        alpha_group = gr.Group(visible=True)
-                        with alpha_group:
-                            alpha_input = gr.Slider(
+                            eps_input = gr.Slider(
                                 minimum=0.0,
-                                maximum=0.1,
-                                value=2/255,
+                                maximum=1.0,
+                                value=8/255,
                                 step=1/255,
-                                label="Alpha (α) - Step Size"
-                            )
-                        
-                        steps_group = gr.Group(visible=True)
-                        with steps_group:
-                            steps_input = gr.Slider(
-                                minimum=1,
-                                maximum=100,
-                                value=10,
-                                step=1,
-                                label="Steps - NNumber of Iterations"
+                                label="Epsilon (ε) - Maximum Perturbation"
                             )
+                    
+                            # Parâmetros iterativos
+                            alpha_group = gr.Group(visible=True)
+                            with alpha_group:
+                                alpha_input = gr.Slider(
+                                    minimum=0.0,
+                                    maximum=0.1,
+                                    value=2/255,
+                                    step=1/255,
+                                    label="Alpha (α) - Step Size"
+                                )
                         
+                            steps_group = gr.Group(visible=True)
+                            with steps_group:
+                                steps_input = gr.Slider(
+                                    minimum=1,
+                                    maximum=100,
+                                    value=10,
+                                    step=1,
+                                    label="Steps - NNumber of Iterations"
+                                )
+                    
                         # Função para atualizar visibilidade dos parâmetros
                         def update_attack_params(attack_type):
                             if attack_type == "FGSM":
@@ -477,15 +485,15 @@ def create_app():
                             outputs=[alpha_group, steps_group]
                         )
                         
-                        gr.Markdown("#### Attention Configuration")
+                        with gr.Column():
+                            gr.Markdown("#### Attention Rollout Configuration")
                         
-                        head_fusion_attack = gr.Radio(
-                            choices=["mean", "max", "min"],
-                            value="max",
-                            label="Head Fusion"
-                        )
-                        
-                        with gr.Row():
+                            head_fusion_attack = gr.Radio(
+                                choices=["mean", "max", "min"],
+                                value="max",
+                                label="Head Fusion"
+                            )
+
                             discard_ratio_attack = gr.Slider(
                                 minimum=0.0,
                                 maximum=1.0,
@@ -501,22 +509,10 @@ def create_app():
                                 label="Alpha Overlay"
                             )
                         
-                        attack_btn = gr.Button("Execute Full Analysis", variant="primary", size="lg")
-                    
-                    with gr.Column(scale=2):
-                        output_text_attack = gr.Markdown(label="Result")
-                        
-                        with gr.Row():
-                            with gr.Column():
-                                gr.Markdown("**Final Adversarial Image**")
-                                final_adv_image = gr.Image(type="pil", show_label=False)
-                            with gr.Column():
-                                gr.Markdown("**All Iterations**")
-                                iteration_gallery = gr.Gallery(
-                                    columns=5,
-                                    height="auto",
-                                    show_label=False
-                                )
+                    attack_btn = gr.Button("Execute Full Analysis", variant="primary", size="lg")
+                
+                with gr.Column(scale=2):
+                    output_text_attack = gr.Markdown(label="Result")
                 
                 # Seção de Evolução da Atenção
                 gr.Markdown("---")
@@ -538,10 +534,11 @@ def create_app():
                         maximum=10,
                         step=1,
                         value=0,
-                        label="Attack Iteration"
+                        label="Attack Iteration",
+                        info="0 for the original image, last is the final adversarial image"
                     )
                     show_attention_checkbox = gr.Checkbox(
-                        value=False,
+                        value=True,
                         label="Show Attention Heatmap",
                         info="Overlay attention map on images"
                     )
@@ -584,11 +581,7 @@ def create_app():
                         attack_type, eps_input, alpha_input, steps_input,
                         discard_ratio_attack, head_fusion_attack, alpha_overlay_attack
                     ],
-                    outputs=[iteration_images_state, output_text_attack, final_adv_image, attention_overlays_state]
-                ).then(
-                    fn=lambda x: x,
-                    inputs=[iteration_images_state],
-                    outputs=[iteration_gallery]
+                    outputs=[iteration_images_state, output_text_attack, attention_overlays_state]
                 ).then(
                     fn=lambda imgs: gr.update(maximum=len(imgs)-1 if imgs else 0, value=0),
                     inputs=[iteration_images_state],
@@ -626,8 +619,12 @@ def create_app():
 
 if __name__ == "__main__":
     app = create_app()
-    app.launch(
-        server_name="0.0.0.0",
-        server_port=7861,
-        share=False
-    )
+    try:
+        app.launch(
+            server_name="0.0.0.0",
+            server_port=7861,
+            share=False
+        )
+
+    except KeyboardInterrupt:
+        print("\nShutting down gracefully...")

utils/attacks.py

@@ -3,6 +3,41 @@ import torchattacks
 from PIL import Image
 from typing import List, Tuple
 import numpy as np
+def capture_outputs_and_attentions(model, x_norm: torch.Tensor):
+    """Executa um forward único capturando atenções via hooks nas camadas de atenção do ViT.
+    Retorna (outputs, attentions_list) onde attentions_list é lista de tensores [B,H,T,T] por camada.
+    Funciona para modelos do timm com atributo 'blocks' e submódulo 'attn'.
+    """
+    attentions: List[torch.Tensor] = []
+
+    def make_attention_hook():
+        def hook(module, input, output):
+            x = input[0]
+            B, N, C = x.shape
+            if not (hasattr(module, 'qkv') and hasattr(module, 'num_heads')):
+                return
+            qkv = module.qkv(x).reshape(B, N, 3, module.num_heads, C // module.num_heads).permute(2, 0, 3, 1, 4)
+            q, k, v = qkv.unbind(0)
+            scale = (C // module.num_heads) ** -0.5
+            attn = (q @ k.transpose(-2, -1)) * scale
+            attn = attn.softmax(dim=-1)
+            attentions.append(attn.detach())
+        return hook
+
+    hooks = []
+    if hasattr(model, 'blocks'):
+        for block in model.blocks:
+            if hasattr(block, 'attn'):
+                hooks.append(block.attn.register_forward_hook(make_attention_hook()))
+
+    model.eval()
+    outputs = model(x_norm)
+
+    for h in hooks:
+        h.remove()
+
+    attentions = [a.cpu() for a in attentions]
+    return outputs, attentions
 
 
 def denormalize_imagenet(tensor: torch.Tensor) -> torch.Tensor:
@@ -57,6 +92,8 @@ class FGSM(torchattacks.FGSM):
         super().__init__(model, eps=eps)
         self.iteration_images: List[Image.Image] = []
         self.iteration_tensors: List[torch.Tensor] = []
+        # Atenções por iteração (iteração 0: original, iteração 1: adversarial)
+        self.attentions_per_iter: List[List[torch.Tensor]] = []
     
     def forward(self, images, labels) -> Tuple[torch.Tensor, List[Image.Image]]:
         """
@@ -77,6 +114,7 @@ class FGSM(torchattacks.FGSM):
         
         self.iteration_images = []
         self.iteration_tensors = []
+        self.attentions_per_iter = []
         
         # Salvar imagem original
         pil_img_orig = tensor_to_pil(images_denorm[0], denormalize=False)
@@ -85,7 +123,9 @@ class FGSM(torchattacks.FGSM):
         
         # Calcular gradiente
         images.requires_grad = True
-        outputs = self.get_logits(images)
+        # Capturar atenções e logits para imagem original
+        outputs, attentions0 = capture_outputs_and_attentions(self.model, images)
+        self.attentions_per_iter.append([att for att in attentions0])
         
         if self.targeted:
             target_labels = self.get_target_label(images, labels)
@@ -106,6 +146,10 @@ class FGSM(torchattacks.FGSM):
         pil_img_adv = tensor_to_pil(adv_images_denorm[0], denormalize=False)
         self.iteration_images.append(pil_img_adv)
         self.iteration_tensors.append(adv_images.clone().detach())
+
+        # Capturar atenções para imagem adversarial final
+        outputs_adv, attentions1 = capture_outputs_and_attentions(self.model, adv_images)
+        self.attentions_per_iter.append([att for att in attentions1])
         
         return adv_images, self.iteration_images
 
@@ -120,6 +164,7 @@ class PGDIterations(torchattacks.PGD):
         super().__init__(model, eps=eps, alpha=alpha, steps=steps, random_start=random_start)
         self.iteration_images: List[Image.Image] = []
         self.iteration_tensors: List[torch.Tensor] = []
+        self.attentions_per_iter: List[List[torch.Tensor]] = []
 
     def forward(self, images, labels) -> Tuple[torch.Tensor, List[Image.Image]]:
         """
@@ -154,17 +199,21 @@ class PGDIterations(torchattacks.PGD):
 
         self.iteration_images = []
         self.iteration_tensors = []
+        self.attentions_per_iter = []
         
         # Salvar iteração 0 (imagem original)
         pil_img_orig = tensor_to_pil(images_denorm[0], denormalize=False)
         self.iteration_images.append(pil_img_orig)
         self.iteration_tensors.append(images.clone().detach())
+        # Atenções da imagem original
+        outputs0, attentions0 = capture_outputs_and_attentions(self.model, images)
+        self.attentions_per_iter.append([att for att in attentions0])
 
         for _ in range(self.steps):
             # Normalizar para passar pelo modelo
             adv_images = (adv_images_denorm - mean) / std
             adv_images.requires_grad = True
-            outputs = self.get_logits(adv_images)
+            outputs, attentions = capture_outputs_and_attentions(self.model, adv_images)
 
             # Calculate loss
             if self.targeted:
@@ -188,12 +237,13 @@ class PGDIterations(torchattacks.PGD):
             pil_img = tensor_to_pil(adv_images_denorm[0], denormalize=False)
             self.iteration_images.append(pil_img)
             self.iteration_tensors.append(adv_images_normalized.clone().detach())
+            # Atenções desta iteração
+            self.attentions_per_iter.append([att for att in attentions])
 
         # Retornar imagem normalizada para o modelo
         adv_images = (adv_images_denorm - mean) / std
         return adv_images, self.iteration_images
-
-
+    
 class SAGA(torch.nn.Module):
     """
     SAGA: Self-Attention Gradient Attack
@@ -205,6 +255,170 @@ class SAGA(torch.nn.Module):
     Baseado em: https://github.com/MetaMain/ViTRobust
     Paper: "On the Robustness of Vision Transformers to Adversarial Examples" (ICCV 2021)
     """
+
+    def __init__(self, model, eps=8/255, steps=10, discard_ratio: float = 0.0, head_fusion: str = "mean"):
+        """Implementação correta do SAGA baseada no código original (SelfAttentionGradientAttack).
+
+        Parâmetros:
+        - model: Vision Transformer (deve expor atenções via forward ou função auxiliar em visualization utils)
+        - eps: orçamento L_inf máximo (em pixel space [0,1])
+        - steps: número de iterações (FGSM iterativo)
+        - discard_ratio: razão de descarte usada no attention rollout
+        - head_fusion: estratégia de fusão de heads ('mean','max','min')
+        """
+        super().__init__()
+        self.model = model
+        self.eps = eps
+        self.steps = steps
+        self.eps_step = self.eps / max(1, steps)
+        self.discard_ratio = discard_ratio
+        self.head_fusion = head_fusion
+        self.device = next(model.parameters()).device
+        self.iteration_images: List[Image.Image] = []
+        self.iteration_tensors: List[torch.Tensor] = []
+        self.attention_masks_cache: List[np.ndarray] = []
+        # Cache opcional: atenções por camada/head em cada iteração
+        # Formato: lista por iteração; cada item é a lista de tensores [B, H, T, T] por camada
+        self.attentions_per_iter: List[List[torch.Tensor]] = []
+
+    def _attention_map(self, images_norm: torch.Tensor, save: bool = False) -> torch.Tensor:
+        """Extrai mapa de atenção (rollout) e retorna tensor expandido [B,3,H,W] em [0,1].
+        images_norm: imagens já normalizadas para o forward do modelo.
+        """
+        # Esta função agora assume que as atenções foram capturadas no mesmo forward
+        # e serão passadas externamente; mantida para compatibilidade se necessário.
+        raise RuntimeError("_attention_map should not be called directly; use integrated forward attention capture.")
+
+    def _capture_outputs_and_attentions(self, x_norm: torch.Tensor):
+        """Executa um forward único capturando atenções via hooks nas camadas de atenção do ViT.
+        Retorna (outputs, attentions_list) onde attentions_list é lista de tensores [B,H,T,T] por camada.
+        """
+        attentions: List[torch.Tensor] = []
+
+        def make_attention_hook():
+            def hook(module, input, output):
+                # input[0] é o embedding antes de atenção (B, N, C)
+                x = input[0]
+                B, N, C = x.shape
+                if not (hasattr(module, 'qkv') and hasattr(module, 'num_heads')):
+                    return
+                qkv = module.qkv(x).reshape(B, N, 3, module.num_heads, C // module.num_heads).permute(2, 0, 3, 1, 4)
+                q, k, v = qkv.unbind(0)
+                scale = (C // module.num_heads) ** -0.5
+                attn = (q @ k.transpose(-2, -1)) * scale
+                attn = attn.softmax(dim=-1)
+                attentions.append(attn.detach())
+            return hook
+
+        hooks = []
+        if not hasattr(self.model, 'blocks'):
+            outputs = self.model(x_norm)
+            return outputs, []
+        for block in self.model.blocks:
+            if hasattr(block, 'attn'):
+                hooks.append(block.attn.register_forward_hook(make_attention_hook()))
+
+        self.model.eval()
+        outputs = self.model(x_norm)
+
+        for h in hooks:
+            h.remove()
+
+        # mover atenções para CPU para cache leve
+        attentions = [a.cpu() for a in attentions]
+        return outputs, attentions
+
+    def forward(self, images, labels) -> Tuple[torch.Tensor, List[Image.Image]]:
+        """Executa o ataque SAGA (FGSM iterativo com ponderação por atenção).
+
+        Fluxo por iteração:
+        1. Normaliza a imagem adversarial atual.
+        2. Calcula loss e gradiente.
+        3. Extrai mapa de atenção da imagem atual e pondera gradiente.
+        4. Aplica passo FGSM (sign) em pixel space [0,1].
+        5. Projeta em L_inf (clamp delta) e clip final para [0,1].
+        6. Salva imagem e tensor normalizado.
+        """
+        images = images.clone().detach().to(self.device)
+        labels = labels.clone().detach().to(self.device)
+
+        # Mean/std ImageNet para conversão entre espaços
+        mean = torch.tensor([0.485, 0.456, 0.406], device=self.device).view(1, 3, 1, 1)
+        std = torch.tensor([0.229, 0.224, 0.225], device=self.device).view(1, 3, 1, 1)
+
+        # Pixel space [0,1]
+        images_denorm = images * std + mean
+        adv_denorm = images_denorm.clone().detach()
+
+        # Reset buffers
+        self.iteration_images = []
+        self.iteration_tensors = []
+        self.attention_masks_cache = []
+        self.attentions_per_iter = []
+
+        # Iteração 0 (imagem original)
+        self.iteration_images.append(tensor_to_pil(images_denorm[0], denormalize=False))
+        self.iteration_tensors.append(images.clone().detach())
+        # Atenção da imagem original: captura integrada
+        outputs0, attentions0 = self._capture_outputs_and_attentions(images)
+        # Guardar atenções brutas
+        self.attentions_per_iter.append([att for att in attentions0])
+        # Gerar máscara de rollout para cache visual
+        from utils.visualization import attention_rollout
+        import cv2
+        b, _, h, w = images.shape
+        mask0 = attention_rollout(attentions0, discard_ratio=self.discard_ratio, head_fusion=self.head_fusion)
+        mask0_resized = cv2.resize(mask0, (w, h))
+        self.attention_masks_cache.append(mask0.copy())
+
+        loss_fn = torch.nn.CrossEntropyLoss()
+
+        for _ in range(self.steps):
+            # Normalizar para forward
+            adv_norm = (adv_denorm - mean) / std
+            adv_norm.requires_grad = True
+            outputs, attentions = self._capture_outputs_and_attentions(adv_norm)
+            if isinstance(outputs, tuple):  # compatibilidade com modelos que retornam extras
+                outputs = outputs[0]
+            loss = loss_fn(outputs, labels)
+            grad = torch.autograd.grad(loss, adv_norm, retain_graph=False, create_graph=False)[0]
+
+            # Atenção da imagem adversarial atual (já capturada)
+            # Cache de atenções por camada/head
+            self.attentions_per_iter.append([att for att in attentions])
+            # Rollout para gerar mapa usado na ponderação
+            mask = attention_rollout(attentions, discard_ratio=self.discard_ratio, head_fusion=self.head_fusion)
+            mask_resized = cv2.resize(mask, (adv_norm.shape[-1], adv_norm.shape[-2]))
+            mmax = mask_resized.max() if mask_resized.max() > 0 else 1.0
+            mask_resized = (mask_resized / mmax).astype('float32')
+            att_map = torch.from_numpy(mask_resized).to(self.device).unsqueeze(0).unsqueeze(0).repeat(adv_norm.size(0), 3, 1, 1)
+            # Cache visual
+            self.attention_masks_cache.append(mask.copy())
+            grad_weighted = grad * att_map
+
+            # FGSM step em pixel space (sign do gradiente normalizado equivale ao do desnormalizado)
+            adv_denorm = adv_denorm.detach() + self.eps_step * grad_weighted.sign()
+
+            # Projeção na bola L_inf de raio eps em relação à imagem original
+            delta = torch.clamp(adv_denorm - images_denorm, min=-self.eps, max=self.eps)
+            adv_denorm = torch.clamp(images_denorm + delta, 0.0, 1.0).detach()
+
+            # Salvar artefatos desta iteração
+            self.iteration_images.append(tensor_to_pil(adv_denorm[0], denormalize=False))
+            self.iteration_tensors.append(((adv_denorm - mean) / std).clone().detach())
+
+        # Retorna tensor normalizado final
+        adv_final = (adv_denorm - mean) / std
+        return adv_final, self.iteration_images
+                
+
+
+class AttentionWeightedPGD(torch.nn.Module):
+    """
+    [Deprecated]
+    Implementação errada do ataque SAGA, mas que consegue fazer ataques
+    adversariais eficazes em ViTs usando mapas de atenção para pesar o gradiente.
+    """
     def __init__(self, model, eps=0.03, steps=10):
         super().__init__()
         self.model = model

utils/model_loader.py

@@ -110,17 +110,20 @@ def build_model_from_checkpoint(checkpoint: Any, device: Optional[torch.device]
         elif 'state_dict' in checkpoint:
             state_dict = checkpoint['state_dict']
             num_classes = infer_num_classes(state_dict)
+            # TODO: fazer tratamento dinâmico para timm, pytorch, etc.
             model = timm.create_model('vit_base_patch16_224', pretrained=False, num_classes=num_classes)
             model.load_state_dict(state_dict)
         elif 'model_state_dict' in checkpoint:
             # Novo formato com class_names embutidas
             state_dict = checkpoint['model_state_dict']
             num_classes = infer_num_classes(state_dict)
+            # TODO: fazer tratamento dinâmico para timm, pytorch, etc.
             model = timm.create_model('vit_base_patch16_224', pretrained=False, num_classes=num_classes)
             model.load_state_dict(state_dict)
         else:
             # assume dict é um state_dict
             num_classes = infer_num_classes(checkpoint)
+            # TODO: fazer tratamento dinâmico para timm, pytorch, etc.
             model = timm.create_model('vit_base_patch16_224', pretrained=False, num_classes=num_classes)
             model.load_state_dict(checkpoint)
     else:
@@ -147,13 +150,16 @@ def load_model_and_labels(
     device = device or DEVICE_DEFAULT
     checkpoint = load_checkpoint(model_path, device=device)
     class_names_ckpt = extract_class_names(checkpoint)
-    class_names_file = load_class_names_from_file(labels_file)
-    class_names = class_names_file or class_names_ckpt
-    source: Optional[str] = None
-    if class_names_file:
-        source = 'file'
-    elif class_names_ckpt:
-        source = 'checkpoint'
+    # class_names_file = load_class_names_from_file(labels_file)
+    # class_names = class_names_file or class_names_ckpt
+    # source: Optional[str] = None
+    # if class_names_file:
+    #     source = 'file'
+    # elif class_names_ckpt:
+    #     source = 'checkpoint'
+
+    class_names = class_names_ckpt
+    source = 'checkpoint' if class_names_ckpt else None
 
     model = build_model_from_checkpoint(checkpoint, device=device)
     return model, class_names, source

utils/visualization.py

@@ -77,9 +77,9 @@ def extract_attention_maps(model, image: torch.Tensor) -> list:
     return attentions
 
 
-def attention_rollout(attentions: list, 
-                     discard_ratio: float = 0.9, 
-                     head_fusion: str = 'max') -> np.ndarray:
+def attention_rollout(attentions: list,
+                      discard_ratio: float = 0.9,
+                      head_fusion: str = 'max') -> np.ndarray:
     """
     Implementa Attention Rollout seguindo a implementação original.
     
@@ -93,12 +93,12 @@ def attention_rollout(attentions: list,
     Returns:
         mask: Array numpy [grid_size, grid_size] com valores normalizados [0, 1]
     """
-    # Inicializar com matriz identidade (CORREÇÃO 1)
+    # Inicializar com matriz identidade
     result = torch.eye(attentions[0].size(-1))
     
     with torch.no_grad():
         for attention in attentions:
-            # Agregar heads (CORREÇÃO 2: usar axis=1, não dim=0)
+            # Agregar heads
             if head_fusion == 'mean':
                 attention_heads_fused = attention.mean(axis=1)
             elif head_fusion == 'max':
@@ -107,12 +107,20 @@ def attention_rollout(attentions: list,
                 attention_heads_fused = attention.min(axis=1)[0]
             else:
                 raise ValueError(f"head_fusion deve ser 'mean', 'max' ou 'min'")
-
-            # Descartar atenções fracas, mas proteger CLS token
-            flat = attention_heads_fused.view(attention_heads_fused.size(0), -1)
-            _, indices = flat.topk(int(flat.size(-1) * discard_ratio), -1, False)
-            indices = indices[indices != 0]  # Proteger CLS token
-            flat[0, indices] = 0
+            # Aplicar descarte condicional das atenções fracas por amostra
+            if discard_ratio > 0.0:
+                bsz, tokens, _ = attention_heads_fused.shape
+                flat = attention_heads_fused.view(bsz, -1)
+                k = int(flat.size(-1) * discard_ratio)
+                if k > 0:
+                    # Menores valores (largest=False)
+                    vals, idxs = torch.topk(flat, k, dim=-1, largest=False)
+                    for b in range(bsz):
+                        idxs_b = idxs[b]
+                        # proteger CLS (posição 0 nas matrizes quadradas)
+                        idxs_b = idxs_b[idxs_b != 0]
+                        flat[b, idxs_b] = 0
+                    attention_heads_fused = flat.view(bsz, tokens, tokens)
 
             # Adicionar identidade e normalizar
             I = torch.eye(attention_heads_fused.size(-1))
@@ -124,7 +132,6 @@ def attention_rollout(attentions: list,
             # Rollout recursivo
             result = torch.matmul(a, result)
     
-    # CORREÇÃO 4: Extrair atenção do CLS token (batch, CLS, patches)
     # Look at the total attention between the class token and the image patches
     mask = result[0, 0, 1:]
     
@@ -188,7 +195,8 @@ def extract_attention_for_iterations(
     head_fusion: str = 'max'
 ) -> list:
     """
-    Extrai mapas de atenção para cada iteração do ataque PGD.
+    [Deprecated when cached attentions are present]
+    Extrai mapas de atenção para cada iteração do ataque usando hooks.
     
     Args:
         model: Modelo ViT
@@ -217,6 +225,49 @@ def extract_attention_for_iterations(
     return attention_masks
 
 
+def rollout_from_cached_attentions(
+    attentions_per_iter: list,
+    discard_ratio: float = 0.9,
+    head_fusion: str = 'max'
+) -> list:
+    """
+    Gera máscaras de atenção por iteração a partir de atenções já capturadas no ataque.
+
+    Args:
+        attentions_per_iter: Lista por iteração; cada item é a lista de tensores [B, H, T, T] por camada
+        discard_ratio: Proporção de atenções fracas a descartar
+        head_fusion: Como agregar heads ('mean', 'max', 'min')
+
+    Returns:
+        Lista de máscaras de atenção [grid, grid] normalizadas [0, 1]
+    """
+    attention_masks = []
+
+    if attentions_per_iter is None or len(attentions_per_iter) == 0:
+        return attention_masks
+
+    for layer_attns in attentions_per_iter:
+        # layer_attns: lista de tensores por camada [B, H, T, T]
+        # Garantir CPU e detach
+        attentions_cpu = []
+        for att in layer_attns:
+            if isinstance(att, torch.Tensor):
+                attentions_cpu.append(att.detach().cpu())
+            else:
+                # já é CPU numpy/tensor? tentar converter via torch.as_tensor
+                attentions_cpu.append(torch.as_tensor(att))
+
+        # Aplicar rollout padrão sobre a lista de camadas
+        mask = attention_rollout(
+            attentions_cpu,
+            discard_ratio=discard_ratio,
+            head_fusion=head_fusion
+        )
+        attention_masks.append(mask)
+
+    return attention_masks
+
+
 def create_iteration_attention_overlays(
     iteration_images: list,
     attention_masks: list,