train_subset.py

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, Subset
import models
import os
from tqdm import tqdm
import argparse
import random
import numpy as np
import logging
from pathlib import Path
from tqdm import tqdm
from sklearn.metrics.pairwise import cosine_similarity
import matplotlib.pyplot as plt


def set_seed(seed):
    """Set all random seeds for reproducibility"""
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False


def filter_dataset_by_labels(dataset, selected_labels):
    """
    Filter dataset to only include samples with specified labels.
    Also remaps labels to start from 0 for the filtered dataset.

    Args:
        dataset: PyTorch dataset
        selected_labels: List of label indices to include

    Returns:
        Subset of the original dataset containing only selected labels with remapped labels
    """
    # Create a label mapping from original labels to new sequential labels
    label_mapping = {
        old_label: new_label
        for new_label, old_label in enumerate(sorted(selected_labels))
    }

    # Get all labels from the dataset
    all_labels = []
    for i in range(len(dataset)):
        _, label = dataset[i]
        all_labels.append(label)

    # Find indices of samples with selected labels
    selected_indices = []
    for i, label in enumerate(all_labels):
        if label in selected_labels:
            selected_indices.append(i)

    print(f"Original dataset size: {len(dataset)}")
    print(f"Selected labels: {selected_labels}")
    print(f"Filtered dataset size: {len(selected_indices)}")
    print(f"Label mapping: {label_mapping}")
    print(f"Label distribution in filtered dataset:")

    # Count samples per selected label
    label_counts = {}
    for idx in selected_indices:
        _, label = dataset[idx]
        label_counts[label] = label_counts.get(label, 0) + 1

    for label in sorted(selected_labels):
        count = label_counts.get(label, 0)
        print(f"  Label {label} -> {label_mapping[label]}: {count} samples")

    # Create a custom dataset that remaps labels
    class LabelRemappedDataset(torch.utils.data.Dataset):
        def __init__(self, dataset, indices, label_mapping):
            self.dataset = dataset
            self.indices = indices
            self.label_mapping = label_mapping

        def __len__(self):
            return len(self.indices)

        def __getitem__(self, idx):
            data, label = self.dataset[self.indices[idx]]
            # Remap the label to the new sequential range
            new_label = self.label_mapping[label]
            return data, new_label

    return LabelRemappedDataset(dataset, selected_indices, label_mapping)


def compute_classifier_alignment(model, test_loader, device, num_classes):
    """
    Compute alignment between classifier weights and class-mean features.
    Returns alignment scores and feature statistics.
    """
    model.eval()

    # Extract features and labels
    all_features = []
    all_labels = []

    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)

            # Get features before the classifier (assuming ResNet18 structure)
            # Remove the classifier layer temporarily
            features = model.avgpool(
                model.layer4(
                    model.layer3(
                        model.layer2(
                            model.layer1(model.relu(model.bn1(model.conv1(data))))
                        )
                    )
                )
            )
            features = features.view(features.size(0), -1)  # Flatten

            all_features.append(features.cpu().numpy())
            all_labels.append(target.cpu().numpy())

    # Concatenate all features and labels
    all_features = np.concatenate(all_features, axis=0)
    all_labels = np.concatenate(all_labels, axis=0)

    # Compute class-mean features
    class_mean_features = []
    for class_idx in range(num_classes):
        class_mask = all_labels == class_idx
        if np.sum(class_mask) > 0:
            class_mean = np.mean(all_features[class_mask], axis=0)
            class_mean_features.append(class_mean)
        else:
            # If no samples for this class, use zero vector
            class_mean_features.append(np.zeros(all_features.shape[1]))

    class_mean_features = np.array(
        class_mean_features
    )  # Shape: (num_classes, feature_dim)

    # Get classifier weights
    classifier_weights = (
        model.fc.weight.detach().cpu().numpy()
    )  # Shape: (num_classes, feature_dim)

    # Compute cosine similarity between classifier weights and class-mean features
    alignment_scores = []
    for i in range(num_classes):
        weight_vec = classifier_weights[i]
        feature_vec = class_mean_features[i]

        # Normalize vectors
        weight_norm = np.linalg.norm(weight_vec)
        feature_norm = np.linalg.norm(feature_vec)

        if weight_norm > 0 and feature_norm > 0:
            cosine_sim = np.dot(weight_vec, feature_vec) / (weight_norm * feature_norm)
            alignment_scores.append(cosine_sim)
        else:
            alignment_scores.append(0.0)

    alignment_scores = np.array(alignment_scores)

    # Compute overall alignment metrics
    mean_alignment = np.mean(alignment_scores)
    std_alignment = np.std(alignment_scores)
    min_alignment = np.min(alignment_scores)
    max_alignment = np.max(alignment_scores)

    return {
        "alignment_scores": alignment_scores,
        "mean_alignment": mean_alignment,
        "std_alignment": std_alignment,
        "min_alignment": min_alignment,
        "max_alignment": max_alignment,
        "class_mean_features": class_mean_features,
        "classifier_weights": classifier_weights,
    }


def save_alignment_plot(alignment_history, save_path):
    """Save alignment metrics over epochs as a plot"""
    epochs = list(alignment_history.keys())
    mean_alignments = [alignment_history[epoch]["mean_alignment"] for epoch in epochs]
    std_alignments = [alignment_history[epoch]["std_alignment"] for epoch in epochs]

    plt.figure(figsize=(10, 6))
    plt.plot(epochs, mean_alignments, "b-", label="Mean Alignment")
    plt.fill_between(
        epochs,
        [m - s for m, s in zip(mean_alignments, std_alignments)],
        [m + s for m, s in zip(mean_alignments, std_alignments)],
        alpha=0.3,
        label="±1 Std Dev",
    )
    plt.xlabel("Epoch")
    plt.ylabel("Cosine Similarity")
    plt.title("Classifier Weight Alignment with Class Features")
    plt.legend()
    plt.grid(True, alpha=0.3)
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()


def train_model(seed, args, output_dir):
    """Train model with specific seed"""
    print(f"\nTraining with seed {seed}")
    print("=" * 60)

    # Set seeds for reproducibility
    set_seed(seed)

    # Device configuration
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    verify_output = (
        f'outputs/seed{seed}_{args.dataset}/"logs"/"verification_results.txt'
    )

    # Dataset selection and normalization
    if args.dataset == "cifar10":
        num_classes = 10
        mean = (0.4914, 0.4822, 0.4465)
        std = (0.2023, 0.1994, 0.2010)
        train_dataset = datasets.CIFAR10(
            root="../data/",
            train=True,
            transform=transforms.Compose(
                [
                    transforms.RandomCrop(32, padding=4),
                    transforms.RandomHorizontalFlip(),
                    transforms.ToTensor(),
                    transforms.Normalize(mean, std),
                ]
            ),
            download=True,
        )
        test_dataset = datasets.CIFAR10(
            root="../data/",
            train=False,
            transform=transforms.Compose(
                [transforms.ToTensor(), transforms.Normalize(mean, std)]
            ),
            download=True,
        )
        model_out_name = "cifar10_model"
    elif args.dataset == "cifar100":
        num_classes = 100
        mean = (0.5071, 0.4867, 0.4408)
        std = (0.2675, 0.2565, 0.2761)

        # Training data with augmentation
        train_transform = transforms.Compose(
            [
                transforms.RandomCrop(32, padding=4),  # Random crop with padding
                transforms.RandomHorizontalFlip(p=0.5),  # Random horizontal flip
                transforms.ColorJitter(
                    brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1
                ),  # Color jittering
                transforms.ToTensor(),
                transforms.Normalize(mean, std),
            ]
        )

        # Validation data with minimal preprocessing
        test_transform = transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize(mean, std)]
        )

        train_dataset = datasets.CIFAR100(
            root="../data/", train=True, transform=train_transform, download=True
        )
        test_dataset = datasets.CIFAR100(
            root="../data/", train=False, transform=test_transform, download=True
        )
        model_out_name = "cifar100_model"
    elif args.dataset == "tinyimagenet":
        num_classes = 200
        mean = (0.485, 0.456, 0.406)
        std = (0.229, 0.224, 0.225)
        train_dir = os.path.join(args.tinyimagenet_root, "train")
        val_dir = os.path.join(args.tinyimagenet_root, "val", "images")
        train_dataset = datasets.ImageFolder(
            root=train_dir,
            transform=transforms.Compose(
                [
                    transforms.Resize((64, 64)),
                    transforms.ToTensor(),
                    transforms.Normalize(mean, std),
                ]
            ),
        )
        test_dataset = datasets.ImageFolder(
            root=val_dir,
            transform=transforms.Compose(
                [
                    transforms.Resize((64, 64)),
                    transforms.ToTensor(),
                    transforms.Normalize(mean, std),
                ]
            ),
        )
        model_out_name = "tinyimagenet_model"
    else:
        raise ValueError(f"Unsupported dataset: {args.dataset}")

    # Filter datasets by selected labels
    if args.selected_labels is not None:
        print(f"\nFiltering datasets to include only labels: {args.selected_labels}")
        train_dataset = filter_dataset_by_labels(train_dataset, args.selected_labels)
        test_dataset = filter_dataset_by_labels(test_dataset, args.selected_labels)

        # Update num_classes to match the number of selected labels
        num_classes = len(args.selected_labels)
        print(f"Updated number of classes: {num_classes}")

    train_loader = DataLoader(
        train_dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=16,  # 8–16 depending on CPU
        pin_memory=True,
        persistent_workers=True,
        prefetch_factor=4,
    )
    test_loader = DataLoader(
        test_dataset,
        batch_size=args.batch_size,
        num_workers=16,
        pin_memory=True,
        persistent_workers=True,
        prefetch_factor=4,
    )

    # Initialize model
    # model = models.SimpleConvNet().to(device)
    model = models.ResNet18()
    """
    ResNet18网络的7x7降采样卷积和池化操作容易丢失一部分信息,
    所以在实验中我们将7x7的降采样层和最大池化层去掉,替换为一个3x3的降采样卷积,
    同时减小该卷积层的步长和填充大小
    """
    model.conv1 = nn.Conv2d(
        in_channels=3, out_channels=64, kernel_size=3, stride=1, padding=1, bias=False
    )
    model.maxpool = nn.Identity()
    model.fc = torch.nn.Linear(512, num_classes)
    model = model.to(device)

    # Load pre-trained weights if specified
    if args.classifier_only:
        print(f"Loading pre-trained weights from: {args.classifier_only}")
        if not os.path.exists(args.classifier_only):
            raise FileNotFoundError(
                f"Model weights file not found: {args.classifier_only}"
            )

        # Load the state dict
        state_dict = torch.load(args.classifier_only, map_location=device)
        model.load_state_dict(state_dict)
        print("Pre-trained weights loaded successfully")

        # Freeze all layers except the classifier (fc layer)
        print("Freezing all layers except the classifier...")
        for name, param in model.named_parameters():
            if "fc" not in name:  # Freeze all layers except the final classifier
                param.requires_grad = False
                print(f"Frozen: {name}")
            else:
                print(f"Trainable: {name}")

    criterion = nn.CrossEntropyLoss()

    # Only optimize trainable parameters
    if args.classifier_only:
        optimizer = optim.SGD(
            filter(lambda p: p.requires_grad, model.parameters()),
            lr=args.learning_rate,
            momentum=0.9,
        )
        print("Optimizer configured to train only classifier parameters")
    else:
        optimizer = optim.SGD(
            model.parameters(), lr=args.learning_rate, momentum=0.9, weight_decay=5e-4
        )

    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)

    # Initialize alignment tracking
    alignment_history = {}
    alignment_check_freq = args.alignment_check_freq

    best_acc = 0
    for epoch in range(args.epochs):
        model.train()
        train_correct = 0
        train_total = 0
        train_loss = 0

        # Training progress bar
        train_pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{args.epochs} [Train]")
        for data, target in train_pbar:
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()

            # Calculate training accuracy
            _, predicted = torch.max(output.data, 1)
            train_total += target.size(0)
            train_correct += (predicted == target).sum().item()
            train_loss += loss.item()

            # Update progress bar
            current_lr = optimizer.param_groups[0]["lr"]
            train_pbar.set_postfix(
                {
                    "loss": f"{train_loss/(train_pbar.n+1):.4f}",
                    "acc": f"{100.*train_correct/train_total:.2f}%",
                    "lr": f"{current_lr:.6f}",
                }
            )

        # Evaluate
        model.eval()
        val_correct = 0
        val_total = 0
        val_loss = 0

        # Validation progress bar
        val_pbar = tqdm(test_loader, desc=f"Epoch {epoch+1}/{args.epochs} [Valid]")
        with torch.no_grad():
            for data, target in val_pbar:
                data, target = data.to(device), target.to(device)
                output = model(data)
                loss = criterion(output, target)
                val_loss += loss.item()

                _, predicted = torch.max(output.data, 1)
                val_total += target.size(0)
                val_correct += (predicted == target).sum().item()

                # Update progress bar
                val_pbar.set_postfix(
                    {
                        "loss": f"{val_loss/(val_pbar.n+1):.4f}",
                        "acc": f"{100.*val_correct/val_total:.2f}%",
                    }
                )

        # Step the scheduler
        if scheduler is not None:
            scheduler.step()

        # Compute classifier alignment if enabled and at check frequency
        if args.track_alignment and (epoch + 1) % alignment_check_freq == 0:
            print(f"\nComputing classifier alignment at epoch {epoch + 1}...")
            alignment_results = compute_classifier_alignment(
                model, test_loader, device, num_classes
            )
            alignment_history[epoch + 1] = alignment_results

            # Print alignment summary
            tqdm.write(f"Alignment Metrics (Epoch {epoch + 1}):")
            tqdm.write(f"  Mean: {alignment_results['mean_alignment']:.4f}")
            tqdm.write(f"  Std:  {alignment_results['std_alignment']:.4f}")
            tqdm.write(f"  Min:  {alignment_results['min_alignment']:.4f}")
            tqdm.write(f"  Max:  {alignment_results['max_alignment']:.4f}")

            # Log alignment metrics
            logging.info(
                f"Epoch{epoch+1} Alignment - Mean: {alignment_results['mean_alignment']:.4f}, "
                f"Std: {alignment_results['std_alignment']:.4f}, "
                f"Min: {alignment_results['min_alignment']:.4f}, "
                f"Max: {alignment_results['max_alignment']:.4f}"
            )

        # Print epoch summary
        train_acc = 100.0 * train_correct / train_total
        val_acc = 100.0 * val_correct / val_total
        current_lr = optimizer.param_groups[0]["lr"]
        tqdm.write(f"\nEpoch {epoch+1}/{args.epochs}:")
        tqdm.write(
            f"Train Loss: {train_loss/len(train_loader):.4f}, Train Acc: {train_acc:.2f}%"
        )
        tqdm.write(
            f"Valid Loss: {val_loss/len(test_loader):.4f}, Valid Acc: {val_acc:.2f}%"
        )
        tqdm.write(f"Learning Rate: {current_lr:.6f}")
        logging.info(
            f"Epoch{epoch+1}'s loss: {train_loss/len(train_loader):.4f}, Train Acc: {train_acc:.2f}%"
            f" | Test Acc: {val_acc:.2f}% | LR: {current_lr:.6f}"
        )

        # Save best model
        if val_acc > best_acc:
            best_acc = val_acc
            torch.save(
                model.state_dict(),
                f"{output_dir}/{model_out_name}_best.pth",
            )
            tqdm.write(f"Saved best model with validation accuracy: {best_acc:.2f}%")
        tqdm.write("-" * 60)

    # Save final model
    torch.save(
        model.state_dict(),
        f"{output_dir}/{model_out_name}_final.pth",
    )

    # Save alignment plot if alignment was tracked
    if args.track_alignment and alignment_history:
        alignment_plot_path = f"{output_dir}/alignment_plot.png"
        save_alignment_plot(alignment_history, alignment_plot_path)
        print(f"Alignment plot saved to: {alignment_plot_path}")

        # Save alignment data
        alignment_data_path = f"{output_dir}/alignment_data.npz"
        np.savez(
            alignment_data_path,
            epochs=np.array(list(alignment_history.keys())),
            mean_alignments=np.array(
                [
                    alignment_history[epoch]["mean_alignment"]
                    for epoch in alignment_history.keys()
                ]
            ),
            std_alignments=np.array(
                [
                    alignment_history[epoch]["std_alignment"]
                    for epoch in alignment_history.keys()
                ]
            ),
            min_alignments=np.array(
                [
                    alignment_history[epoch]["min_alignment"]
                    for epoch in alignment_history.keys()
                ]
            ),
            max_alignments=np.array(
                [
                    alignment_history[epoch]["max_alignment"]
                    for epoch in alignment_history.keys()
                ]
            ),
        )
        print(f"Alignment data saved to: {alignment_data_path}")

    return best_acc


def setup_logging(run_dir: Path) -> None:
    """Configure logging to file and console for each run"""
    log_file = run_dir / "logs/run.log"
    (run_dir / "logs").mkdir(exist_ok=True)
    # 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 main():
    parser = argparse.ArgumentParser(
        description="Train cifar10 or cifar100 models with multiple seeds on subset of labels"
    )
    parser.add_argument(
        "--seeds",
        nargs="+",
        type=int,
        default=[0, 1, 2],
        help="List of seeds to use for training",
    )
    parser.add_argument(
        "--batch-size", type=int, default=512, help="Batch size for training"
    )
    parser.add_argument(
        "--epochs", type=int, default=100, help="Number of epochs to train"
    )
    parser.add_argument(
        "--learning-rate", type=float, default=0.4, help="Learning rate"
    )
    parser.add_argument(
        "--dataset",
        type=str,
        default="cifar10",
        choices=["cifar10", "cifar100", "tinyimagenet"],
        help="Dataset to use (cifar10, cifar100, or tinyimagenet)",
    )
    parser.add_argument(
        "--tinyimagenet-root",
        type=str,
        default="../data/tiny-imagenet-200",
        help="Root directory for TinyImageNet dataset (default: ../tiny-imagenet-200)",
    )
    parser.add_argument(
        "--selected-labels",
        nargs="+",
        type=int,
        default=None,
        help="List of label indices to include in training (e.g., 0 1 2 for first three classes). If not specified, uses all labels.",
    )
    parser.add_argument(
        "--classifier-only",
        type=str,
        default=None,
        help="Path to pre-trained model weights to resume from (only trains classifier when specified)",
    )
    parser.add_argument(
        "--lr-scheduler",
        action="store_true",
        default=False,
        help="Enable learning rate scheduler (StepLR)",
    )
    parser.add_argument(
        "--lr-step-size",
        type=int,
        default=15,
        help="Step size for learning rate scheduler (default: 75)",
    )
    parser.add_argument(
        "--lr-gamma",
        type=float,
        default=0.1,
        help="Gamma for learning rate scheduler (default: 0.1)",
    )
    parser.add_argument(
        "--track-alignment",
        action="store_true",
        default=False,
        help="Enable tracking of classifier alignment metrics during training",
    )
    parser.add_argument(
        "--alignment-check-freq",
        type=int,
        default=1,
        help="Frequency (in epochs) to check and save classifier alignment (default: 10)",
    )
    args = parser.parse_args()

    # Validate selected labels
    if args.selected_labels is not None:
        if args.dataset == "cifar10" and max(args.selected_labels) >= 10:
            raise ValueError(
                "CIFAR-10 has 10 classes (0-9), but selected labels exceed this range"
            )
        elif args.dataset == "cifar100" and max(args.selected_labels) >= 100:
            raise ValueError(
                "CIFAR-100 has 100 classes (0-99), but selected labels exceed this range"
            )
        elif args.dataset == "tinyimagenet" and max(args.selected_labels) >= 200:
            raise ValueError(
                "TinyImageNet has 200 classes (0-199), but selected labels exceed this range"
            )

        print(f"Training on subset of labels: {args.selected_labels}")
        print(f"Number of classes: {len(args.selected_labels)}")
    else:
        print("Training on all available labels")

    # Store results
    results = {}

    # Train for each seed
    for seed in args.seeds:
        # Create output directory with subset information
        if args.selected_labels is not None:
            labels_str = "_".join(map(str, sorted(args.selected_labels)))
            output_dir = f"outputs/seed{seed}_{args.dataset}_subset_{labels_str}"
        else:
            output_dir = f"outputs/seed{seed}_{args.dataset}"

        print(f"Creating output directory: {output_dir}")
        os.makedirs(output_dir, exist_ok=True)
        setup_logging(Path(output_dir))
        best_acc = train_model(seed, args, output_dir)
        results[seed] = best_acc

    # Print summary of results
    print("\nTraining Complete!")
    print("=" * 60)
    print("Results Summary:")
    for seed, acc in results.items():
        print(f"Seed {seed}: Best accuracy = {acc:.2f}%")

    # Calculate statistics
    accuracies = list(results.values())
    mean_acc = np.mean(accuracies)
    std_acc = np.std(accuracies)
    print(f"\nMean accuracy: {mean_acc:.2f}% ± {std_acc:.2f}%")


if __name__ == "__main__":
    main()


# python src/train_subset.py --dataset cifar10 --selected-labels 1 2 3 4 5 6 7 8 9 --epochs 100 --batch-size 64 --seeds 1