{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.12"},"kaggle":{"accelerator":"none","dataSources":[],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false},"papermill":{"default_parameters":{},"duration":30931.326158,"end_time":"2026-01-21T06:14:24.530873","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2026-01-20T21:38:53.204715","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"db722de7","cell_type":"code","source":"# Install py7zr for extracting Kaggle dataset archives\n!pip install py7zr","metadata":{"execution":{"iopub.execute_input":"2026-01-20T21:38:55.672472Z","iopub.status.busy":"2026-01-20T21:38:55.671837Z","iopub.status.idle":"2026-01-20T21:39:12.402842Z","shell.execute_reply":"2026-01-20T21:39:12.402091Z"},"papermill":{"duration":16.73808,"end_time":"2026-01-20T21:39:12.404665","exception":false,"start_time":"2026-01-20T21:38:55.666585","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"bcfc6afa","cell_type":"markdown","source":"# FAST ProgFL (Sequential) + FedBN for CIFAR-10\nThis notebook implements a Federated Learning simulation using Progressive Layer Training (ProgFL) and FedBN on the CIFAR-10 dataset.\n\n### Key Features:\n- **FedAvg Baseline**: Standard Federated Averaging.\n- **Progressive Layer Training**: Sequentially training layers of the model to reduce communication costs.\n- **FedBN**: Keeping Batch Normalization statistics local to handle non-IID data.\n- **Warmup Rounds**: Initial rounds of standard FedAvg before switching to progressive training.\n- **Optimization**: Uses AMP (Automatic Mixed Precision) and optimized data loading for speed.","metadata":{"papermill":{"duration":0.003811,"end_time":"2026-01-20T21:39:12.412413","exception":false,"start_time":"2026-01-20T21:39:12.408602","status":"completed"},"tags":[]}},{"id":"015cf3e2","cell_type":"markdown","source":"## Theoretical Background and References\n\nThis implementation integrates several advanced Federated Learning (FL) concepts to address key challenges: **Communication Efficiency** and **Data Heterogeneity (Non-IID)**.\n\n### 1. Federated Averaging (FedAvg)\nThe baseline algorithm used is **FedAvg**, where a central server coordinates a global model update by aggregating local updates from clients.\n*   **Reference**: McMahan, B., et al. (2017). *Communication-Efficient Learning of Deep Networks from Decentralized Data*. AISTATS. [Link](https://arxiv.org/abs/1602.05629)\n\n### 2. Addressing Non-IID Data: FedBN\nData across clients is often Non-Independent and Identically Distributed (Non-IID). This code uses **FedBN (Federated Batch Normalization)**, which keeps Batch Normalization (BN) parameters local to each client while aggregating other model parameters. This allows the model to adapt to local feature shifts without incurring extra communication or computational costs.\n*   **Logic**: In the code, `is_bn_key` identifies BN parameters. During aggregation, these are excluded (`if use_fedbn and is_bn_key(k): continue`). Clients restore their specific BN state before training.\n*   **Reference**: Li, X., et al. (2021). *FedBN: Federated Learning on Non-IID Features via Local Batch Normalization*. ICLR. [Link](https://arxiv.org/abs/2102.07623)\n\n### 3. Communication Efficiency: Progressive Layer Training (ProgFL)\nTo reduce the communication bottleneck, this implementation uses a **Sequential Progressive Layer Training** strategy. Instead of updating and transmitting the full model every round, only a specific subset of layers (groups) is trained and communicated.\n*   **Logic**: The model is split into groups (e.g., `conv1`, `layer1`, `fc`). In each round, only one group is active (`active_group_idx`). Other layers are frozen (`freeze_all` then `activate_group`). This significantly reduces `comm_mb`.\n*   **Context**: This aligns with research on *Layer-wise Federated Learning* and *Sparse Updates* to minimize bandwidth usage.\n\n### 4. Data Distribution (Dirichlet)\nTo simulate realistic Non-IID settings, the CIFAR-10 dataset is partitioned using a **Dirichlet distribution** ($\\alpha$). A lower $\\alpha$ indicates higher heterogeneity.\n*   **Reference**: Hsu, T. M. H., et al. (2019). *Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification*. [Link](https://arxiv.org/abs/1909.06335)","metadata":{"papermill":{"duration":0.003554,"end_time":"2026-01-20T21:39:12.419615","exception":false,"start_time":"2026-01-20T21:39:12.416061","status":"completed"},"tags":[]}},{"id":"89659510","cell_type":"code","source":"import os, time, random, math\nfrom dataclasses import dataclass\nfrom typing import List, Dict, Optional, Set, Tuple\n\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader, Subset\n\nimport torchvision\nimport torchvision.transforms as T\nimport matplotlib.pyplot as plt\n\nfrom torch.amp import autocast, GradScaler\n\nimport shutil\nimport os\n","metadata":{"execution":{"iopub.execute_input":"2026-01-20T21:39:12.428280Z","iopub.status.busy":"2026-01-20T21:39:12.427598Z","iopub.status.idle":"2026-01-20T21:39:21.814665Z","shell.execute_reply":"2026-01-20T21:39:21.814042Z"},"papermill":{"duration":9.393507,"end_time":"2026-01-20T21:39:21.816521","exception":false,"start_time":"2026-01-20T21:39:12.423014","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"b48ffae5","cell_type":"code","source":"# ============================================================\n# 0) CONFIG (ALL HYPERPARAMETERS HERE)\n# ============================================================\n@dataclass\nclass CFG:\n    # System\n    SEED: int = 42\n    DEVICE: str = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n    OUTDIR: str = \"./fl_runs\"\n    PROJECT: str = \"cifar10_progfl_seq_fedbn_fast\"\n\n    # Dataset\n    DATA_DIR: str = \"./data\"\n    NUM_CLASSES: int = 10\n    Global_EXTRACT_PATH:str=None\n\n    # FL simulation\n    NUM_CLIENTS: int = 20\n    CLIENT_FRACTION: float = 0.3\n    NONIID_ALPHA: float = (\n        0.5  # Controls data heterogeneity (Lower = more Non-IID) [Hsu et al., 2019]\n    )\n\n    # Training\n    ROUNDS: int = 200\n    LOCAL_EPOCHS: int = 5\n    BATCH_SIZE: int = 512  # ↑ T4 likes larger batch for utilization\n    TEST_BS: int = 512\n    LR: float = 0.01\n    MOMENTUM: float = 0.9\n    WEIGHT_DECAY: float = 5e-4\n\n    # Progressive FL\n    WARMUP_ROUNDS: int = 10  # FedAvg warmup then progressive\n    PROG_GROUPS_PER_ROUND: int = 1\n\n    # FedBN\n    USE_FEDBN: bool = False\n\n    # Evaluation frequency (HUGE speedup)\n    EVAL_EVERY: int = 5  # evaluate every N rounds\n\n    # DataLoader performance (Kaggle/Colab)\n    NUM_WORKERS: int = 2\n    PIN_MEMORY: bool = True\n    PERSISTENT_WORKERS: bool = True\n    PREFETCH_FACTOR: int = 2\n\n    # AMP\n    USE_AMP: bool = True\n\n    # cuDNN / TF32 speed knobs\n    CUDNN_BENCHMARK: bool = True\n    TF32: bool = True\n    CHANNELS_LAST: bool = True\n\n    # Reporting\n    SAVE_CSV: bool = True\n    SAVE_PLOTS: bool = True\n    EXPORT_PDF: bool = True\n    DPI: int = 250\n\n    # Multi-seed (publication)\n    MULTI_SEED: bool = True  # Turn ON later for CI plots\n    SEEDS: Tuple[int, ...] = (42, 123, 999)\n    CI_Z: float = 1.96\n\n    # Experiments\n    MAX_EXPERIMENTS: int = 10\n\n    EXPERIMENTS: Tuple[Dict, ...] = (\n        dict(name=\"FedAvg\", progressive=False, fedbn=False, warmup=0),\n        dict(\n            name=\"ProgFL_Sequential_FedBN_Warmup\",\n            progressive=True,\n            fedbn=True,\n            warmup=10,\n        ),\n        # # Ablations (enable later if needed):\n        dict(name=\"ProgFL_Sequential_FedBN_NoWarmup\", progressive=True, fedbn=True, warmup=0),\n        dict(name=\"ProgFL_Sequential_NoFedBN_Warmup\", progressive=True, fedbn=False, warmup=10),\n    )\n\n\ncfg = CFG()","metadata":{"execution":{"iopub.execute_input":"2026-01-20T21:39:21.826747Z","iopub.status.busy":"2026-01-20T21:39:21.826186Z","iopub.status.idle":"2026-01-20T21:39:21.913108Z","shell.execute_reply":"2026-01-20T21:39:21.912299Z"},"papermill":{"duration":0.093916,"end_time":"2026-01-20T21:39:21.914711","exception":false,"start_time":"2026-01-20T21:39:21.820795","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"1e9bade1","cell_type":"markdown","source":"### Model Architecture & Progressive Grouping\nHere we define the ResNet18 model. Crucially, for **ProgFL**, we define `layer_groups`.\n*   **Logic**: The model is segmented into logical blocks (Input Conv, Layer 1-4, FC).\n*   **Purpose**: By training only one block at a time, we drastically reduce the number of parameters sent over the network in each round.","metadata":{"papermill":{"duration":0.003772,"end_time":"2026-01-20T21:39:21.922345","exception":false,"start_time":"2026-01-20T21:39:21.918573","status":"completed"},"tags":[]}},{"id":"0db73b10","cell_type":"code","source":"# ============================================================\n# 1) UTILITIES\n# ============================================================\ndef set_seed(seed: int):\n    \"\"\"Sets the seed for reproducibility.\"\"\"\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed_all(seed)\n\n\ndef ensure_dir(path: str):\n    \"\"\"Creates a directory if it does not exist.\"\"\"\n    os.makedirs(path, exist_ok=True)\n\n\ndef accuracy(logits: torch.Tensor, y: torch.Tensor) -> float:\n    \"\"\"Calculates accuracy from logits and labels.\"\"\"\n    pred = logits.argmax(dim=1)\n    return (pred == y).float().mean().item()\n\n\ndef savefig_both(path_no_ext: str):\n    \"\"\"Saves a figure in both PNG and PDF formats.\"\"\"\n    png_path = path_no_ext + \".png\"\n    plt.tight_layout()\n    plt.savefig(png_path, dpi=cfg.DPI)\n    if cfg.EXPORT_PDF:\n        plt.savefig(path_no_ext + \".pdf\")\n    plt.close()\n\n\ndef is_bn_key(k: str) -> bool:\n    \"\"\"Checks if a parameter key belongs to Batch Normalization.\"\"\"\n    return (\n        \"running_mean\" in k\n        or \"running_var\" in k\n        or \"num_batches_tracked\" in k\n        or \".bn\" in k\n        or k.startswith(\"bn\")\n    )\n\n\ndef model_bytes(sd: Dict[str, torch.Tensor], keys: Optional[Set[str]] = None) -> int:\n    \"\"\"Calculates the size of the model (or specific keys) in bytes.\"\"\"\n    total = 0\n    for k, v in sd.items():\n        if keys is not None and k not in keys:\n            continue\n        total += v.numel() * v.element_size()\n    return total","metadata":{"execution":{"iopub.execute_input":"2026-01-20T21:39:21.930840Z","iopub.status.busy":"2026-01-20T21:39:21.930583Z","iopub.status.idle":"2026-01-20T21:39:21.939208Z","shell.execute_reply":"2026-01-20T21:39:21.938501Z"},"papermill":{"duration":0.014637,"end_time":"2026-01-20T21:39:21.940693","exception":false,"start_time":"2026-01-20T21:39:21.926056","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"6de44f8b","cell_type":"code","source":"# ============================================================\n# 2) MODEL\n# ============================================================\n# def build_model() -> nn.Module:\n#     \"\"\"Builds a ResNet18 model adapted for CIFAR-10.\"\"\"\n#     model = torchvision.models.resnet18(num_classes=cfg.NUM_CLASSES)\n#     model.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1, bias=False)\n#     model.maxpool = nn.Identity()\n\n#     if cfg.CHANNELS_LAST:\n#         model = model.to(memory_format=torch.channels_last)\n\n#     return model\n\n# ============================================================\n# 2) MODEL\n# ============================================================\ndef build_model() -> nn.Module:\n    \"\"\"\n    Builds a ResNet18 model adapted for CIFAR-10.\n    USES PRETRAINED WEIGHTS (IMAGENET) for faster convergence.\n    \"\"\"\n    # 1. Load the model with ImageNet weights\n    # 'weights=\"DEFAULT\"' automatically downloads the best available weights\n    model = torchvision.models.resnet18(weights=\"DEFAULT\")\n    \n    # 2. Modify the First Layer (Input)\n    # ImageNet expects 224x224 images, but CIFAR-10 is 32x32.\n    # We keep the weights but change the stride/padding to avoid downsampling too fast.\n    # Note: We must re-initialize this layer because the architecture changes slightly.\n    model.conv1 = nn.Conv2d(\n        in_channels=3, \n        out_channels=64, \n        kernel_size=3,  # Changed from 7 to 3 for small images\n        stride=1,       # Changed from 2 to 1 to keep resolution\n        padding=1,      # Changed from 3 to 1\n        bias=False\n    )\n    \n    # 3. Remove the MaxPool layer\n    # For 32x32 images, early maxpooling destroys too much information.\n    model.maxpool = nn.Identity()\n    \n    # 4. Modify the Final Layer (Output)\n    # ImageNet has 1000 classes; CIFAR-10 has 10.\n    # We replace the final fully connected layer.\n    model.fc = nn.Linear(model.fc.in_features, cfg.NUM_CLASSES)\n\n    # Performance optimization\n    if cfg.CHANNELS_LAST:\n        model = model.to(memory_format=torch.channels_last)\n\n    return model\n\n\n# def define_layer_groups(model: nn.Module) -> List[List[str]]:\n#     \"\"\"Defines groups of layers for progressive training.\"\"\"\n#     groups = []\n#     groups.append(\n#         [\n#             n\n#             for n, _ in model.named_parameters()\n#             if n.startswith(\"conv1\") or n.startswith(\"bn1\")\n#         ]\n#     )\n#     groups.append([n for n, _ in model.named_parameters() if n.startswith(\"layer1\")])\n#     groups.append([n for n, _ in model.named_parameters() if n.startswith(\"layer2\")])\n#     groups.append([n for n, _ in model.named_parameters() if n.startswith(\"layer3\")])\n#     groups.append([n for n, _ in model.named_parameters() if n.startswith(\"layer4\")])\n#     groups.append([n for n, _ in model.named_parameters() if n.startswith(\"fc\")])\n#     return groups\n\ndef define_layer_groups(model: nn.Module) -> List[List[str]]:\n    \"\"\"\n    Defines fewer, larger groups for better stability.\n    Group 0: Low-level features (Conv1, Layer1, Layer2)\n    Group 1: High-level features (Layer3, Layer4)\n    Note: 'fc' is intentionally excluded here; we will add it manually in the loop.\n    \"\"\"\n    groups = []\n    \n    # Group 1: Low Level Features\n    # Matches 'conv1', 'bn1', 'layer1', 'layer2'\n    g1 = [\n        n for n, _ in model.named_parameters() \n        if any(x in n for x in [\"conv1\", \"bn1\", \"layer1\", \"layer2\"])\n    ]\n    groups.append(g1)\n\n    # Group 2: High Level Features\n    # Matches 'layer3', 'layer4'\n    g2 = [\n        n for n, _ in model.named_parameters() \n        if any(x in n for x in [\"layer3\", \"layer4\"])\n    ]\n    groups.append(g2)\n    \n    return groups\n\n\ndef freeze_all(model: nn.Module):\n    \"\"\"Freezes all parameters in the model.\"\"\"\n    for p in model.parameters():\n        p.requires_grad = False\n\n\ndef activate_group(model: nn.Module, group_names: Set[str]):\n    \"\"\"Unfreezes parameters belonging to a specific group.\"\"\"\n    d = dict(model.named_parameters())\n    for n in group_names:\n        if n in d:\n            d[n].requires_grad = True\n\n\ndef remove_directory(dir_path):\n    # Check if the path exists and is actually a directory\n    if os.path.exists(dir_path) and os.path.isdir(dir_path):\n        try:\n            # Removes the directory and all its contents\n            shutil.rmtree(dir_path)\n            print(f\"Success: '{dir_path}' has been deleted.\")\n        except OSError as e:\n            print(f\"Error: Could not delete '{dir_path}'. Reason: {e}\")\n    else:\n        print(f\"Skipped: '{dir_path}' does not exist or is not a directory.\")\n\n# Example usage:\n# remove_directory(\"temp_folder\")","metadata":{"execution":{"iopub.execute_input":"2026-01-20T21:39:21.949736Z","iopub.status.busy":"2026-01-20T21:39:21.949231Z","iopub.status.idle":"2026-01-20T21:39:21.959218Z","shell.execute_reply":"2026-01-20T21:39:21.958632Z"},"papermill":{"duration":0.016254,"end_time":"2026-01-20T21:39:21.960656","exception":false,"start_time":"2026-01-20T21:39:21.944402","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"ada025e1","cell_type":"code","source":"# ============================================================\n# 3) DATA\n# ============================================================\nimport py7zr\nfrom PIL import Image\nfrom sklearn.model_selection import train_test_split\n\n\nclass KaggleCIFAR10Dataset(torch.utils.data.Dataset):\n    \"\"\"Custom Dataset for loading CIFAR-10 from Kaggle input.\"\"\"\n\n    def __init__(self, img_dir, labels_df, transform=None, samples_per_class=50):\n        self.img_dir = img_dir\n\n        # Filter data to keep specific number of samples per class\n        if samples_per_class is not None:\n            labels_df = labels_df.groupby(\"label\").head(samples_per_class)\n\n        self.labels_df = labels_df.reset_index(drop=True)\n        self.transform = transform\n\n        # Standard CIFAR-10 classes\n        classes = [\n            \"airplane\",\n            \"automobile\",\n            \"bird\",\n            \"cat\",\n            \"deer\",\n            \"dog\",\n            \"frog\",\n            \"horse\",\n            \"ship\",\n            \"truck\",\n        ]\n        self.class_to_idx = {cls: i for i, cls in enumerate(classes)}\n\n        # Pre-compute targets for dirichlet_split\n        self.targets = [\n            self.class_to_idx[row[\"label\"]] for _, row in self.labels_df.iterrows()\n        ]\n\n    def __len__(self):\n        return len(self.labels_df)\n\n    def __getitem__(self, idx):\n        row = self.labels_df.iloc[idx]\n        img_id = row[\"id\"]\n        target = self.targets[idx]\n\n        img_path = os.path.join(self.img_dir, f\"{img_id}.png\")\n        # Open image\n        img = Image.open(img_path).convert(\"RGB\")\n\n        if self.transform:\n            img = self.transform(img)\n\n        return img, target\n\n\ndef get_data(test_size=0.1, train_size=None):\n    \"\"\"\n    Prepares CIFAR-10 data from Kaggle input.\n    Extracts train.7z if needed, splits into Train/Test, and returns datasets.\n    \"\"\"\n    INPUT_DIR = \"/kaggle/input/cifar-10\"\n    WORKING_DIR = \"/kaggle/working\"\n    TRAIN_ARCHIVE = os.path.join(INPUT_DIR, \"train.7z\")\n    LABELS_CSV = os.path.join(INPUT_DIR, \"trainLabels.csv\")\n    EXTRACT_PATH = os.path.join(WORKING_DIR, \"train\")\n    cfg.Global_EXTRACT_PATH=EXTRACT_PATH\n\n    # 1. Extract Data if not exists\n    if not os.path.exists(EXTRACT_PATH):\n        print(f\"Extracting {TRAIN_ARCHIVE} to {WORKING_DIR}...\")\n        with py7zr.SevenZipFile(TRAIN_ARCHIVE, mode=\"r\") as z:\n            z.extractall(path=WORKING_DIR)\n        print(\"Extraction complete.\")\n    else:\n        print(\"Data already extracted.\")\n\n    # 2. Load Labels\n    df = pd.read_csv(LABELS_CSV)\n\n    # 3. Split into Train (45k) and Test (5k) to simulate standard setup\n    # Note: Standard CIFAR-10 is 50k train, 10k test. Here we only have 50k labeled.\n    # We'll use 45k for FL training and 5k for global evaluation.\n    train_df, test_df = train_test_split(\n        df,\n        test_size=test_size,\n        train_size=train_size,\n        random_state=cfg.SEED,\n        stratify=df[\"label\"],\n    )\n\n    # 4. Transforms\n    transform_train = T.Compose(\n        [\n            T.RandomCrop(32, padding=4),\n            T.RandomHorizontalFlip(),\n            T.ToTensor(), \n            T.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261)),\n        ]\n    )\n    transform_test = T.Compose(\n        [T.ToTensor(), T.Normalize((0.4914, 0.4822, 0.4465), (0.247, 0.243, 0.261))]\n    )\n\n    # 5. Create Datasets\n    trainset = KaggleCIFAR10Dataset(EXTRACT_PATH, train_df, transform=transform_train,samples_per_class=None)\n    testset = KaggleCIFAR10Dataset(EXTRACT_PATH, test_df, transform=transform_test,samples_per_class=None)\n\n    testloader = DataLoader(\n        testset,\n        batch_size=cfg.TEST_BS,\n        shuffle=False,\n        num_workers=cfg.NUM_WORKERS,\n        pin_memory=cfg.PIN_MEMORY,\n        persistent_workers=cfg.PERSISTENT_WORKERS,\n        prefetch_factor=cfg.PREFETCH_FACTOR,\n    )\n    \n\n    return trainset, testloader\n\n\ndef dirichlet_split(trainset, num_clients: int, alpha: float) -> List[List[int]]:\n    \"\"\"\n    Splits the dataset among clients using a Dirichlet distribution for non-IID data.\n\n    Reference: Hsu et al., \"Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification\".\n    \"\"\"\n    y = np.array(trainset.targets)\n    idx_by_class = [np.where(y == k)[0] for k in range(cfg.NUM_CLASSES)]\n    client_indices = [[] for _ in range(num_clients)]\n\n    for k in range(cfg.NUM_CLASSES):\n        np.random.shuffle(idx_by_class[k])\n        prop = np.random.dirichlet(alpha=np.ones(num_clients) * alpha)\n        prop = prop / prop.sum()\n\n        splits = (np.cumsum(prop) * len(idx_by_class[k])).astype(int)[:-1]\n        chunks = np.split(idx_by_class[k], splits)\n\n        for cid in range(num_clients):\n            client_indices[cid].extend(chunks[cid].tolist())\n\n    for cid in range(num_clients):\n        np.random.shuffle(client_indices[cid])\n\n    return client_indices\n\n\n# def make_client_loaders(trainset, client_indices: List[List[int]]) -> List[DataLoader]:\n#     \"\"\"Creates DataLoaders for each client.\"\"\"\n#     loaders = []\n#     for idxs in client_indices:\n#         ds = Subset(trainset, idxs)\n#         loader = DataLoader(\n#             ds,\n#             batch_size=cfg.BATCH_SIZE,\n#             shuffle=True,\n#             num_workers=cfg.NUM_WORKERS,\n#             pin_memory=cfg.PIN_MEMORY,\n#             persistent_workers=cfg.PERSISTENT_WORKERS,\n#             prefetch_factor=cfg.PREFETCH_FACTOR,\n#         )\n#         loaders.append(loader)\n#     return loaders\n\ndef make_client_loaders(trainset, client_indices: List[List[int]]) -> List[DataLoader]:\n    \"\"\"Creates DataLoaders for each client, handling empty clients safely.\"\"\"\n    loaders = []\n    for i, idxs in enumerate(client_indices):\n        # FIX: If a client has 0 samples, give them 1 random sample to prevent crash\n        # This is a common hack in Non-IID simulations to avoid breaking the loop\n        if len(idxs) == 0:\n            print(f\"WARNING: Client {i} received 0 samples. Assigning 1 random sample.\")\n            idxs = [0] # Assign index 0 arbitrarily just to keep it alive\n            \n        ds = Subset(trainset, idxs)\n        loader = DataLoader(\n            ds,\n            batch_size=cfg.BATCH_SIZE,\n            shuffle=True,\n            num_workers=cfg.NUM_WORKERS,\n            pin_memory=cfg.PIN_MEMORY,\n            persistent_workers=cfg.PERSISTENT_WORKERS,\n            prefetch_factor=cfg.PREFETCH_FACTOR,\n        )\n        loaders.append(loader)\n    return loaders\n","metadata":{"execution":{"iopub.execute_input":"2026-01-20T21:39:21.969102Z","iopub.status.busy":"2026-01-20T21:39:21.968806Z","iopub.status.idle":"2026-01-20T21:39:23.423896Z","shell.execute_reply":"2026-01-20T21:39:23.423062Z"},"papermill":{"duration":1.461438,"end_time":"2026-01-20T21:39:23.425766","exception":false,"start_time":"2026-01-20T21:39:21.964328","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"c914744f","cell_type":"markdown","source":"### Data Loading Note\nWe use `torchvision.datasets.CIFAR10` to load the data. This is the standard PyTorch method and fetches the same dataset as the Kaggle CIFAR-10 competition.\n*   **Automatic Download**: The `download=True` flag ensures the data is downloaded automatically if not present.\n*   **Format**: It loads the data directly into memory as efficient arrays, avoiding the need to extract and read thousands of individual image files from disk (which is common with the raw Kaggle download).","metadata":{"papermill":{"duration":0.004212,"end_time":"2026-01-20T21:39:23.433900","exception":false,"start_time":"2026-01-20T21:39:23.429688","status":"completed"},"tags":[]}},{"id":"4f09c8c0","cell_type":"code","source":"# ============================================================\n# 4) EVAL\n# ============================================================\n@torch.no_grad()\ndef evaluate(model: nn.Module, testloader: DataLoader) -> Tuple[float, float]:\n    \"\"\"Evaluates the model on the test set.\"\"\"\n    model.eval()\n    crit = nn.CrossEntropyLoss()\n    losses, accs = [], []\n\n    for x, y in testloader:\n        x = x.to(cfg.DEVICE, non_blocking=True)\n        y = y.to(cfg.DEVICE, non_blocking=True)\n        if cfg.CHANNELS_LAST:\n            x = x.contiguous(memory_format=torch.channels_last)\n\n        logits = model(x)\n        loss = crit(logits, y)\n        losses.append(loss.item())\n        accs.append(accuracy(logits, y))\n\n    return float(np.mean(losses)), float(np.mean(accs))\n\n\n# @torch.no_grad()\n# def evaluate(model: nn.Module, testloader: DataLoader):\n#     model.eval()\n#     criterion = nn.CrossEntropyLoss()\n\n#     total_loss = 0.0\n#     correct = 0\n#     total = 0\n\n#     for x, y in testloader:\n#         x = x.to(cfg.DEVICE, non_blocking=True)\n#         y = y.to(cfg.DEVICE, non_blocking=True)\n\n#         if cfg.CHANNELS_LAST:\n#             x = x.contiguous(memory_format=torch.channels_last)\n\n#         logits = model(x)\n#         loss = criterion(logits, y)\n\n#         total_loss += loss.item() * y.size(0)\n#         _, preds = torch.max(logits, 1)\n#         correct += (preds == y).sum().item()\n#         total += y.size(0)\n\n#     avg_loss = float(total_loss / total)\n#     accuracy = float(correct / total)\n\n#     return avg_loss, accuracy\n","metadata":{"execution":{"iopub.execute_input":"2026-01-20T21:39:23.443209Z","iopub.status.busy":"2026-01-20T21:39:23.442608Z","iopub.status.idle":"2026-01-20T21:39:23.449432Z","shell.execute_reply":"2026-01-20T21:39:23.448645Z"},"papermill":{"duration":0.013085,"end_time":"2026-01-20T21:39:23.450933","exception":false,"start_time":"2026-01-20T21:39:23.437848","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"0286ae5c","cell_type":"markdown","source":"### Training Loop Logic (The Core Algorithm)\n\nThis function `run_one_experiment` orchestrates the entire FL process. Here is the step-by-step logic mapping to the research concepts:\n\n1.  **Initialization**:\n    *   `client_bn_states`: A cache to store the local Batch Norm statistics for each client (Crucial for **FedBN**).\n\n2.  **Round Loop**:\n    *   **Client Selection**: A random fraction (`CLIENT_FRACTION`) of clients is selected.\n    *   **Progressive Schedule (ProgFL)**:\n        *   We determine which layer group is active: `active_group_idx = rnd % n_groups`.\n        *   This implements a sequential training schedule (e.g., Conv1 -> Layer1 -> ... -> FC -> Conv1...).\n\n3.  **Client Local Training**:\n    *   **Global Sync**: Client loads the global model parameters.\n    *   **FedBN Step (Restore)**: If FedBN is active, the client *overwrites* the global BN statistics with its own cached local BN statistics. This ensures the client normalizes features based on its own data distribution (addressing Non-IID).\n    *   **ProgFL Step (Freeze)**: We freeze all layers (`freeze_all`) and then unfreeze only the active group (`activate_group`).\n    *   **Training**: Standard SGD on the local dataset.\n    *   **FedBN Step (Save)**: After training, the updated BN statistics are saved back to the client's local cache.\n    *   **Communication**: The client sends back the model updates.\n        *   *Optimization*: We only calculate the size of the active parameters for `comm_bytes`.\n\n4.  **Server Aggregation**:\n    *   The server averages the weights received from clients.\n    *   **FedBN Logic**: BN parameters are *excluded* from aggregation. The global model keeps its own BN stats (or they are irrelevant if clients always use local ones).\n    *   **ProgFL Logic**: Only the active layers are updated in the global model.\n\n5.  **Evaluation**:\n    *   The global model is evaluated on the held-out test set.","metadata":{"papermill":{"duration":0.003848,"end_time":"2026-01-20T21:39:23.458562","exception":false,"start_time":"2026-01-20T21:39:23.454714","status":"completed"},"tags":[]}},{"id":"202e442e","cell_type":"code","source":"# ============================================================\n# 5) TRAINING LOOP (FAST, NO DEEPCOPY)\n# ============================================================\ndef run_one_experiment(exp_cfg: Dict, seed: int) -> pd.DataFrame:\n    \"\"\"\n    Runs a single FL experiment with the given configuration and seed.\n\n    Steps:\n    1. Setup: Initialize data, models, and clients.\n    2. Loop over rounds:\n       - Select clients.\n       - Determine active layers (for ProgFL).\n       - Client Training:\n         - Load global weights.\n         - Restore local BN stats (if FedBN).\n         - Train local model (only active layers if ProgFL).\n         - Save local BN stats.\n         - Calculate communication cost.\n       - Server Aggregation:\n         - Aggregate weights from clients (weighted average).\n         - Update global model.\n       - Evaluation: Evaluate global model on test set.\n    3. Save results and plots.\n    \"\"\"\n    set_seed(seed)\n\n    exp_name = exp_cfg[\"name\"]\n    outdir = os.path.join(cfg.OUTDIR, cfg.PROJECT, f\"{exp_name}_seed{seed}\")\n    ensure_dir(outdir)\n\n    trainset, testloader = get_data()\n    client_indices = dirichlet_split(trainset, cfg.NUM_CLIENTS, cfg.NONIID_ALPHA)\n    client_loaders = make_client_loaders(trainset, client_indices)\n\n    # Global model (server)\n    global_model = build_model().to(cfg.DEVICE)\n    layer_groups = define_layer_groups(global_model)\n    criterion = nn.CrossEntropyLoss()\n\n    # Reusable client model (NO deepcopy)\n    client_model = build_model().to(cfg.DEVICE)\n    scaler = GradScaler(\"cuda\", enabled=cfg.USE_AMP)\n\n    n_groups = len(layer_groups)\n\n    # FedBN cache: store BN state per client (CPU)\n    client_bn_states: List[Optional[Dict[str, torch.Tensor]]] = [\n        None for _ in range(cfg.NUM_CLIENTS)\n    ]\n\n    history = []\n    pbar = tqdm(range(cfg.ROUNDS), desc=f\"{exp_name} (seed {seed})\", ncols=120)\n\n    prev_test_loss, prev_test_acc = None, None\n\n    for rnd in pbar:\n        round_start = time.time()\n        num_sel = max(1, int(cfg.CLIENT_FRACTION * cfg.NUM_CLIENTS))\n        selected = np.random.choice(cfg.NUM_CLIENTS, num_sel, replace=False)\n\n        progressive = exp_cfg[\"progressive\"]\n        use_fedbn = exp_cfg[\"fedbn\"]\n        warmup_rounds = int(exp_cfg[\"warmup\"])\n\n        progressive_this_round = progressive and (rnd >= warmup_rounds)\n\n        # active_group_idx = -1\n        # active_param_names: Optional[Set[str]] = None\n        # if progressive_this_round:\n        #     active_group_idx = rnd % n_groups\n        #     active_param_names = set(layer_groups[active_group_idx])\n            \n        # --- REPLACEMENT BLOCK START ---\n        active_group_idx = -1\n        active_param_names: Optional[Set[str]] = None\n        \n        if progressive_this_round:\n            active_group_idx = rnd % n_groups\n            # Start with the layer group (Low or High)\n            active_param_names = set(layer_groups[active_group_idx])\n            \n            # CRITICAL FIX: Always add the Classifier (fc) to the active set\n            # This ensures the \"Head\" is always connected to the \"Body\"\n            fc_names = [n for n, _ in global_model.named_parameters() if \"fc\" in n]\n            active_param_names.update(fc_names)\n        # --- REPLACEMENT BLOCK END ---\n\n        # Cache global state on CPU for comm accounting\n        global_sd_cpu = {\n            k: v.detach().cpu() for k, v in global_model.state_dict().items()\n        }\n\n        client_sds = []  # CPU state dicts\n        client_weights = []  # dataset size weights\n\n        local_losses, local_accs, local_times = [], [], []\n        comm_bytes = 0\n\n        for cid in selected:\n            # Load global weights\n            client_model.load_state_dict(global_model.state_dict(), strict=True)\n\n            # Restore BN (FedBN)\n            if use_fedbn and client_bn_states[cid] is not None:\n                sd = client_model.state_dict()\n                for k, v in client_bn_states[cid].items():\n                    if k in sd:\n                        sd[k].copy_(v.to(sd[k].device))\n                client_model.load_state_dict(sd, strict=True)\n\n            # Freeze for progressive training\n            if progressive_this_round:\n                freeze_all(client_model)\n                activate_group(client_model, active_param_names)\n            else:\n                for p in client_model.parameters():\n                    p.requires_grad = True\n\n            # Optimizer for active params only\n            optimizer = optim.SGD(\n                filter(lambda p: p.requires_grad, client_model.parameters()),\n                lr=cfg.LR,\n                momentum=cfg.MOMENTUM,\n                weight_decay=cfg.WEIGHT_DECAY,\n            )\n\n            # # Local train\n            # client_model.train()\n\n            # Local train\n            # --- START FIX: Smart Mode Switching ---\n            client_model.train()  # 1. Default: Set everything to train\n            \n            if progressive_this_round:\n                client_model.eval()   # 2. Lock everything to EVAL (stops BN updates)\n                \n                # 3. Always unlock the Classifier\n                client_model.fc.train()\n                \n                # 4. Unlock ONLY the active group\n                if active_group_idx == 0:  # Low Level\n                    client_model.conv1.train()\n                    client_model.bn1.train()\n                    client_model.layer1.train()\n                    client_model.layer2.train()\n                elif active_group_idx == 1: # High Level\n                    client_model.layer3.train()\n                    client_model.layer4.train()\n            # --- END FIX ---\n\n            # t0 = time.time()\n            # ... (The rest of your code continues unchanged)\n            t0 = time.time()\n            tot_loss, tot_acc, steps = 0.0, 0.0, 0\n\n            for _ in range(cfg.LOCAL_EPOCHS):\n                for x, y in client_loaders[cid]:\n                    x = x.to(cfg.DEVICE, non_blocking=True)\n                    y = y.to(cfg.DEVICE, non_blocking=True)\n                    if cfg.CHANNELS_LAST:\n                        x = x.contiguous(memory_format=torch.channels_last)\n\n                    optimizer.zero_grad(set_to_none=True)\n\n                    with autocast(\"cuda\", enabled=cfg.USE_AMP):\n                        logits = client_model(x)\n                        loss = criterion(logits, y)\n\n                    scaler.scale(loss).backward()\n                    scaler.step(optimizer)\n                    scaler.update()\n\n                    tot_loss += loss.item()\n                    tot_acc += accuracy(logits.detach(), y)\n                    steps += 1\n\n            elapsed = time.time() - t0\n            local_losses.append(tot_loss / max(1, steps))\n            local_accs.append(tot_acc / max(1, steps))\n            local_times.append(elapsed)\n\n            # Save BN state back to FedBN cache\n            if use_fedbn:\n                bn_sd = {}\n                for k, v in client_model.state_dict().items():\n                    if is_bn_key(k):\n                        bn_sd[k] = v.detach().cpu().clone()\n                client_bn_states[cid] = bn_sd\n\n            # Export client weights to CPU\n            cm_sd_cpu = {\n                k: v.detach().cpu() for k, v in client_model.state_dict().items()\n            }\n            client_sds.append(cm_sd_cpu)\n            client_weights.append(len(client_indices[cid]))\n\n            # Communication accounting (client->server upload)\n            if progressive_this_round:\n                upload_keys = set(active_param_names)\n                if use_fedbn:\n                    upload_keys = {k for k in upload_keys if not is_bn_key(k)}\n                comm_bytes += model_bytes(cm_sd_cpu, upload_keys)\n            else:\n                if use_fedbn:\n                    keys = {k for k in cm_sd_cpu.keys() if not is_bn_key(k)}\n                    comm_bytes += model_bytes(cm_sd_cpu, keys)\n                else:\n                    comm_bytes += model_bytes(cm_sd_cpu, None)\n\n        # Server aggregation\n        w = np.array(client_weights, dtype=np.float32)\n        w = (w / w.sum()).tolist()\n\n        global_sd = global_model.state_dict()\n        for k in list(global_sd.keys()):\n            if use_fedbn and is_bn_key(k):\n                continue\n            if (\n                progressive_this_round\n                and (active_param_names is not None)\n                and (k not in active_param_names)\n            ):\n                continue\n            global_sd[k] = sum(wi * csd[k] for wi, csd in zip(w, client_sds))\n        global_model.load_state_dict(global_sd, strict=True)\n\n        comm_mb = comm_bytes / (1024 * 1024)\n        round_time = time.time() - round_start\n\n        # Evaluate less frequently\n        if (rnd % cfg.EVAL_EVERY) == 0 or rnd == cfg.ROUNDS - 1:\n            test_loss, test_acc = evaluate(global_model, testloader)\n            prev_test_loss, prev_test_acc = test_loss, test_acc\n        else:\n            test_loss, test_acc = prev_test_loss, prev_test_acc\n\n        history.append(\n            dict(\n                round=rnd,\n                test_loss=test_loss,\n                test_acc=test_acc,\n                round_time=round_time,\n                comm_mb=comm_mb,\n                active_group=active_group_idx,\n                progressive=int(progressive_this_round),\n                warmup=int(rnd < warmup_rounds),\n                mean_local_acc=float(np.mean(local_accs)),\n                mean_local_loss=float(np.mean(local_losses)),\n                mean_local_time=float(np.mean(local_times)),\n            )\n        )\n\n        pbar.set_postfix(\n            {\n                \"acc\": f\"{test_acc:.3f}\",\n                \"commMB\": f\"{comm_mb:.2f}\",\n                \"grp\": active_group_idx if progressive_this_round else \"ALL\",\n                \"WU\": \"Y\" if rnd < warmup_rounds else \"N\",\n            }\n        )\n\n    df = pd.DataFrame(history)\n    if cfg.SAVE_CSV:\n        df.to_csv(os.path.join(outdir, \"history.csv\"), index=False)\n\n    if cfg.SAVE_PLOTS:\n        plotdir = os.path.join(outdir, \"plots\")\n        ensure_dir(plotdir)\n\n        plt.figure()\n        plt.plot(df[\"round\"], df[\"test_acc\"])\n        plt.xlabel(\"Round\")\n        plt.ylabel(\"Test Accuracy\")\n        plt.title(f\"Accuracy vs Round ({exp_name})\")\n        savefig_both(os.path.join(plotdir, \"acc_vs_round\"))\n\n        plt.figure()\n        plt.plot(df[\"round\"], df[\"comm_mb\"].cumsum(), label=\"CumComm\")\n        plt.xlabel(\"Round\")\n        plt.ylabel(\"Cumulative Comm (MB)\")\n        plt.title(f\"Cumulative Communication ({exp_name})\")\n        savefig_both(os.path.join(plotdir, \"cumcomm_vs_round\"))\n\n        plt.figure()\n        plt.plot(df[\"comm_mb\"].cumsum(), df[\"test_acc\"])\n        plt.xlabel(\"Cumulative Communication (MB)\")\n        plt.ylabel(\"Accuracy\")\n        plt.title(f\"Accuracy vs Cumulative Communication ({exp_name})\")\n        savefig_both(os.path.join(plotdir, \"acc_vs_cumcomm\"))\n\n    return df","metadata":{"execution":{"iopub.execute_input":"2026-01-20T21:39:23.467522Z","iopub.status.busy":"2026-01-20T21:39:23.467236Z","iopub.status.idle":"2026-01-20T21:39:23.494263Z","shell.execute_reply":"2026-01-20T21:39:23.493516Z"},"papermill":{"duration":0.033548,"end_time":"2026-01-20T21:39:23.495791","exception":false,"start_time":"2026-01-20T21:39:23.462243","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"a03d70ff","cell_type":"code","source":"# ============================================================\n# 6) SUMMARY + LATEX\n# ============================================================\ndef compute_summary(df: pd.DataFrame) -> Dict:\n    \"\"\"Computes summary statistics for an experiment.\"\"\"\n    best_acc = df[\"test_acc\"].max()\n    best_round = int(df.loc[df[\"test_acc\"].idxmax(), \"round\"])\n    total_comm = df[\"comm_mb\"].sum()\n    total_time = df[\"round_time\"].sum()\n    auc_acc = float(\n        np.trapz(df[\"test_acc\"].values, df[\"round\"].values)\n        / (df[\"round\"].iloc[-1] + 1e-12)\n    )\n    return dict(\n        best_acc=best_acc,\n        best_round=best_round,\n        total_comm=total_comm,\n        total_time=total_time,\n        auc_acc=auc_acc,\n    )\n\n\ndef generate_latex_table(\n    summary_df: pd.DataFrame, outdir: str, fname=\"latex_table.tex\"\n):\n    \"\"\"Generates a LaTeX table from the summary dataframe.\"\"\"\n    ensure_dir(outdir)\n    path = os.path.join(outdir, fname)\n\n    with open(path, \"w\") as f:\n        f.write(\"\\\\begin{table}[t]\\n\\\\centering\\n\\\\small\\n\")\n        f.write(\"\\\\begin{tabular}{l c c c}\\n\\\\toprule\\n\")\n        f.write(\"Experiment & Best Acc & Total Comm (MB) & Total Time (s) \\\\\\\\\\n\")\n        f.write(\"\\\\midrule\\n\")\n\n        for _, row in summary_df.iterrows():\n            f.write(\n                f\"{row['experiment']} & {row['best_acc']:.4f} & {row['total_comm']:.2f} & {row['total_time']:.2f} \\\\\\\\\\n\"\n            )\n\n        f.write(\"\\\\bottomrule\\n\\\\end{tabular}\\n\")\n        f.write(\"\\\\caption{Federated Learning results.}\\n\")\n        f.write(\"\\\\label{tab:fl}\\n\\\\end{table}\\n\")\n\n    print(f\"LaTeX table saved: {os.path.abspath(path)}\")","metadata":{"execution":{"iopub.execute_input":"2026-01-20T21:39:23.505193Z","iopub.status.busy":"2026-01-20T21:39:23.504475Z","iopub.status.idle":"2026-01-20T21:39:23.511575Z","shell.execute_reply":"2026-01-20T21:39:23.510803Z"},"papermill":{"duration":0.013556,"end_time":"2026-01-20T21:39:23.513178","exception":false,"start_time":"2026-01-20T21:39:23.499622","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"1d98be00","cell_type":"code","source":"# ============================================================\n# 7) RUN ALL\n# ============================================================\ndef run_all():\n    \"\"\"Runs all configured experiments.\"\"\"\n    ensure_dir(os.path.join(cfg.OUTDIR, cfg.PROJECT))\n\n    seeds = cfg.SEEDS if cfg.MULTI_SEED else (cfg.SEED,)\n    all_summaries = []\n\n    for exp in cfg.EXPERIMENTS[: cfg.MAX_EXPERIMENTS]:\n        for sd in seeds:\n            df = run_one_experiment(exp, sd)\n            summ = compute_summary(df)\n            summ[\"experiment\"] = exp[\"name\"]\n            summ[\"seed\"] = sd\n            all_summaries.append(summ)\n\n    summary_df = pd.DataFrame(all_summaries)\n    summary_path = os.path.join(cfg.OUTDIR, cfg.PROJECT, \"ALL_SUMMARIES.csv\")\n    summary_df.to_csv(summary_path, index=False)\n\n    combined_dir = os.path.join(cfg.OUTDIR, cfg.PROJECT, \"COMBINED\")\n    ensure_dir(combined_dir)\n\n    generate_latex_table(summary_df, combined_dir)\n\n    print(\"\\nDONE.\")\n    print(f\"Summary CSV: {os.path.abspath(summary_path)}\")\n    print(f\"Outputs: {os.path.abspath(combined_dir)}\")\n\n\nif __name__ == \"__main__\":\n    # speed knobs\n    if cfg.CUDNN_BENCHMARK:\n        torch.backends.cudnn.benchmark = True\n    if cfg.TF32 and torch.cuda.is_available():\n        torch.set_float32_matmul_precision(\"high\")\n\n    run_all()\n    \n    remove_directory(cfg.Global_EXTRACT_PATH)","metadata":{"execution":{"iopub.execute_input":"2026-01-20T21:39:23.522525Z","iopub.status.busy":"2026-01-20T21:39:23.521787Z","iopub.status.idle":"2026-01-21T06:14:21.133658Z","shell.execute_reply":"2026-01-21T06:14:21.132751Z"},"papermill":{"duration":30897.618286,"end_time":"2026-01-21T06:14:21.135352","exception":false,"start_time":"2026-01-20T21:39:23.517066","status":"completed"},"tags":[]},"outputs":[],"execution_count":null},{"id":"2775fe45","cell_type":"markdown","source":"## Flow and Summary\n\n### Implementation Flow\n1.  **Initialization**:\n    -   The `CFG` class defines all hyperparameters (clients, rounds, learning rate, etc.).\n    -   The `run_all` function iterates through defined experiments (e.g., FedAvg, ProgFL).\n\n2.  **Experiment Execution (`run_one_experiment`)**:\n    -   **Data Setup**: CIFAR-10 is loaded and split among clients using a Dirichlet distribution (Non-IID).\n    -   **Model Setup**: A global ResNet18 model is initialized. A reusable client model is created to save memory.\n    -   **Training Loop**:\n        -   **Client Selection**: A fraction of clients is selected for the round.\n        -   **Progressive Logic**: If enabled, only specific layers are active (unfrozen) based on the round number.\n        -   **Client Update**:\n            -   Global weights are loaded into the client model.\n            -   If FedBN is used, the client's specific Batch Norm statistics are restored.\n            -   The client trains on its local data (only active layers are updated).\n            -   Updated weights and BN stats are saved.\n        -   **Aggregation**: The server aggregates the weights from all participating clients. BN stats are *not* aggregated (FedBN).\n        -   **Evaluation**: The global model is evaluated on the test set periodically.\n\n3.  **Reporting**:\n    -   Results (accuracy, loss, communication cost) are saved to CSV.\n    -   Plots are generated for accuracy and communication.\n    -   A summary LaTeX table is created comparing experiments.\n\n### Summary\nThis implementation efficiently simulates Federated Learning by reusing a single client model instance to save memory. It incorporates **FedBN** to handle non-IID data by keeping batch normalization statistics local to each client. **Progressive Layer Training** is implemented to reduce communication costs by training only a subset of layers in each round. The code is optimized for speed using AMP and efficient data loading.","metadata":{"papermill":{"duration":0.08707,"end_time":"2026-01-21T06:14:21.314569","exception":false,"start_time":"2026-01-21T06:14:21.227499","status":"completed"},"tags":[]}}]}