utils.py

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets
from tqdm import tqdm
import numpy as np
from sklearn.metrics import confusion_matrix

# from color_transforms import apply_color_transform, apply_random_color
from pathlib import Path
import datetime
import logging
import argparse
import matplotlib.pyplot as plt
import seaborn as sns
from typing import Tuple
import os
import glob
from collections import defaultdict
from typing import Dict, Tuple
import models

# from opts import OPT as opt
import random


def set_random_seed(seed: int) -> None:
    """Set random seed for reproducibility across all libraries"""
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # for multi-GPU
    torch.backends.cudnn.enabled = True
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(seed)
    # seed the python.
    os.environ["PYTHONSEED"] = str(seed)
    # seed everything about the program.
    random.seed(seed)


def create_run_dir(args, method_names) -> Path:
    """Create a uniquely named directory for this run"""
    # Get timestamp
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")

    # Create run name
    # method_names = {0: "vanilla", 1: "noise", 2: "kl", 3: "kl1"}
    method_name = method_names[args.method]
    if args.mode == "CR":
        # Format: output/YYYYMMDD_HHMMSS_method_forgetN
        run_name = f"{timestamp}_{method_name}_{args.dataset}_forget{args.forget_class}"
    elif args.mode == "HR":
        run_name = f"HR/{timestamp}_{method_name}_{args.dataset}"
    if args.eval_only:
        run_name = f"{timestamp}_eval_only"

    # Create directory structure
    run_dir = Path(args.output_dir) / run_name
    run_dir.mkdir(parents=True, exist_ok=True)

    # Create subdirectories
    (run_dir / "models").mkdir(exist_ok=True)
    (run_dir / "logs").mkdir(exist_ok=True)

    return run_dir


def create_run_color_dir(args, mode) -> Path:
    """Create a uniquely named directory for this run"""
    # Get timestamp
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")

    # Create run name
    method_names = {0: "vanilla", 1: "noise", 2: "kl"}
    method_name = method_names[args.method]

    # Format: output/YYYYMMDD_HHMMSS_method_forgetN
    run_name = f"{timestamp}_{method_name}_{mode}"
    if args.eval_only:
        run_name = f"{timestamp}_eval_only"

    # Create directory structure
    run_dir = Path(args.output_dir) / run_name
    run_dir.mkdir(parents=True, exist_ok=True)

    # Create subdirectories
    (run_dir / "models").mkdir(exist_ok=True)
    (run_dir / "logs").mkdir(exist_ok=True)

    return run_dir


def plot_confusion_matrix(cm, title, save_path):
    """Plot and save confusion matrix"""
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues")
    plt.title(title)
    plt.ylabel("True Label")
    plt.xlabel("Predicted Label")
    plt.savefig(save_path, bbox_inches="tight", pad_inches=0.1)
    plt.close()


def plot_accuracy_comparison(orig_acc, unl_acc, save_path, forget_class):
    """Plot and save accuracy comparison bar chart"""
    plt.figure(figsize=(12, 6))
    x = np.arange(10)
    width = 0.35

    plt.bar(x - width / 2, orig_acc, width, label="Original Model")
    plt.bar(x + width / 2, unl_acc, width, label="Unlearned Model")

    plt.xlabel("Digit")
    plt.ylabel("Accuracy (%)")
    plt.title(f"Per-class Accuracy Comparison (forget_class={forget_class})")
    plt.xticks(x)
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.savefig(save_path, bbox_inches="tight", pad_inches=0.1)
    plt.close()


def setup_logging(run_dir: Path) -> None:
    """Configure logging to file and console for each run"""
    log_file = run_dir / "logs/run.log"

    # Get root logger
    logger = logging.getLogger()
    logger.setLevel(logging.INFO)

    # Remove existing handlers
    for handler in logger.handlers[:]:
        logger.removeHandler(handler)

    # Create formatter
    formatter = logging.Formatter("%(asctime)s - %(message)s")

    # File handler
    file_handler = logging.FileHandler(log_file)
    file_handler.setFormatter(formatter)

    # Console handler
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(formatter)

    # Add handlers
    logger.addHandler(file_handler)
    logger.addHandler(console_handler)


def plot_accuracy_comparison_experiments(orig_acc, unl_acc, save_path):
    """绘制不同实验之间原模型和unlearn模型准确率的对比柱状图"""
    plt.figure(figsize=(12, 6))
    modes = list(orig_acc.keys())
    x = np.arange(len(modes))
    width = 0.35

    orig_values = [orig_acc[mode] for mode in modes]
    unl_values = [unl_acc[mode] for mode in modes]

    orig_bars = plt.bar(x - width / 2, orig_values, width, label="Original Model")
    unl_bars = plt.bar(x + width / 2, unl_values, width, label="Unlearned Model")

    plt.xlabel("color_mode")
    plt.ylabel("Accuracy (%)")
    plt.title(f"Comparations between original and unlearned model")
    plt.xticks(x, modes)
    plt.legend()
    plt.grid(True, alpha=0.3)

    # 在每个柱子上添加数据标签
    for bar in orig_bars:
        height = bar.get_height()
        plt.text(
            bar.get_x() + bar.get_width() / 2.0,
            height,
            f"{height:.2f}%",
            ha="center",
            va="bottom",
            fontsize=10,
        )
    for bar in unl_bars:
        height = bar.get_height()
        plt.text(
            bar.get_x() + bar.get_width() / 2.0,
            height,
            f"{height:.2f}%",
            ha="center",
            va="bottom",
            fontsize=10,
        )

    plt.savefig(save_path, bbox_inches="tight", pad_inches=0.1)
    plt.close()


def load_model(model_path, num_classes=10):
    model = models.ResNet18()
    model.conv1 = nn.Conv2d(
        in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1, bias=False
    )
    model.fc = torch.nn.Linear(512, num_classes)
    model.load_state_dict(torch.load(model_path, weights_only=True))

    return model


def find_model_path(seed: int, mode, method: str, class_id: int) -> str:
    """通过特征匹配查找模型路径"""
    if mode == "CR":
        base_dir = f"../outputs/seed{seed}/cifar10/{method}"
        if not os.path.exists(base_dir):
            return None

        target_suffix = f"{method}_forget{class_id}"
        candidate_dirs = []
        # 遍历所有子目录
        for dir_name in os.listdir(base_dir):
            dir_path = os.path.join(base_dir, dir_name)

            # 验证是否为有效目录
            if not os.path.isdir(dir_path):
                continue

            # 分割目录名特征
            parts = dir_name.split("_")
            if len(parts) < 2:
                continue

            # 匹配方法名和遗忘类
            if "_".join(parts[-2:]) == target_suffix:
                candidate_dirs.append(dir_path)

        # 按目录名排序（假设目录名含时间戳）
        candidate_dirs.sort(reverse=True)

        # 检查模型文件
        for dir_path in candidate_dirs:
            model_path = os.path.join(
                dir_path, "models", f"cifar10_model_best_{method}_forget{class_id}.pth"
            )
            if os.path.exists(model_path):
                return model_path
    elif mode == "HR":
        base_dir = f"../outputs_seed{seed}/HR"
        if not os.path.exists(base_dir):
            return None

        target_suffix = f"{method}"
        candidate_dirs = []
        # 遍历所有子目录
        for dir_name in os.listdir(base_dir):
            dir_path = os.path.join(base_dir, dir_name)

            # 验证是否为有效目录
            if not os.path.isdir(dir_path):
                continue

            # 分割目录名特征
            parts = dir_name.split("_")
            if len(parts) < 2:
                continue

            # 匹配方法名和遗忘类
            if "_".join(parts[-1:]) == target_suffix:
                candidate_dirs.append(dir_path)

        # 按目录名排序（假设目录名含时间戳）
        # candidate_dirs.sort(reverse=True)

        # 检查模型文件
        for dir_path in candidate_dirs:
            model_path = os.path.join(
                dir_path, "models", f"cifar10_model_best_{method}.pth"
            )
            if os.path.exists(model_path):
                return model_path

    return None


def calculate_differences(avg_results: Dict) -> Tuple[Dict, np.ndarray]:
    """计算方法间的差异"""
    # 方法间两两差异（按类别）
    class_diffs = defaultdict(dict)

    # 方法间整体差异矩阵
    methods = list(avg_results.keys())
    n_methods = len(methods)
    diff_matrix = np.zeros((n_methods, n_methods))

    for i, m1 in enumerate(methods):
        for j, m2 in enumerate(methods):
            if i >= j:
                continue
            # 按类别计算差异
            class_diffs[(m1, m2)] = {
                c: avg_results[m1][c] - avg_results[m2][c] for c in avg_results[m1]
            }
            # 计算整体平均差异
            diff_matrix[i, j] = np.mean(list(class_diffs[(m1, m2)].values()))
            diff_matrix[j, i] = -diff_matrix[i, j]

    return class_diffs, diff_matrix


def accuracy(net, loader, single_class=False):
    """Return accuracy on a dataset given by the data loader."""
    correct = 0
    total = 0

    total_sc = torch.zeros((opt.num_classes))
    correct_sc = torch.zeros((opt.num_classes))

    pred_all = []
    target_all = []

    for inputs, targets in loader:
        inputs, targets = inputs.to(opt.device), targets.to(opt.device)
        outputs = net(inputs)
        _, predicted = outputs.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
        if single_class:
            pred_all.append(predicted.detach().cpu())
            target_all.append(targets.detach().cpu())

    if single_class:
        pred_all = torch.cat(pred_all)
        target_all = torch.cat(target_all)
        for i in range(opt.num_classes):
            buff_tar = target_all[target_all == i]
            buff_pred = pred_all[target_all == i]
            total_sc[i] = buff_tar.shape[0]
            correct_sc[i] = (buff_pred == i).sum().item()
        return correct / total, correct_sc / total_sc
    else:
        return correct / total