kozistr/pytorch_optimizer

pytorch‑optimizer – 一个即用型 PyTorch 优化器、调度器与损失函数的工具箱

什么是它 – 一个 Python 包，通过单一且一致的 API，将超过一百种研究级优化器（例如 AdamP, Ranger21, Lion, SAM）以及一些学习率调度器和损失函数封装在一起。它旨在为任何使用 PyTorch 训练深度学习模型的开发者，让其无需编写自定义的样板代码即可尝试最新的优化技巧。

为什么你可能会感兴趣

广泛的覆盖范围 – 该库收集了许多近期发表的变体，这些变体通常分散在各个独立的仓库或论文中。
统一的接口 – 所有组件都可以直接实例化，通过名称查找（load_optimizer('adamp')），或使用辅助工具 create_optimizer() 来构建，该工具还可串接选用的技巧，例如梯度中心化（Gradient Centralization）、Lookahead 或为了加速而使用的 foreach。
可选的生态系统钩子 – 当安装了 bitsandbytes, q‑galore‑torch 与 torchao 时，可以与它们协同工作，为您提供无需额外代码的低精度或量化优化功能。
生产就绪 – 经过 CI 测试、型别注解，并已在 PyPI 上发布（Apache‑2.0 许可证）。

安装方式

pip install pytorch-optimizer   # requires Python ≥3.8 and PyTorch ≥1.10

如果您需要选用的整合功能，请单独安装它们（例如 pip install bitsandbytes）。

快速开始

from pytorch_optimizer import AdamP, create_optimizer, load_optimizer

model = MyNet()
# 1️⃣ 直接使用类
opt = AdamP(model.parameters(), lr=1e-3)

# 2️⃣ 通过字符串名称加载
opt = load_optimizer('adamp')(model.parameters(), lr=1e-3)


# 3️⃣ 使用同时添加常用技巧的辅助工具
opt = create_optimizer(
    model,
    optimizer_ 
    optimizer_name='adamp',
    lr=1e-3,
    weight_decay=1e-3,
    use_gc=True,          # Gradient Centralization
    use_lookahead=True,   # Lookahead wrapper
)

You can also fetch the class via torch.hub:

opt_cls = torch.hub.load('kozistr/pytorch_optimizer', 'adamp')
opt = opt_cls(model.parameters(), lr=1e-3)

探索有哪些可用选项

from pytorch_optimizer import (
    get_supported_optimizers,
    get_supported_lr_schedulers,
    get_supported_loss_functions,
)

print(get_supported_optimizers())               # list of all 100+ optimizers

print(    print(get_supported_optimizers('adam*'))        # filter by pattern
print(get_supported_lr_schedulers('cosine*'))   # scheduler names
print(get_supported_loss_functions('*focal*')) # loss names

The README 包含一个长表格，将每个优化器链接到其原始代码仓库和论文，以便您可以追溯研究来源。

典型使用场景工作流程

挑选一个符合您训练方案的优化器（例如 Ranger21）。
可选性地启用 use_gc 或是 use_lookahead 以获得额外的稳定性。
将优化器接入您平常的 PyTorch 训练循环中 – 不需要其他变动。
如果您需要学习率调度，请使用 get_supported_lr_schedulers() 取得一个，并将其传递给 torch.optim.lr_scheduler 或使用库内置的封装器。

谁在维护它 – 由 kozistr 团队积极维护，提供 CI 标志、代码覆盖率报告以及在 PyPI 上的定期发布。

底线 – pytorch‑optimizer 是一个实用、面向生产环境的现代优化器及相关工具的集合，让 PyTorch 用户可以快速且安全地尝试最新的训练算法。