Bugfix/fix warmup scheduler setting error

[zephyr-sh]

✨ Summary

This PR introduces a polynomial LR scheduler with warmup and enhances the warmup scheduler infrastructure with better modularity, validation, and test coverage.

✅ Major Changes

1. Refactored PolynomialLRWarmup

• Migrated to inherit from torch.optim.lr_scheduler.LRScheduler (modern PyTorch API).

• Rewritten with:

  • Explicit closed-form logic via get_closed_form()
  • Strict input validation for warmup_iters, total_iters, and power
  • Support for both legacy API (step(epoch)) and modern chainable API

• Behavior:

  • Linearly increases LR from 0 → base_lr over warmup_iters
  • Then polynomial decay to 0 across remaining total_iters - warmup_iters
  • Fixed to clamp after total_iters

2. Introduced WrappedLRScheduler

• A general-purpose warmup wrapper:

  • Handles linear warmup with configurable multiplier
  • Delegates to any other scheduler (after_scheduler) after warmup

• Works with any LRScheduler (e.g. ReduceLROnPlateau, MultiStepLR, etc.)

• Compatible with step() API (including metric-based schedulers)

• Includes strong argument checks and lifecycle safety (e.g. base LR re-init once)

3. Added MultiStepLRWarmUp factory function

• Combines:

  • Linear warmup via WrappedLRScheduler
  • Followed by decay via MultiStepLR

• Clean and reusable for config-based instantiation

4. Registry Integration

• Registered all schedulers (PolynomialLRWarmup, WrappedLRScheduler, MultiStepLRWarmUp) to OPTIMIZERS registry
• Updated __init__.py to reflect the new module inclusion

---

🧪 Tests

tests/base/optim/test_polynomial_lr_warmup.py

• Validates:

  • Input validation
  • Correct LR output in warmup and decay phases
  • Behavior after total_iters
  • Compatibility with .step() chainable API

tests/base/optim/test_warm_up.py

• Covers:

  • WrappedLRScheduler warmup-only behavior
  • Warmup w/ multiplier > 1
  • Correct delegation to after_scheduler (e.g. MultiStepLR)
  • Correct final LR holding behavior
  • Factory behavior and edge cases for MultiStepLRWarmUp

All tests pass.

---

🧩 Motivation

• PyTorch is progressively deprecating _LRScheduler in favor of LRScheduler.
• Polynomial warmup+decay is a common pattern not yet fully supported natively.
• Reusable warmup wrappers simplify composition and config loading.
• Provides a reliable warmup scheduling foundation for future models.

---

🔄 API Summary

# Polynomial warmup + decay
PolynomialLRWarmup(optimizer, warmup_iters=5, total_iters=30, power=1.0)

# General warmup wrapper
WrappedLRScheduler(optimizer, milestone=5, multiplier=1.0, after_scheduler=...)

# Factory: warmup + MultiStep
MultiStepLRWarmUp(optimizer, milestones=[10, 20], warmup_milestone=3, gamma=0.1)

---

📎 Miscellaneous

• Ensures compatibility with both old and new scheduler API
• Fully documented with docstrings and param comments

---

✅ Checklist

• Feature implemented with modular, typed design
• Integrated with registry
• Unit tests added
• Passes all CI checks
• Compatible with both legacy and chainable .step() API

[github-actions]

Coverage Report

File	Stmts	Miss	Cover	Missing
chameleon
__init__.py	6	0	100%
chameleon/base
__init__.py	7	0	100%
power_module.py	40	2	2	95%
utils.py	46	4	4	91%
chameleon/base/blocks
__init__.py	1	0	100%
conv_block.py	46	2	2	96%
mamba_block.py	2	2	2	0%
vit_block.py	2	2	2	0%
chameleon/base/components
__init__.py	5	0	100%
activation.py	36	4	4	89%
dropout.py	5	0	100%
loss.py	82	34	34	59%
norm.py	19	0	100%
pooling.py	20	0	100%
chameleon/base/layers
__init__.py	5	0	100%
aspp.py	21	0	100%
grl.py	22	0	100%
selayer.py	18	0	100%
vae.py	23	0	100%
weighted_sum.py	23	0	100%
chameleon/base/ops
__init__.py	1	0	100%
positional_encoding.py	14	0	100%
chameleon/base/optim
__init__.py	9	0	100%
polynomial_lr_warmup.py	27	0	100%
warm_up.py	39	1	1	97%
chameleon/metrics
__init__.py	1	0	100%
normalized_levenshtein_similarity.py	51	29	29	43%
chameleon/modules
__init__.py	2	0	100%
chameleon/modules/backbones
__init__.py	3	0	100%
gpunet.py	54	2	2	96%
timm.py	7	0	100%
chameleon/modules/necks
__init__.py	2	0	100%
bifpn.py	87	2	2	98%
fpn.py	58	2	2	97%
pafpn.py	0	0	100%
chameleon/registry
__init__.py	2	0	100%
registry.py	92	12	12	87%
root.py	26	5	5	81%
chameleon/tools
__init__.py	1	0	100%
chameleon/tools/calflops
__init__.py	2	0	100%
calculate_pipline.py	182	58	58	68%
flops_counter.py	52	16	16	69%
pytorch_ops.py	336	153	153	54%
utils.py	112	76	76	32%
TOTAL	1589	406	74%

Tests	Skipped	Failures	Errors	Time
104	0 💤	0 ❌	0 🔥	5.342s ⏱️

[kunkunlin1221]

LGTM