Update inference.py

inference.py

@@ -10,7 +10,6 @@ import os
 import requests
 import time
 from pathlib import Path
-from spaces import GPU
 
 # Check CUDA availability
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -21,6 +20,7 @@ MODEL_URLS = {
     'robust_resnet50': 'https://huggingface.co/madrylab/robust-imagenet-models/resolve/main/resnet50_l2_eps3.ckpt',
     'standard_resnet50': 'https://huggingface.co/madrylab/robust-imagenet-models/resolve/main/resnet50_l2_eps0.ckpt'
 }
+
 IMAGENET_MEAN = [0.485, 0.456, 0.406]
 IMAGENET_STD = [0.229, 0.224, 0.225]
 
@@ -105,12 +105,12 @@ def get_inference_configs(eps=0.5, n_itr=50):
         'loss_function': 'CE',  # Loss function: Cross Entropy
         'n_itr': n_itr,  # Number of iterations
         'eps': eps,  # Maximum perturbation size
-        'step_size': 0.02,  # Step size for each iteration
+        'step_size': 1,  # Step size for each iteration
         'diffusion_noise_ratio': 0.0,  # No diffusion noise
         'initial_inference_noise_ratio': 0.0,  # No initial noise
         'top_layer': 'all',  # Use all layers of the model
-        'inference_normalization': 'on',  # Apply normalization during inference
-        'recognition_normalization': 'on',  # Apply normalization during recognition
+        'inference_normalization': 'off',  # Apply normalization during inference
+        'recognition_normalization': 'off',  # Apply normalization during recognition
         'iterations_to_show': [1, 5, 10, 20, 30, 40, 50, n_itr]  # Specific iterations to visualize
     }
     return config
@@ -121,52 +121,384 @@ class GenerativeInferenceModel:
         self.normalizer = NormalizeByChannelMeanStd(IMAGENET_MEAN, IMAGENET_STD).to(device)
         self.labels = get_imagenet_labels()
         
+    def verify_model_integrity(self, model, model_type):
+        """
+        Verify model integrity by running a test input through it.
+        Returns whether the model passes basic integrity check.
+        """
+        try:
+            print(f"\n=== Running model integrity check for {model_type} ===")
+            # Create a deterministic test input
+            test_input = torch.zeros(1, 3, 224, 224)
+            test_input[0, 0, 100:124, 100:124] = 0.5  # Red square
+            test_input = test_input.to(model.device if hasattr(model, 'device') else 'cpu')
+            
+            # Run forward pass
+            with torch.no_grad():
+                output = model(test_input)
+            
+            # Check output shape
+            if output.shape != (1, 1000):
+                print(f"❌ Unexpected output shape: {output.shape}, expected (1, 1000)")
+                return False
+                
+            # Get top prediction
+            probs = torch.nn.functional.softmax(output, dim=1)
+            confidence, prediction = torch.max(probs, 1)
+            
+            # Calculate basic statistics on output
+            mean = output.mean().item()
+            std = output.std().item()
+            min_val = output.min().item()
+            max_val = output.max().item()
+            
+            print(f"Model integrity check results:")
+            print(f"- Output shape: {output.shape}")
+            print(f"- Top prediction: Class {prediction.item()} with {confidence.item()*100:.2f}% confidence")
+            print(f"- Output statistics: mean={mean:.3f}, std={std:.3f}, min={min_val:.3f}, max={max_val:.3f}")
+            
+            # Basic sanity checks
+            if torch.isnan(output).any():
+                print("❌ Model produced NaN outputs")
+                return False
+                
+            if output.std().item() < 0.1:
+                print("⚠️ Low output variance, model may not be discriminative")
+                
+            print("✅ Model passes basic integrity check")
+            return True
+            
+        except Exception as e:
+            print(f"❌ Model integrity check failed with error: {e}")
+            return False
+    
     def load_model(self, model_type):
+        """Load model from checkpoint or use pretrained model."""
         if model_type in self.models:
             return self.models[model_type]
         
         model_path = download_model(model_type)
         
-        # Create standard ResNet50 model
-        model = models.resnet50()
+        # Create a sequential model with normalizer and ResNet50
+        resnet = models.resnet50()
+        model = nn.Sequential(
+            self.normalizer,  # Normalizer is part of the model sequence
+            resnet
+        )
         
         # Load the model checkpoint
         if model_path:
             print(f"Loading {model_type} model from {model_path}...")
-            checkpoint = torch.load(model_path, map_location=device)
-            
-            # Handle different checkpoint formats
-            if 'model' in checkpoint:
-                # Format from madrylab robust models
-                state_dict = checkpoint['model']
-            elif 'state_dict' in checkpoint:
-                state_dict = checkpoint['state_dict']
-            else:
-                # Direct state dict
-                state_dict = checkpoint
-                
-            # Handle prefix in state dict keys
-            new_state_dict = {}
-            for key, value in state_dict.items():
-                if key.startswith('module.'):
-                    new_key = key[7:]  # Remove 'module.' prefix
+            try:
+                checkpoint = torch.load(model_path, map_location=device)
+                
+                # Print checkpoint structure for better understanding
+                print("\n=== Analyzing checkpoint structure ===")
+                if isinstance(checkpoint, dict):
+                    print(f"Checkpoint contains keys: {list(checkpoint.keys())}")
+                    
+                    # Examine 'model' structure if it exists
+                    if 'model' in checkpoint and isinstance(checkpoint['model'], dict):
+                        model_dict = checkpoint['model']
+                        # Get sample of keys to understand structure
+                        first_keys = list(model_dict.keys())[:5]
+                        print(f"'model' contains keys like: {first_keys}")
+                        
+                        # Check for common prefixes in the model dict
+                        prefixes = set()
+                        for key in list(model_dict.keys())[:100]:  # Check first 100 keys
+                            parts = key.split('.')
+                            if len(parts) > 1:
+                                prefixes.add(parts[0])
+                        if prefixes:
+                            print(f"Common prefixes in model dict: {prefixes}")
+                else:
+                    print(f"Checkpoint is not a dictionary, but a {type(checkpoint)}")
+                
+                # Handle different checkpoint formats
+                if 'model' in checkpoint:
+                    # Format from madrylab robust models
+                    state_dict = checkpoint['model']
+                    print("Using 'model' key from checkpoint")
+                elif 'state_dict' in checkpoint:
+                    state_dict = checkpoint['state_dict']
+                    print("Using 'state_dict' key from checkpoint")
+                else:
+                    # Direct state dict
+                    state_dict = checkpoint
+                    print("Using checkpoint directly as state_dict")
+                
+                # Handle prefix in state dict keys for ResNet part
+                resnet_state_dict = {}
+                prefixes_to_try = ['', 'module.', 'model.', 'attacker.model.']
+                resnet_keys = set(resnet.state_dict().keys())
+                
+                # First check if we can find keys directly in the attacker.model path
+                print("\n=== Phase 1: Checking for specific model structures ===")
+                
+                # Check for 'module.model' structure (seen in actual checkpoint)
+                module_model_keys = [key for key in state_dict.keys() if key.startswith('module.model.')]
+                if module_model_keys:
+                    print(f"Found 'module.model' structure with {len(module_model_keys)} parameters")
+                    # Extract all parameters from module.model
+                    for source_key, value in state_dict.items():
+                        if source_key.startswith('module.model.'):
+                            target_key = source_key[len('module.model.'):]
+                            resnet_state_dict[target_key] = value
+                            
+                    print(f"Extracted {len(resnet_state_dict)} parameters from module.model")
+                
+                # Check for 'attacker.model' structure
+                attacker_model_keys = [key for key in state_dict.keys() if key.startswith('attacker.model.')]
+                if attacker_model_keys:
+                    print(f"Found 'attacker.model' structure with {len(attacker_model_keys)} parameters")
+                    # Extract all parameters from attacker.model
+                    for source_key, value in state_dict.items():
+                        if source_key.startswith('attacker.model.'):
+                            target_key = source_key[len('attacker.model.'):]
+                            resnet_state_dict[target_key] = value
+                            
+                    print(f"Extracted {len(resnet_state_dict)} parameters from attacker.model")
+                    
+                    # Check if 'model' (not attacker.model) exists as a fallback
+                    model_keys = [key for key in state_dict.keys() if key.startswith('model.') and not key.startswith('attacker.model.')]
+                    if model_keys and len(resnet_state_dict) < len(resnet_keys):
+                        print(f"Found additional 'model.' structure with {len(model_keys)} parameters")
+                        # Try to complete missing parameters
+                        for source_key, value in state_dict.items():
+                            if source_key.startswith('model.'):
+                                target_key = source_key[len('model.'):]
+                                if target_key in resnet_keys and target_key not in resnet_state_dict:
+                                    resnet_state_dict[target_key] = value
+                                    
                 else:
-                    new_key = key
-                new_state_dict[new_key] = value
+                    # Check for other known structures
+                    structure_found = False
+                    
+                    # Check for 'model.' prefix
+                    model_keys = [key for key in state_dict.keys() if key.startswith('model.')]
+                    if model_keys:
+                        print(f"Found 'model.' structure with {len(model_keys)} parameters")
+                        for source_key, value in state_dict.items():
+                            if source_key.startswith('model.'):
+                                target_key = source_key[len('model.'):]
+                                resnet_state_dict[target_key] = value
+                        structure_found = True
+                    
+                    # Check for ResNet parameters at the top level
+                    top_level_resnet_keys = 0
+                    for key in resnet_keys:
+                        if key in state_dict:
+                            top_level_resnet_keys += 1
+                            
+                    if top_level_resnet_keys > 0:
+                        print(f"Found {top_level_resnet_keys} ResNet parameters at top level")
+                        for target_key in resnet_keys:
+                            if target_key in state_dict:
+                                resnet_state_dict[target_key] = state_dict[target_key]
+                        structure_found = True
+                    
+                    # If no structure was recognized, try the prefix mapping approach
+                    if not structure_found:
+                        print("No standard model structure found, trying prefix mappings...")
+                        for target_key in resnet_keys:
+                            for prefix in prefixes_to_try:
+                                source_key = prefix + target_key
+                                if source_key in state_dict:
+                                    resnet_state_dict[target_key] = state_dict[source_key]
+                                    break
                 
-            model.load_state_dict(new_state_dict)
+                # If we still can't find enough keys, try a final approach of removing prefixes
+                if len(resnet_state_dict) < len(resnet_keys):
+                    print(f"Found only {len(resnet_state_dict)}/{len(resnet_keys)} parameters, trying prefix removal...")
+                    
+                    # Track matches found through prefix removal
+                    prefix_matches = {prefix: 0 for prefix in ['module.', 'model.', 'attacker.model.', 'attacker.']}
+                    layer_matches = {}  # Track matches by layer type
+                    
+                    # Count parameter keys by layer type for analysis
+                    for key in resnet_keys:
+                        layer_name = key.split('.')[0] if '.' in key else key
+                        if layer_name not in layer_matches:
+                            layer_matches[layer_name] = {'total': 0, 'matched': 0}
+                        layer_matches[layer_name]['total'] += 1
+                    
+                    # Try keys with common prefixes
+                    for source_key, value in state_dict.items():
+                        # Skip if already found
+                        target_key = source_key
+                        matched_prefix = None
+                        
+                        # Try removing various prefixes
+                        for prefix in ['module.', 'model.', 'attacker.model.', 'attacker.']:
+                            if source_key.startswith(prefix):
+                                target_key = source_key[len(prefix):]
+                                matched_prefix = prefix
+                                break
+                        
+                        # If the target key is in the ResNet keys, add it to the state dict
+                        if target_key in resnet_keys and target_key not in resnet_state_dict:
+                            resnet_state_dict[target_key] = value
+                            
+                            # Update match statistics
+                            if matched_prefix:
+                                prefix_matches[matched_prefix] += 1
+                            
+                            # Update layer matches
+                            layer_name = target_key.split('.')[0] if '.' in target_key else target_key
+                            if layer_name in layer_matches:
+                                layer_matches[layer_name]['matched'] += 1
+                    
+                    # Print detailed prefix removal statistics
+                    print("\n=== Prefix Removal Statistics ===")
+                    total_matches = sum(prefix_matches.values())
+                    print(f"Total parameters matched through prefix removal: {total_matches}/{len(resnet_keys)} ({(total_matches/len(resnet_keys))*100:.1f}%)")
+                    
+                    # Show matches by prefix
+                    print("\nMatches by prefix:")
+                    for prefix, count in sorted(prefix_matches.items(), key=lambda x: x[1], reverse=True):
+                        if count > 0:
+                            print(f"  {prefix}: {count} parameters")
+                    
+                    # Show matches by layer type
+                    print("\nMatches by layer type:")
+                    for layer, stats in sorted(layer_matches.items(), key=lambda x: x[1]['total'], reverse=True):
+                        match_percent = (stats['matched'] / stats['total']) * 100 if stats['total'] > 0 else 0
+                        print(f"  {layer}: {stats['matched']}/{stats['total']} ({match_percent:.1f}%)")
+                    
+                    # Check for specific important layers (conv1, layer1, etc.)
+                    critical_layers = ['conv1', 'bn1', 'layer1', 'layer2', 'layer3', 'layer4', 'fc']
+                    print("\nStatus of critical layers:")
+                    for layer in critical_layers:
+                        if layer in layer_matches:
+                            match_percent = (layer_matches[layer]['matched'] / layer_matches[layer]['total']) * 100
+                            status = "✅ COMPLETE" if layer_matches[layer]['matched'] == layer_matches[layer]['total'] else "⚠️ INCOMPLETE"
+                            print(f"  {layer}: {layer_matches[layer]['matched']}/{layer_matches[layer]['total']} ({match_percent:.1f}%) - {status}")
+                        else:
+                            print(f"  {layer}: Not found in model")
+                
+                # Load the ResNet state dict
+                if resnet_state_dict:
+                    try:
+                        # Use strict=False to allow missing keys
+                        result = resnet.load_state_dict(resnet_state_dict, strict=False)
+                        missing_keys, unexpected_keys = result
+                        
+                        # Generate detailed information with better formatting
+                        loading_report = []
+                        loading_report.append(f"\n===== MODEL LOADING REPORT: {model_type} =====")
+                        loading_report.append(f"Total parameters in checkpoint: {len(resnet_state_dict):,}")
+                        loading_report.append(f"Total parameters in model: {len(resnet.state_dict()):,}")
+                        loading_report.append(f"Missing keys: {len(missing_keys):,} parameters")
+                        loading_report.append(f"Unexpected keys: {len(unexpected_keys):,} parameters")
+
+                        # Calculate percentage of parameters loaded
+                        loaded_keys = set(resnet_state_dict.keys()) - set(unexpected_keys)
+                        loaded_percent = (len(loaded_keys) / len(resnet.state_dict())) * 100
+                        
+                        # Determine loading success status
+                        if loaded_percent >= 99.5:
+                            status = "✅ COMPLETE - All important parameters loaded"
+                        elif loaded_percent >= 90:
+                            status = "🟡 PARTIAL - Most parameters loaded, should still function"
+                        elif loaded_percent >= 50:
+                            status = "⚠️ INCOMPLETE - Many parameters missing, may not function properly"
+                        else:
+                            status = "❌ FAILED - Critical parameters missing, will not function properly"
+                            
+                        loading_report.append(f"Successfully loaded: {len(loaded_keys):,} parameters ({loaded_percent:.1f}%)")
+                        loading_report.append(f"Loading status: {status}")
+                        
+                        # If loading is severely incomplete, fall back to PyTorch's pretrained model
+                        if loaded_percent < 50:
+                            loading_report.append("\n⚠️ WARNING: Loading from checkpoint is too incomplete.")
+                            loading_report.append("⚠️ Falling back to PyTorch's pretrained model to avoid broken inference.")
+                            
+                            # Create a new ResNet model with pretrained weights
+                            resnet = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
+                            model = nn.Sequential(self.normalizer, resnet)
+                            loading_report.append("✅ Successfully loaded PyTorch's pretrained ResNet50 model")
+                        
+                        # Show missing keys by layer type
+                        if missing_keys:
+                            loading_report.append("\nMissing keys by layer type:")
+                            layer_types = {}
+                            for key in missing_keys:
+                                # Extract layer type (e.g., 'conv', 'bn', 'layer1', etc.)
+                                parts = key.split('.')
+                                if len(parts) > 0:
+                                    layer_type = parts[0]
+                                    if layer_type not in layer_types:
+                                        layer_types[layer_type] = 0
+                                    layer_types[layer_type] += 1
+                            
+                            # Add counts by layer type
+                            for layer_type, count in sorted(layer_types.items(), key=lambda x: x[1], reverse=True):
+                                loading_report.append(f"  {layer_type}: {count:,} parameters")
+                                
+                            loading_report.append("\nFirst 10 missing keys:")
+                            for i, key in enumerate(sorted(missing_keys)[:10]):
+                                loading_report.append(f"  {i+1}. {key}")
+                        
+                        # Show unexpected keys if any
+                        if unexpected_keys:
+                            loading_report.append("\nFirst 10 unexpected keys:")
+                            for i, key in enumerate(sorted(unexpected_keys)[:10]):
+                                loading_report.append(f"  {i+1}. {key}")
+                                
+                        loading_report.append("========================================")
+                        
+                        # Convert report to string and print it
+                        report_text = "\n".join(loading_report)
+                        print(report_text)
+                        
+                        # Also save to a file for reference
+                        os.makedirs("logs", exist_ok=True)
+                        with open(f"logs/model_loading_{model_type}.log", "w") as f:
+                            f.write(report_text)
+                        
+                        # Look for normalizer parameters as well
+                        if any(key.startswith('attacker.normalize.') for key in state_dict.keys()):
+                            norm_state_dict = {}
+                            for key, value in state_dict.items():
+                                if key.startswith('attacker.normalize.'):
+                                    norm_key = key[len('attacker.normalize.'):]
+                                    norm_state_dict[norm_key] = value
+                            
+                            if norm_state_dict:
+                                try:
+                                    self.normalizer.load_state_dict(norm_state_dict, strict=False)
+                                    print("Successfully loaded normalizer parameters")
+                                except Exception as e:
+                                    print(f"Warning: Could not load normalizer parameters: {e}")
+                    except Exception as e:
+                        print(f"Warning: Error loading ResNet parameters: {e}")
+                        # Fall back to loading without normalizer
+                        model = resnet  # Use just the ResNet model without normalizer
+            except Exception as e:
+                print(f"Error loading model checkpoint: {e}")
+                # Fallback to PyTorch's pretrained model
+                print("Falling back to PyTorch's pretrained model")
+                resnet = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
+                model = nn.Sequential(self.normalizer, resnet)
         else:
             # Fallback to PyTorch's pretrained model
-            model = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
+            print("No checkpoint available, using PyTorch's pretrained model")
+            resnet = models.resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
+            model = nn.Sequential(self.normalizer, resnet)
         
         model = model.to(device)
         model.eval()  # Set to evaluation mode
         
+        # Verify model integrity
+        self.verify_model_integrity(model, model_type)
+        
         # Store the model for future use
         self.models[model_type] = model
         return model
-    @GPU
+    
     def inference(self, image, model_type, config):
+        """Run generative inference on the image."""
         # Load model if not already loaded
         model = self.load_model(model_type)
         
@@ -181,12 +513,24 @@ class GenerativeInferenceModel:
         image_tensor = transform(image).unsqueeze(0).to(device)
         image_tensor.requires_grad = True
         
-        # Normalize the image for model input
-        normalized_tensor = normalize_transform(image_tensor)
+        # Check model structure
+        is_sequential = isinstance(model, nn.Sequential)
         
         # Get original predictions
         with torch.no_grad():
-            output_original = model(normalized_tensor)
+            # If the model is sequential with a normalizer, skip the normalization step
+            if is_sequential and isinstance(model[0], NormalizeByChannelMeanStd):
+                print("Model is sequential with normalization")
+                output_original = model(image_tensor)  # Model includes normalization
+                # Get the core model part (typically at index 1 in Sequential)
+                core_model = model[1]
+            else:
+                print("Model is not sequential with normalization")
+                # Use manual normalization for non-sequential models
+                normalized_tensor = normalize_transform(image_tensor)
+                output_original = model(normalized_tensor)
+                core_model = model
+            
             probs_orig = F.softmax(output_original, dim=1)
             conf_orig, classes_orig = torch.max(probs_orig, 1)
             
@@ -197,7 +541,8 @@ class GenerativeInferenceModel:
         infer_step = InferStep(image_tensor, config['eps'], config['step_size'])
         
         # Storage for inference steps
-        x = image_tensor.clone()
+        # Create a new tensor that requires gradients
+        x = image_tensor.clone().detach().requires_grad_(True)
         all_steps = [image_tensor[0].detach().cpu()]
         
         # Main inference loop
@@ -205,32 +550,72 @@ class GenerativeInferenceModel:
             # Reset gradients
             x.grad = None
             
-            # Normalize input for the model
-            normalized_x = normalize_transform(x)
-            
             # Forward pass
-            output = model(normalized_x)
+            if is_sequential and isinstance(model[0], NormalizeByChannelMeanStd):
+                output = model(x)  # Model includes normalization
+            else:
+                # Use manual normalization for non-sequential models
+                normalized_x = normalize_transform(x)
+                output = model(normalized_x)
             
             # Calculate loss to maximize confidence for least confident classes
-            target_classes = least_confident_classes[:10]  # Use top 10 least confident classes
-            loss = 0
-            for idx in target_classes:
-                target = torch.tensor([idx.item()], device=device)
-                loss = loss - F.cross_entropy(output, target)  # Negative because we want to maximize confidence
-            
-            # Backward pass
-            loss.backward()
-            
-            # Update image
-            with torch.no_grad():
-                step = infer_step.step(x, x.grad)
-                x = x + step
+            try:
+                # Get the least confident classes
+                num_classes = min(10, least_confident_classes.size(1))
+                target_classes = least_confident_classes[0, :num_classes]
+                
+                # Create a combined loss (avoid accumulating in a loop)
+                targets = torch.tensor([idx.item() for idx in target_classes], device=device)
+                
+                # Method 1: Use a single combined loss
+                loss = 0
+                for target in targets:
+                    # Create one-hot target
+                    one_hot = torch.zeros_like(output)
+                    one_hot[0, target] = 1
+                    # Use negative loss to maximize confidence
+                    loss = loss + F.mse_loss(F.softmax(output, dim=1), one_hot)
+                
+                # Method 2: Try direct gradient calculation
+                # Instead of loss.backward(), which might be failing
+                grad = torch.autograd.grad(loss, x, retain_graph=True)[0]
+                
+                if grad is None:
+                    print("Warning: Direct gradient calculation failed")
+                    # Fall back to random perturbation
+                    random_noise = (torch.rand_like(x) - 0.5) * 2 * config['step_size']
+                    x = x + random_noise
+                else:
+                    # Update image with gradient
+                    step = infer_step.step(x, grad)
+                    x = x + step
+                
+                x = infer_step.project(x)
+                    
+            except Exception as e:
+                print(f"Error in gradient calculation: {e}")
+                # Fall back to random perturbation
+                random_noise = (torch.rand_like(x) - 0.5) * 2 * config['step_size']
+                x = x + random_noise
                 x = infer_step.project(x)
             
             # Store step if in iterations_to_show
             if i+1 in config['iterations_to_show'] or i+1 == config['n_itr']:
                 all_steps.append(x[0].detach().cpu())
+        
+        # Print some info about the inference
+        with torch.no_grad():
+            if is_sequential and isinstance(model[0], NormalizeByChannelMeanStd):
+                final_output = model(x)
+            else:
+                normalized_x = normalize_transform(x)
+                final_output = model(normalized_x)
                 
+            final_probs = F.softmax(final_output, dim=1)
+            final_conf, final_classes = torch.max(final_probs, 1)
+            print(f"Original top class: {classes_orig.item()} ({conf_orig.item():.4f})")
+            print(f"Final top class: {final_classes.item()} ({final_conf.item():.4f})")
+            
         # Return final image and all stored steps
         return x[0].detach().cpu(), all_steps