1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_schedule.py",
4  "start_line": 246,
5  "end_line": 269,
6  "content": "def get_lr(epoch, total_epochs, base_lr, config):\n    \"\"\"Compute learning rate for the given epoch.\n\n    Called once per epoch to set the learning rate for all parameter groups.\n\n    Args:\n        epoch: current epoch (0-indexed, ranges from 0 to total_epochs-1)\n        total_epochs: total number of training epochs\n        base_lr: initial learning rate (from --lr flag, default 0.1)\n        config: dict with keys:\n            - arch: str ('resnet20', 'resnet56', 'vgg16bn', 'mobilenetv2')\n            - dataset: str ('cifar10', 'cifar100', 'fmnist')\n\n    Returns:\n        float: learning rate to use for this epoch\n\n    Design considerations:\n        - Warmup phase to stabilize early training\n        - Decay shape (step, cosine, polynomial, exponential, ...)\n        - Final learning rate (decay to zero vs small constant)\n        - Architecture/dataset-aware scheduling\n        - Interaction with momentum and weight decay\n    \"\"\"\n    # Architecture-aware parameters\n    arch = config['arch']\n    dataset = config['dataset']\n    \n    # Warmup epochs: deeper networks need longer warmup\n    if arch == 'resnet20':\n        warmup_epochs = 5\n        plateau_epochs = 20  # epochs at high LR after warmup\n        decay_type = 'cosine'\n    elif arch == 'resnet56':\n        warmup_epochs = 10\n        plateau_epochs = 30\n        decay_type = 'cosine'\n    elif arch == 'mobilenetv2':\n        warmup_epochs = 3\n        plateau_epochs = 15\n        decay_type = 'cosine'\n    else:  # vgg16bn\n        warmup_epochs = 7\n        plateau_epochs = 25\n        decay_type = 'cosine'\n    \n    # Adjust based on dataset difficulty\n    if dataset == 'cifar100':\n        warmup_epochs = int(warmup_epochs * 1.2)  # longer warmup for more classes\n        plateau_epochs = int(plateau_epochs * 1.1)\n    elif dataset == 'fmnist':\n        # FashionMNIST is simpler, shorter warmup\n        warmup_epochs = max(2, warmup_epochs // 2)\n        plateau_epochs = max(10, plateau_epochs // 2)\n    \n    # Ensure warmup + plateau doesn't exceed total epochs\n    warmup_epochs = min(warmup_epochs, total_epochs // 4)\n    plateau_epochs = min(plateau_epochs, total_epochs - warmup_epochs - 10)  # leave room for decay\n    \n    # Phase determination\n    if epoch < warmup_epochs:\n        # Linear warmup from base_lr/10 to base_lr\n        return base_lr * (epoch + 1) / warmup_epochs\n    elif epoch < warmup_epochs + plateau_epochs:\n        # Plateau at base_lr\n        return base_lr\n    else:\n        # Decay phase\n        decay_start_epoch = warmup_epochs + plateau_epochs\n        decay_length = total_epochs - decay_start_epoch\n        progress = (epoch - decay_start_epoch) / decay_length\n        \n        # Cosine decay to base_lr/100\n        final_lr = base_lr / 100.0\n        decay_lr = final_lr + 0.5 * (base_lr - final_lr) * (1 + math.cos(math.pi * progress))\n        \n        # Add small cyclic perturbation (1% magnitude) to help escape local minima\n        cycle = math.sin(2 * math.pi * 3 * progress) * 0.01 * decay_lr\n        return decay_lr + cycle"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_schedule.py",
4  "start_line": 246,
5  "end_line": 323,
6  "content": "def get_lr(epoch, total_epochs, base_lr, config):\n    \"\"\"Adaptive OneCycle schedule with architecture/dataset-aware parameters.\n\n    Phase 1 (warmup): cosine increase from base_lr/div_factor to base_lr.\n    Phase 2 (annealing): cosine decrease from base_lr to final_lr.\n\n    Parameters:\n        pct_start: fraction of total epochs for warmup phase.\n                   Deeper networks have later peak (larger pct_start).\n        div_factor: initial LR divisor, larger for more complex datasets.\n        final_div: final LR divisor, larger for easier datasets.\n    \"\"\"\n    arch = config['arch']\n    dataset = config['dataset']\n    \n    # Default OneCycle parameters (original)\n    pct_start = 0.3\n    div_factor = 25.0\n    final_div = 25.0\n    \n    # Architecture adaptations\n    if arch == 'resnet20':\n        # Shallow network: peak earlier, less aggressive warmup\n        pct_start = 0.25\n        div_factor = 20.0\n        final_div = 20.0\n    elif arch == 'resnet56':\n        # Deeper network: peak later, more aggressive warmup (lower initial LR)\n        pct_start = 0.4\n        div_factor = 30.0\n        final_div = 30.0\n    elif arch == 'mobilenetv2':\n        # Lightweight network: peak earlier, moderate warmup\n        pct_start = 0.2\n        div_factor = 15.0\n        final_div = 15.0\n    else:  # vgg16bn\n        pct_start = 0.35\n        div_factor = 25.0\n        final_div = 25.0\n    \n    # Dataset adaptations\n    if dataset == 'cifar100':\n        # More classes: need slower warmup (larger div_factor) and smaller final LR\n        div_factor *= 1.5\n        final_div *= 1.5\n    elif dataset == 'fmnist':\n        # Simpler dataset: can use more aggressive schedule\n        div_factor *= 0.8\n        final_div *= 0.8\n    \n    # Compute LR\n    min_lr = base_lr / div_factor\n    final_lr = base_lr / final_div\n    \n    progress = epoch / max(total_epochs - 1, 1)\n    \n    if progress <= pct_start:\n        # Warmup phase: cosine from min_lr to base_lr\n        t = progress / pct_start\n        return min_lr + (base_lr - min_lr) * 0.5 * (1 + math.cos(math.pi * (1 - t)))\n    else:\n        # Anneal phase: cosine from base_lr to final_lr\n        t = (progress - pct_start) / (1 - pct_start)\n        return final_lr + (base_lr - final_lr) * 0.5 * (1 + math.cos(math.pi * t))"
7}

1{
2  "op": "replace",
3  "filename": "pytorch-vision/custom_schedule.py",
4  "start_line": 246,
5  "end_line": 310,
6  "content": "def get_lr(epoch, total_epochs, base_lr, config):\n    \"\"\"Adaptive Plateau Cosine (APC) schedule.\n    \n    Warmup (linear) → Plateau (constant high LR) → Cosine decay to final LR.\n    \n    Hyperparameters adapt to architecture depth and dataset complexity.\n    \"\"\"\n    arch = config['arch']\n    dataset = config['dataset']\n    \n    # Warmup fraction of total epochs\n    if arch == 'resnet20':\n        warmup_ratio = 0.03\n        plateau_ratio = 0.15\n    elif arch == 'resnet56':\n        warmup_ratio = 0.05\n        plateau_ratio = 0.20\n    elif arch == 'mobilenetv2':\n        warmup_ratio = 0.02\n        plateau_ratio = 0.10\n    else:  # vgg16bn\n        warmup_ratio = 0.04\n        plateau_ratio = 0.18\n    \n    # Dataset adjustments\n    if dataset == 'cifar100':\n        # More classes, need longer plateau for exploration\n        plateau_ratio *= 1.2\n    elif dataset == 'fmnist':\n        # Simpler dataset, shorter plateau\n        plateau_ratio *= 0.8\n    \n    # Ensure ratios are within reasonable bounds\n    warmup_ratio = min(0.1, max(0.01, warmup_ratio))\n    plateau_ratio = min(0.3, max(0.05, plateau_ratio))\n    \n    warmup_epochs = int(warmup_ratio * total_epochs)\n    plateau_epochs = int(plateau_ratio * total_epochs)\n    \n    # Ensure we have at least 10 epochs for decay\n    if warmup_epochs + plateau_epochs > total_epochs - 10:\n        plateau_epochs = total_epochs - warmup_epochs - 10\n    \n    # Final LR = base_lr / 100\n    final_lr = base_lr / 100.0\n    \n    # Phase determination\n    if epoch < warmup_epochs:\n        # Linear warmup from base_lr/10 to base_lr\n        return base_lr * (epoch + 1) / warmup_epochs\n    elif epoch < warmup_epochs + plateau_epochs:\n        # Plateau at base_lr\n        return base_lr\n    else:\n        # Cosine decay from base_lr to final_lr\n        decay_start_epoch = warmup_epochs + plateau_epochs\n        decay_length = total_epochs - decay_start_epoch\n        progress = (epoch - decay_start_epoch) / decay_length\n        return final_lr + 0.5 * (base_lr - final_lr) * (1 + math.cos(math.pi * progress))"
7}

1{
2  "n": 1
3}