DAY43 复习日

浙大疏锦行-CSDN博客

kaggle找到一个图像数据集，用cnn网络进行训练并且用grad-cam做可视化
进阶：把项目拆分成多个文件

src/config.py: 用于存放项目配置，例如文件路径、学习率、批次大小等。

# src/config.py

# Paths
DATA_DIR = "data/"
OUTPUT_DIR = "outputs/"
MODEL_SAVE_PATH = "outputs/models/"
LOG_DIR = "outputs/logs/"
GRAD_CAM_VISUALS_DIR = "outputs/grad_cam_visuals/"

# Dataset parameters
IMAGE_SIZE = (224, 224) # Example image size
BATCH_SIZE = 32

# Training parameters
LEARNING_RATE = 1e-3
NUM_EPOCHS = 20
DEVICE = "cuda" # or "cpu"

# Add other configurations as needed 

src/dataset_utils.py: 用于处理数据集的加载、预处理和数据增强。

# src/dataset_utils.py

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Dataset
# from PIL import Image
# import os

# Example for a custom dataset if your data is not in standard ImageFolder format
# class CustomImageDataset(Dataset):
#     def __init__(self, img_dir, transform=None):
#         self.img_dir = img_dir
#         self.transform = transform
#         self.img_labels = self._load_labels() # Implement this based on your data
#         self.img_paths = self._load_image_paths() # Implement this

#     def __len__(self):
#         return len(self.img_labels)

#     def __getitem__(self, idx):
#         img_path = self.img_paths[idx]
#         image = Image.open(img_path).convert("RGB")
#         label = self.img_labels[idx]
#         if self.transform:
#             image = self.transform(image)
#         return image, label

#     def _load_labels(self):
#         # Implement logic to load or infer labels
#         pass

#     def _load_image_paths(self):
#         # Implement logic to get all image paths
#         pass

def get_transforms(image_size):
    """Define transformations for training and validation data."""
    train_transform = transforms.Compose([
        transforms.Resize(image_size),
        transforms.RandomHorizontalFlip(),
        transforms.RandomRotation(10),
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    
    val_transform = transforms.Compose([
        transforms.Resize(image_size),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    return train_transform, val_transform

def get_dataloaders(data_dir, batch_size, image_size):
    """Creates DataLoader instances for train and validation sets.
    Assumes data is organized in ImageFolder structure (e.g., data/train/classA, data/val/classA).
    Adjust if using a custom dataset.
    """
    train_transform, val_transform = get_transforms(image_size)
    
    # Example for ImageFolder, adjust paths as per your dataset structure
    # train_dataset = datasets.ImageFolder(os.path.join(data_dir, 'train'), transform=train_transform)
    # val_dataset = datasets.ImageFolder(os.path.join(data_dir, 'val'), transform=val_transform)
    
    # Placeholder datasets - replace with actual dataset loading
    print(f"Reminder: Implement actual dataset loading in src/dataset_utils.py based on your Kaggle dataset structure in {data_dir}")
    # For demonstration, creating dummy datasets. Replace these lines.
    num_train_samples = 100 # Placeholder
    num_val_samples = 20   # Placeholder
    num_classes = 2        # Placeholder (e.g., cats vs dogs)
    train_dataset = torch.utils.data.TensorDataset(
        torch.randn(num_train_samples, 3, image_size[0], image_size[1]), 
        torch.randint(0, num_classes, (num_train_samples,))
    )
    val_dataset = torch.utils.data.TensorDataset(
        torch.randn(num_val_samples, 3, image_size[0], image_size[1]), 
        torch.randint(0, num_classes, (num_val_samples,))
    )
    
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True)
    
    return train_loader, val_loader

# Example usage (for testing this module independently):
# if __name__ == '__main__':
#     from config import DATA_DIR, BATCH_SIZE, IMAGE_SIZE
#     # Ensure DATA_DIR points to a valid dataset structured for ImageFolder or adapt CustomImageDataset
#     # For example, DATA_DIR could be 'path/to/dogs-vs-cats/train_subset' if you create such a subset
#     # train_loader, val_loader = get_dataloaders(DATA_DIR, BATCH_SIZE, IMAGE_SIZE)
#     # print(f"Loaded {len(train_loader.dataset)} training images and {len(val_loader.dataset)} validation images.")
#     # for images, labels in train_loader:
#     #     print("Batch shape:", images.shape)
#     #     print("Labels:", labels)
#     #     break 

src/model_arch.py: 用于定义您的CNN模型结构。

# src/model_arch.py

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleCNN(nn.Module):
    def __init__(self, num_classes=2): # Assuming 2 classes (e.g., cats vs dogs)
        super(SimpleCNN, self).__init__()
        # Convolutional Layer 1
        self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1, stride=1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        # Output: 16 x H/2 x W/2 (e.g., 16x112x112 if input 224x224)

        # Convolutional Layer 2
        self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1, stride=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        # Output: 32 x H/4 x W/4 (e.g., 32x56x56)

        # Convolutional Layer 3
        self.conv3 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1, stride=1)
        self.relu3 = nn.ReLU()
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        # Output: 64 x H/8 x W/8 (e.g., 64x28x28)
        
        # Placeholder for the first fully connected layer's input features.
        # This needs to be calculated based on the input image size.
        # For input 224x224, after 3 pooling layers (224/2/2/2 = 28), it becomes 64 * 28 * 28
        # We will store this calculation to be done in forward pass or make it an argument
        self.fc_input_features = None # Will be determined dynamically in forward or init

        # Fully Connected Layers
        # self.fc1 = nn.Linear(self.fc_input_features, 512) # example, fc_input_features needs calculation
        self.fc1 = None # To be defined after knowing input size
        self.relu4 = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        self.fc2 = nn.Linear(512, num_classes) 

        # For Grad-CAM, we need to access activations and gradients easily.
        # We can register hooks or store them if needed, but PyTorch handles grads well.
        # The target layer for Grad-CAM is often the last convolutional layer.
        self.feature_extractor = nn.Sequential(
            self.conv1, self.relu1, self.pool1,
            self.conv2, self.relu2, self.pool2,
            self.conv3, self.relu3, self.pool3
        )

        self.classifier = None # Will combine fc1, relu4, dropout, fc2

    def forward(self, x):
        # x shape: (batch_size, 3, H, W)
        x = self.feature_extractor(x)
        # x shape after feature_extractor: (batch_size, 64, H/8, W/8)

        if self.fc1 is None:
            # Dynamically calculate the input features for the first FC layer
            # This is useful if the image size can vary, though fixed is more common for CNNs.
            batch_size, num_channels, height, width = x.size()
            self.fc_input_features = num_channels * height * width
            self.fc1 = nn.Linear(self.fc_input_features, 512).to(x.device) # Ensure fc1 is on the same device
            # Re-define classifier now that fc1 is initialized
            self.classifier = nn.Sequential(
                self.fc1,
                self.relu4,
                self.dropout,
                self.fc2
            ).to(x.device) # Ensure classifier parts are on the same device
            # Also move fc2 to the correct device if it wasn't part of `self.classifier` initially
            self.fc2 = self.fc2.to(x.device)

        x = x.view(x.size(0), -1)  # Flatten the feature maps
        # x shape: (batch_size, 64 * H/8 * W/8)
        x = self.classifier(x)
        return x

    def get_target_layer_for_grad_cam(self):
        # Typically, the last convolutional layer is used for Grad-CAM.
        # In this model, it's self.conv3 within the feature_extractor sequential block.
        # We need to return the actual layer module, not just its name.
        return self.feature_extractor[-3] # self.conv3 is the 3rd to last item (conv3, relu3, pool3)

# Example: Using a pre-trained model like ResNet18 and modifying its classifier
# from torchvision import models
# def get_pretrained_model(num_classes=2, freeze_features=True):
#     model = models.resnet18(pretrained=True)
#     if freeze_features:
#         for param in model.parameters():
#             param.requires_grad = False
    
#     num_ftrs = model.fc.in_features
#     model.fc = nn.Linear(num_ftrs, num_classes)
#     return model

# Example usage (for testing this module independently):
# if __name__ == '__main__':
#     # Test SimpleCNN
#     dummy_input = torch.randn(4, 3, 224, 224) # Batch of 4, 3 channels, 224x224 image
#     simple_model = SimpleCNN(num_classes=2)
#     output = simple_model(dummy_input)
#     print("SimpleCNN output shape:", output.shape) # Expected: (4, 2)
#     target_layer = simple_model.get_target_layer_for_grad_cam()
#     print("Target layer for Grad-CAM (SimpleCNN):", target_layer)

    # Test Pretrained ResNet18
    # resnet_model = get_pretrained_model(num_classes=2)
    # output_resnet = resnet_model(dummy_input)
    # print("ResNet18 output shape:", output_resnet.shape) # Expected: (4, 2)
    # For ResNet, the target layer for Grad-CAM is usually model.layer4[-1]
    # print("Target layer for Grad-CAM (ResNet18 model.layer4[-1]):", resnet_model.layer4[-1]) 

src/trainer.py: 包含模型训练、验证和保存/加载检查点的逻辑。

# src/trainer.py

import torch
import torch.nn as nn
import torch.optim as optim
# import os
# from .config import DEVICE, MODEL_SAVE_PATH, LOG_DIR # Assuming relative import if config is in the same directory/package
# from .utils import save_checkpoint, load_checkpoint, plot_training_curves # Example utility functions

# Fallback if relative imports fail (e.g. when running script directly)
# try:
#     from .config import DEVICE, MODEL_SAVE_PATH, LOG_DIR
#     from .utils import save_checkpoint, load_checkpoint #, plot_training_curves
# except ImportError:
# print("Running trainer.py directly or utils/config not found with relative import, using direct import for config.")
# This is tricky because cwd matters. For now, let's assume they are imported in main.py
# For independent testing, you might need to adjust sys.path or use absolute imports if src is a package.

class Trainer:
    def __init__(self, model, train_loader, val_loader, config):
        self.model = model.to(config.DEVICE)
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.config = config
        self.criterion = nn.CrossEntropyLoss() # Example loss for classification
        self.optimizer = optim.Adam(self.model.parameters(), lr=config.LEARNING_RATE)
        self.best_val_loss = float('inf')
        self.history = {'train_loss': [], 'val_loss': [], 'train_acc': [], 'val_acc': []}

    def _train_one_epoch(self):
        self.model.train()
        epoch_loss = 0
        correct_predictions = 0
        total_samples = 0

        for batch_idx, (data, target) in enumerate(self.train_loader):
            data, target = data.to(self.config.DEVICE), target.to(self.config.DEVICE)
            
            self.optimizer.zero_grad()
            output = self.model(data)
            loss = self.criterion(output, target)
            loss.backward()
            self.optimizer.step()
            
            epoch_loss += loss.item() * data.size(0)
            _, predicted = torch.max(output.data, 1)
            correct_predictions += (predicted == target).sum().item()
            total_samples += target.size(0)

            if batch_idx % 10 == 0: # Log every 10 batches
                print(f'    Batch {batch_idx+1}/{len(self.train_loader)}, Loss: {loss.item():.4f}')

        avg_epoch_loss = epoch_loss / total_samples
        avg_epoch_acc = correct_predictions / total_samples
        return avg_epoch_loss, avg_epoch_acc

    def _validate_one_epoch(self):
        self.model.eval()
        epoch_loss = 0
        correct_predictions = 0
        total_samples = 0

        with torch.no_grad():
            for data, target in self.val_loader:
                data, target = data.to(self.config.DEVICE), target.to(self.config.DEVICE)
                output = self.model(data)
                loss = self.criterion(output, target)
                
                epoch_loss += loss.item() * data.size(0)
                _, predicted = torch.max(output.data, 1)
                correct_predictions += (predicted == target).sum().item()
                total_samples += target.size(0)
        
        avg_epoch_loss = epoch_loss / total_samples
        avg_epoch_acc = correct_predictions / total_samples
        return avg_epoch_loss, avg_epoch_acc

    def train(self):
        print(f"Starting training for {self.config.NUM_EPOCHS} epochs on {self.config.DEVICE}...")
        for epoch in range(self.config.NUM_EPOCHS):
            print(f"Epoch {epoch+1}/{self.config.NUM_EPOCHS}")
            train_loss, train_acc = self._train_one_epoch()
            val_loss, val_acc = self._validate_one_epoch()
            
            self.history['train_loss'].append(train_loss)
            self.history['train_acc'].append(train_acc)
            self.history['val_loss'].append(val_loss)
            self.history['val_acc'].append(val_acc)
            
            print(f"  Epoch {epoch+1} Summary: Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f} | Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}")

            if val_loss < self.best_val_loss:
                self.best_val_loss = val_loss
                # Create model_save_path directory if it does not exist
                # if not os.path.exists(self.config.MODEL_SAVE_PATH):
                #    os.makedirs(self.config.MODEL_SAVE_PATH)
                # save_checkpoint(self.model, self.optimizer, epoch, os.path.join(self.config.MODEL_SAVE_PATH, 'best_model.pth'))
                print(f"  Saved new best model checkpoint at epoch {epoch+1} (Val Loss: {val_loss:.4f})")
                print(f"  Reminder: Implement save_checkpoint in utils.py and uncomment relevant lines in trainer.py")
        
        print("Training finished.")
        # plot_training_curves(self.history, os.path.join(self.config.LOG_DIR, 'training_curves.png'))
        print(f"  Reminder: Implement plot_training_curves in utils.py and uncomment relevant lines in trainer.py")
        return self.history

# Example usage (for testing this module independently):
# if __name__ == '__main__':
#     # This requires config, model_arch, dataset_utils to be importable
#     # You might need to add `src` to PYTHONPATH or run from the project root: python -m src.trainer
#     import sys
#     sys.path.append('..') # Adds project root to path if running from src/
#     from config import BATCH_SIZE, IMAGE_SIZE, DATA_DIR, DEVICE, LEARNING_RATE, NUM_EPOCHS, MODEL_SAVE_PATH, LOG_DIR
#     from dataset_utils import get_dataloaders
#     from model_arch import SimpleCNN

#     # Create a dummy config object if not importing from a central config file
#     class DummyConfig:
#         DEVICE = DEVICE
#         LEARNING_RATE = LEARNING_RATE
#         NUM_EPOCHS = 3 # Shorten for test
#         MODEL_SAVE_PATH = MODEL_SAVE_PATH
#         LOG_DIR = LOG_DIR
#         # Add any other config attributes used by Trainer

#     cfg = DummyConfig()
    
#     train_loader, val_loader = get_dataloaders(DATA_DIR, BATCH_SIZE, IMAGE_SIZE)
#     model = SimpleCNN(num_classes=2) # Assuming 2 classes for the dummy data in dataset_utils
    
#     trainer = Trainer(model, train_loader, val_loader, cfg)
#     trainer.train()
#     print("Trainer test finished.") 

src/grad_cam.py: 实现Grad-CAM算法。

# src/grad_cam.py

import torch
import torch.nn.functional as F
import cv2 # OpenCV for image manipulation
import numpy as np
# import matplotlib.pyplot as plt # For displaying

class GradCAM:
    def __init__(self, model, target_layer):
        self.model = model
        self.target_layer = target_layer
        self.gradients = None
        self.activations = None
        
        self._register_hooks()

    def _register_hooks(self):
        def backward_hook(module, grad_input, grad_output):
            self.gradients = grad_output[0]
            return None

        def forward_hook(module, input, output):
            self.activations = output
            return None

        self.target_layer.register_forward_hook(forward_hook)
        self.target_layer.register_backward_hook(backward_hook)

    def _get_cam_weights(self, grads):
        # Global Average Pooling of gradients
        return torch.mean(grads, dim=[2, 3], keepdim=True)

    def generate_cam(self, input_tensor, class_idx=None, use_relu=True):
        """
        Generates the Class Activation Map (CAM).
        Args:
            input_tensor: Preprocessed input image tensor (usually 1xCxHxW).
            class_idx: Index of the class to generate CAM for. If None, uses the predicted class.
            use_relu: Whether to apply ReLU to the CAM to keep only positive contributions.
        Returns:
            cam_heatmap: A numpy array of the CAM, normalized to [0, 1].
        """
        self.model.eval() # Ensure model is in eval mode
        self.model.zero_grad() # Zero out gradients

        # Forward pass to get activations
        output = self.model(input_tensor)
        # output shape: (1, num_classes)

        if class_idx is None:
            class_idx = torch.argmax(output, dim=1).item()

        # Backward pass to get gradients
        # Create a one-hot like score for the target class
        score = output[0, class_idx]
        score.backward()

        if self.gradients is None or self.activations is None:
            raise RuntimeError("Gradients or activations not found. Hooks might not have been called.")

        # Get weights from gradients (Global Average Pooling)
        # self.gradients shape: (batch_size_usually_1, num_channels_target_layer, H_feature, W_feature)
        weights = self._get_cam_weights(self.gradients)
        # weights shape: (1, num_channels_target_layer, 1, 1)

        # Weighted sum of activations
        # self.activations shape: (batch_size_usually_1, num_channels_target_layer, H_feature, W_feature)
        cam = torch.sum(weights * self.activations, dim=1, keepdim=True)
        # cam shape: (1, 1, H_feature, W_feature)

        if use_relu:
            cam = F.relu(cam)

        # Resize CAM to original image size
        # input_tensor shape: (1, C, H_orig, W_orig)
        # cam shape before upsample: (1, 1, H_feature, W_feature)
        cam = F.interpolate(cam, size=(input_tensor.shape[2], input_tensor.shape[3]), mode='bilinear', align_corners=False)
        # cam shape after upsample: (1, 1, H_orig, W_orig)

        # Normalize CAM
        cam = cam.squeeze().cpu().detach().numpy()
        cam = cam - np.min(cam)
        cam = cam / (np.max(cam) + 1e-8) # Add epsilon to avoid division by zero
        return cam

def overlay_cam_on_image(original_img_np, cam_heatmap, alpha=0.5):
    """
    Overlays the CAM heatmap on the original image.
    Args:
        original_img_np: Original image as a NumPy array (H, W, C) in range [0, 255].
        cam_heatmap: CAM heatmap (H, W) normalized to [0, 1].
        alpha: Transparency factor for the heatmap.
    Returns:
        overlaid_image: NumPy array of the original image with heatmap overlay.
    """
    heatmap_colored = cv2.applyColorMap(np.uint8(255 * cam_heatmap), cv2.COLORMAP_JET)
    heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB) # Matplotlib uses RGB
    
    if original_img_np.shape[:2] != heatmap_colored.shape[:2]:
        heatmap_colored = cv2.resize(heatmap_colored, (original_img_np.shape[1], original_img_np.shape[0]))

    overlaid_image = cv2.addWeighted(np.uint8(original_img_np), 1 - alpha, heatmap_colored, alpha, 0)
    return overlaid_image

# Example usage (for testing this module independently):
# if __name__ == '__main__':
#     import sys
#     sys.path.append('..') # Adds project root to path if running from src/
#     from model_arch import SimpleCNN # Assuming SimpleCNN is defined
#     from torchvision import transforms
#     from PIL import Image
#     from config import IMAGE_SIZE

#     # 1. Load a model (ensure it's trained or has some weights for meaningful CAM)
#     num_classes = 2 # Example
#     model = SimpleCNN(num_classes=num_classes)
#     # model.load_state_dict(torch.load('path_to_your_trained_model.pth')) # Load trained weights
#     model.eval()

#     # 2. Get the target layer for Grad-CAM
#     target_layer = model.get_target_layer_for_grad_cam()
#     if target_layer is None:
#         raise ValueError("Target layer not found in model. Ensure get_target_layer_for_grad_cam() is implemented.")

#     # 3. Initialize GradCAM
#     grad_cam = GradCAM(model=model, target_layer=target_layer)

#     # 4. Load and preprocess an image
#     # Replace with your image path and appropriate transformations
#     # try:
#     #     img = Image.open("path_to_your_image.jpg").convert('RGB')
#     # except FileNotFoundError:
#     #     print("Test image not found. Grad-CAM test will use a dummy image.")
#     #     img = Image.fromarray(np.uint8(np.random.rand(224, 224, 3) * 255)) # Dummy image
    
#     # Use a dummy image for now as we don't have a specific image path
#     img = Image.fromarray(np.uint8(np.random.rand(IMAGE_SIZE[0], IMAGE_SIZE[1], 3) * 255)) # Dummy image

#     # Preprocessing transformations (should match training)
#     preprocess = transforms.Compose([
#         transforms.Resize(IMAGE_SIZE),
#         transforms.ToTensor(),
#         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
#     ])
#     input_tensor = preprocess(img).unsqueeze(0) # Add batch dimension

#     # 5. Generate CAM
#     # You might want to specify class_idx if you know it, otherwise it uses the predicted class
#     cam_heatmap = grad_cam.generate_cam(input_tensor, class_idx=None) 

#     # 6. Overlay CAM on original image (after denormalizing and converting to displayable format)
#     # To properly display original_img_np, it should be denormalized if normalization was applied
#     # For simplicity, we use the 'img' before ToTensor and normalization for visualization here
#     original_img_np = np.array(img.resize((IMAGE_SIZE[1], IMAGE_SIZE[0]))) # H, W, C
#     if original_img_np.max() <= 1.0: # If ToTensor was already applied and values are 0-1
#         original_img_np = (original_img_np * 255).astype(np.uint8)
    
#     overlaid_image = overlay_cam_on_image(original_img_np, cam_heatmap)

#     # 7. Display (Optional - requires matplotlib)
#     # fig, ax = plt.subplots(1, 2, figsize=(10, 5))
#     # ax[0].imshow(original_img_np)
#     # ax[0].set_title("Original Image")
#     # ax[0].axis('off')
#     # ax[1].imshow(overlaid_image)
#     # ax[1].set_title("Grad-CAM Overlay")
#     # ax[1].axis('off')
#     # plt.show()
#     # print("Grad-CAM generated. If matplotlib is installed and un-commented, an image should display.")
#     # To save the image instead:
#     # cv2.imwrite("outputs/grad_cam_visuals/sample_grad_cam.jpg", cv2.cvtColor(overlaid_image, cv2.COLOR_RGB2BGR))
#     print("Grad-CAM test finished. `overlaid_image` contains the result.")
#     print("Reminder: Implement saving the visual in main.py or notebook, and ensure output directories exist.") 

src/utils.py: 存放一些辅助函数，例如保存/加载模型、绘图等。

# src/utils.py

import torch
import os
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image

def save_checkpoint(model, optimizer, epoch, filepath):
    """Saves model checkpoint."""
    # Ensure directory exists
    os.makedirs(os.path.dirname(filepath), exist_ok=True)
    
    checkpoint = {
        'epoch': epoch,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
    }
    torch.save(checkpoint, filepath)
    print(f"Checkpoint saved to {filepath}")

def load_checkpoint(model, optimizer, filepath, device):
    """Loads model checkpoint."""
    if not os.path.exists(filepath):
        print(f"Checkpoint file not found: {filepath}")
        return None, None # Or raise an error
        
    checkpoint = torch.load(filepath, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])
    if optimizer and 'optimizer_state_dict' in checkpoint:
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    epoch = checkpoint.get('epoch', 0) # Default to epoch 0 if not found
    print(f"Checkpoint loaded from {filepath}, epoch {epoch}")
    return model, optimizer, epoch

def plot_training_curves(history, save_path):
    """Plots training and validation loss and accuracy curves."""
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    
    fig, axes = plt.subplots(2, 1, figsize=(10, 8))
    
    epochs = range(1, len(history['train_loss']) + 1)
    
    # Plot Loss
    axes[0].plot(epochs, history['train_loss'], 'bo-', label='Training Loss')
    axes[0].plot(epochs, history['val_loss'], 'ro-', label='Validation Loss')
    axes[0].set_title('Training and Validation Loss')
    axes[0].set_xlabel('Epochs')
    axes[0].set_ylabel('Loss')
    axes[0].legend()
    axes[0].grid(True)
    
    # Plot Accuracy
    axes[1].plot(epochs, history['train_acc'], 'bo-', label='Training Accuracy')
    axes[1].plot(epochs, history['val_acc'], 'ro-', label='Validation Accuracy')
    axes[1].set_title('Training and Validation Accuracy')
    axes[1].set_xlabel('Epochs')
    axes[1].set_ylabel('Accuracy')
    axes[1].legend()
    axes[1].grid(True)
    
    plt.tight_layout()
    plt.savefig(save_path)
    print(f"Training curves saved to {save_path}")
    plt.close(fig)

def denormalize_image(tensor, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
    """Denormalizes a tensor image with mean and standard deviation."""
    mean = torch.tensor(mean).view(3, 1, 1)
    std = torch.tensor(std).view(3, 1, 1)
    if tensor.is_cuda:
        mean = mean.to(tensor.device)
        std = std.to(tensor.device)
    tensor = tensor * std + mean
    tensor = torch.clamp(tensor, 0, 1) # Clamp values to [0, 1]
    return tensor

def tensor_to_pil(tensor_image):
    """Converts a PyTorch tensor (C, H, W) to a PIL Image."""
    # Denormalize if needed, assuming tensor is already in [0,1] range or denormalized
    # If the tensor is (B, C, H, W), take the first image
    if tensor_image.ndim == 4 and tensor_image.shape[0] == 1:
        tensor_image = tensor_image.squeeze(0)
        
    pil_image = transforms.ToPILImage()(tensor_image.cpu())
    return pil_image

# You might want to add more utility functions here, e.g., for specific dataset parsing, advanced logging, etc. 

outputs 文件夹及其子文件夹 (models, logs, grad_cam_visuals) 已经创建好了。
接下来是 notebooks 文件夹和里面的实验性 Jupyter Notebook。

main.py: 主脚本，用于整合和运行整个流程：数据准备、模型训练、评估和可选的Grad-CAM可视化。

# main.py

import torch
import os
from PIL import Image
import numpy as np
import cv2 # For saving Grad-CAM images

# It's good practice to set up argument parsing for a main script
import argparse

# Assuming src is in PYTHONPATH or running from project root
from src import config
from src import dataset_utils
from src import model_arch
from src import trainer
from src import grad_cam as gradcam_module # Renamed to avoid conflict with a potential variable name
from src import utils

def run_training(cfg):
    """Runs the training process."""
    print("--- Starting Training Phase ---")
    train_loader, val_loader = dataset_utils.get_dataloaders(
        data_dir=cfg.DATA_DIR,
        batch_size=cfg.BATCH_SIZE,
        image_size=cfg.IMAGE_SIZE
    )

    # Determine num_classes from the dataset if possible, or set in config
    # For dummy data in dataset_utils, it's 2. For real data, this needs to be accurate.
    num_classes = 2 # TODO: Make this configurable or infer from data
    print(f"Number of classes: {num_classes}")

    model = model_arch.SimpleCNN(num_classes=num_classes).to(cfg.DEVICE)
    # model = model_arch.get_pretrained_model(num_classes=num_classes).to(cfg.DEVICE) # Alternative
    print(f"Using model: {model.__class__.__name__}")

    model_trainer = trainer.Trainer(model=model, 
                                    train_loader=train_loader, 
                                    val_loader=val_loader, 
                                    config=cfg)
    
    history = model_trainer.train()
    utils.plot_training_curves(history, os.path.join(cfg.LOG_DIR, 'training_curves_main.png'))
    
    print("--- Training Phase Finished ---")
    return model, model_trainer.optimizer # Return optimizer for potential checkpoint loading

def run_grad_cam_visualization(cfg, model, num_samples=5):
    """Runs Grad-CAM visualization on a few samples from the validation set."""
    print("--- Starting Grad-CAM Visualization Phase ---")
    model.eval()
    
    target_layer = model.get_target_layer_for_grad_cam()
    if target_layer is None:
        print("Could not get target layer for Grad-CAM. Skipping visualization.")
        return

    grad_cam_visualizer = gradcam_module.GradCAM(model=model, target_layer=target_layer)
    
    # Get a few images from val_loader (or a specific test set)
    _, val_loader = dataset_utils.get_dataloaders(
        data_dir=cfg.DATA_DIR, 
        batch_size=num_samples, # Get a batch of `num_samples`
        image_size=cfg.IMAGE_SIZE
    )
    
    try:
        images, labels = next(iter(val_loader))
    except StopIteration:
        print("Validation loader is empty. Cannot generate Grad-CAM visuals.")
        # Fallback to dummy image if val_loader is empty for some reason
        print("Using a dummy image for Grad-CAM demonstration as val_loader was empty.")
        images = torch.randn(num_samples, 3, cfg.IMAGE_SIZE[0], cfg.IMAGE_SIZE[1])
        labels = torch.zeros(num_samples) # Dummy labels

    images = images.to(cfg.DEVICE)

    # Ensure the visualization directory exists
    os.makedirs(cfg.GRAD_CAM_VISUALS_DIR, exist_ok=True)

    for i in range(min(num_samples, images.size(0))):
        img_tensor = images[i].unsqueeze(0) # Process one image at a time (needs batch dim)
        original_pil = utils.tensor_to_pil(utils.denormalize_image(images[i].cpu()))
        original_np = np.array(original_pil)

        # Generate CAM (use predicted class by default)
        cam_heatmap = grad_cam_visualizer.generate_cam(img_tensor, class_idx=None)
        overlaid_image = gradcam_module.overlay_cam_on_image(original_np, cam_heatmap, alpha=0.6)

        # Save the image
        predicted_class = torch.argmax(model(img_tensor), dim=1).item()
        actual_label = labels[i].item() if labels is not None else 'N/A'
        save_filename = f"grad_cam_sample_{i}_pred_{predicted_class}_actual_{actual_label:.0f}.png"
        save_path = os.path.join(cfg.GRAD_CAM_VISUALS_DIR, save_filename)
        
        # OpenCV expects BGR, PIL/Matplotlib use RGB
        cv2.imwrite(save_path, cv2.cvtColor(overlaid_image, cv2.COLOR_RGB2BGR))
        print(f"Saved Grad-CAM visualization to {save_path}")

    print("--- Grad-CAM Visualization Phase Finished ---")

def main():
    parser = argparse.ArgumentParser(description='CNN Image Classification and Grad-CAM Visualization')
    parser.add_argument('--mode', type=str, default='train_gradcam', choices=['train_only', 'gradcam_only', 'train_gradcam'], 
                        help='Run mode: train_only, gradcam_only, or train_gradcam')
    parser.add_argument('--model_load_path', type=str, default=None, 
                        help='Path to load a pre-trained model checkpoint for gradcam_only mode or to resume training.')
    # Add more CLI arguments as needed (e.g., overriding config parameters)
    args = parser.parse_args()

    # Create necessary output directories defined in config
    os.makedirs(config.MODEL_SAVE_PATH, exist_ok=True)
    os.makedirs(config.LOG_DIR, exist_ok=True)
    os.makedirs(config.GRAD_CAM_VISUALS_DIR, exist_ok=True)

    trained_model = None
    optimizer_state = None

    if args.mode in ['train_only', 'train_gradcam']:
        # If a model_load_path is provided for training, it implies resuming training
        # For simplicity, this example does not explicitly implement resuming optimizer state from checkpoint in main
        # but the utils.load_checkpoint function supports it.
        if args.model_load_path and os.path.exists(args.model_load_path):
            print(f"Resuming training or using pre-trained model from: {args.model_load_path}")
            # This part needs careful handling of num_classes if model is pre-trained
            num_classes = 2 # TODO: Get this from config or model checkpoint
            model_shell = model_arch.SimpleCNN(num_classes=num_classes).to(config.DEVICE) # Create a shell model
            # Create a dummy optimizer shell; its state will be loaded if available in checkpoint
            dummy_optimizer = torch.optim.Adam(model_shell.parameters(), lr=config.LEARNING_RATE)
            loaded_model, _, _ = utils.load_checkpoint(model_shell, dummy_optimizer, args.model_load_path, config.DEVICE)
            if loaded_model:
                 trained_model = loaded_model # This model is now ready for further training or Grad-CAM
            else:
                print(f"Failed to load model from {args.model_load_path}. Starting fresh training.")
                trained_model, _ = run_training(config) # optimizer state also returned if needed
        else:
            trained_model, _ = run_training(config)
    
    if args.mode in ['gradcam_only', 'train_gradcam']:
        if args.mode == 'gradcam_only':
            if not args.model_load_path or not os.path.exists(args.model_load_path):
                print("Error: For 'gradcam_only' mode, a valid --model_load_path must be provided.")
                return
            # Load the model for Grad-CAM
            num_classes = 2 # TODO: Get this from config or model checkpoint
            model_to_visualize = model_arch.SimpleCNN(num_classes=num_classes).to(config.DEVICE)
            # Optimizer is not strictly needed for visualization only, pass None
            loaded_model, _, _ = utils.load_checkpoint(model_to_visualize, None, args.model_load_path, config.DEVICE)
            if not loaded_model:
                print(f"Error: Failed to load model from {args.model_load_path} for Grad-CAM.")
                return
            trained_model = loaded_model
        
        if trained_model is None:
            print("Error: Model not available for Grad-CAM. Train a model first or provide a valid path.")
            return
            
        run_grad_cam_visualization(config, trained_model, num_samples=5)

if __name__ == '__main__':
    print(f"PyTorch Version: {torch.__version__}")
    print(f"OpenCV Version: {cv2.__version__}")
    main() 

main.py 脚本已经创建完毕。它包含以下功能：
参数解析: 使用 argparse 来处理命令行参数，例如运行模式 (train_only, gradcam_only, train_gradcam) 和预训练模型的路径。
训练流程 (run_training): 初始化数据加载器、模型、训练器，然后开始训练，并在结束后绘制训练曲线。
Grad-CAM 可视化流程 (run_grad_cam_visualization): 加载训练好的模型，从验证集中选取样本图像，生成并保存Grad-CAM叠加图。
主函数 (main): 根据命令行参数协调执行训练和/或Grad-CAM可视化。

我最近在学习深度学习，现在有个任务：kaggle找到一个图像数据集，用cnn网络进行训练并且用grad-cam做可视化
进阶：把项目拆分成多个文件
并且请你新建一个文件夹，这些文件按照规则存放在该文件夹中

好的，这个任务非常有意思！我们可以分几步来完成：
选择一个Kaggle图像数据集：Kaggle上有许多适合图像分类的数据集。一个常见的例子是 “猫狗分类” (Cats and Dogs)。如果您没有特定的数据集，我们可以先用这个作为例子。您需要先从Kaggle下载数据集。通常可以使用Kaggle API来下载，例如：kaggle competitions download -c dogs-vs-cats (这只是一个例子，具体命令取决于数据集)。下载后，需要解压到我们稍后创建的 data 文件夹中。
项目结构：我会为您创建一个新的项目文件夹，并在其中组织好代码文件和数据。
代码实现：我们会逐步创建CNN模型、数据加载、训练逻辑和Grad-CAM可视化部分。
我们先来创建项目文件夹结构。我建议的文件夹结构如下：
Apply
您觉得这个结构怎么样？如果可以，我现在就为您创建这些文件夹和空的 Python 文件。
首先，我将在您的工作区 D:/Apython_study/ 中创建一个名为 kaggle_cnn_gradcam 的新文件夹。
path should be a path.relative()d string, but got “D:/Apython_study/kaggle_cnn_gradcam/README.md”
抱歉，刚才创建文件时路径有点问题。我来修正一下。
首先，我将在您的工作区 D:/Apython_study/ 中创建一个名为 kaggle_cnn_gradcam 的新文件夹，并在其中添加一个 README.md 文件。
很好，README.md 和项目主文件夹 kaggle_cnn_gradcam 已经创建好了。
接下来，我们来创建 requirements.txt 文件，列出一些常用的库。请注意，这些库和版本是初步的，后续根据实际代码可能会调整。
requirements.txt 文件已经创建。
现在，我们来创建规划好的子文件夹和 Python 脚本文件。我会在每个 Python 文件中添加一个简短的注释来说明其用途。
创建 data 文件夹: 这个文件夹需要您手动将下载的Kaggle数据集（例如解压后的图片）放进去。为了让这个文件夹在git中能够被追踪（如果需要的话），我会在里面创建一个空的 .gitkeep 文件。
data 文件夹已准备好。
接下来是 src 文件夹及其内部的 Python 文件：
src/config.py: 用于存放项目配置，例如文件路径、学习率、批次大小等。
src/config.py 已创建。
src/dataset_utils.py: 用于处理数据集的加载、预处理和数据增强。
src/dataset_utils.py 已创建。请注意，您需要根据您选择的Kaggle数据集的实际文件结构来调整其中的数据加载逻辑。
src/model_arch.py: 用于定义您的CNN模型结构。
src/model_arch.py 已创建。其中包含一个简单的三层CNN模型 (SimpleCNN)，并动态计算全连接层的输入大小。我还添加了一个获取Grad-CAM目标层的方法，并注释掉了使用预训练模型 (如ResNet) 的示例代码。
src/trainer.py: 包含模型训练、验证和保存/加载检查点的逻辑。
src/trainer.py 已创建。其中包含一个 Trainer 类，封装了训练和验证的循环，以及记录损失和准确率的逻辑。保存检查点和绘制训练曲线的功能已经留出位置，但需要您在 utils.py 中实现具体函数并取消注释相关行。
src/grad_cam.py: 实现Grad-CAM算法。
src/grad_cam.py 已创建。这个文件里包含了一个 GradCAM 类，它通过注册钩子 (hooks) 来获取目标层的激活值和梯度，然后计算并生成CAM热力图。同时还提供了一个 overlay_cam_on_image 函数，用于将热力图叠加到原始图像上。
src/utils.py: 存放一些辅助函数，例如保存/加载模型、绘图等。
src/utils.py 已创建，其中包含了保存/加载检查点、绘制训练曲线图、图像反归一化以及将张量转换为PIL图像等辅助函数。我已添加了之前遗漏的 from torchvision import transforms 导入。
现在，我们来创建 outputs 文件夹及其子文件夹，用于存放训练过程中的产出物。这些文件夹通常在代码运行时按需创建，但我们可以先创建空的 .gitkeep 文件来占位。
outputs/models/
outputs/logs/
outputs/grad_cam_visuals/
outputs 文件夹及其子文件夹 (models, logs, grad_cam_visuals) 已经创建好了。
接下来是 notebooks 文件夹和里面的实验性 Jupyter Notebook。
notebooks/experiment.ipynb: 一个Jupyter Notebook，可以用来逐步执行数据加载、模型训练、评估和Grad-CAM可视化的整个流程。
notebooks/experiment.ipynb 骨架文件已创建。这个notebook提供了一个引导式的环境，您可以用来运行和调试代码的各个部分。
最后，我们创建 main.py 文件，这将是项目的主执行脚本。
main.py: 主脚本，用于整合和运行整个流程：数据准备、模型训练、评估和可选的Grad-CAM可视化。
main.py 脚本已经创建完毕。它包含以下功能：
参数解析: 使用 argparse 来处理命令行参数，例如运行模式 (train_only, gradcam_only, train_gradcam) 和预训练模型的路径。
训练流程 (run_training): 初始化数据加载器、模型、训练器，然后开始训练，并在结束后绘制训练曲线。
Grad-CAM 可视化流程 (run_grad_cam_visualization): 加载训练好的模型，从验证集中选取样本图像，生成并保存Grad-CAM叠加图。
主函数 (main): 根据命令行参数协调执行训练和/或Grad-CAM可视化。
至此，整个项目的骨架结构和基础代码模块已经创建完毕。文件夹结构如下：
Apply
接下来的步骤：
下载Kaggle数据集: 您需要选择一个Kaggle图像分类数据集（例如猫狗分类），下载它，并解压到 kaggle_cnn_gradcam/data/ 文件夹中。
调整数据加载: 打开 src/dataset_utils.py，根据您下载的数据集的具体文件结构和格式，修改 get_dataloaders 函数。您可能需要使用 torchvision.datasets.ImageFolder (如果数据是按类别分文件夹存放的) 或实现一个自定义的 Dataset 类。确保 num_classes (类别的数量) 在 main.py (或从config加载) 和模型定义中设置正确。
配置调整: 检查 src/config.py 中的设置（如 IMAGE_SIZE, BATCH_SIZE, LEARNING_RATE, NUM_EPOCHS, DEVICE），并根据您的需求和硬件进行调整。
安装依赖: 在您的项目环境中，通过 pip install -r requirements.txt 安装所需的库。
运行项目:
您可以使用 notebooks/experiment.ipynb 来逐步执行和调试代码。
或者直接运行主脚本：python main.py --mode train_gradcam (或其他模式)。