GPU显存不足可用8bit adam优化器节省显存

最新推荐文章于 2024-09-24 14:22:37 发布

nangua06

最新推荐文章于 2024-09-24 14:22:37 发布

阅读量645

点赞数 3

文章标签：深度学习人工智能

本文链接：https://blog.csdn.net/nangua06/article/details/137152222

版权

本文介绍了如何通过引入8-bitAdam优化器（如AdamW8bit）来节省深度学习训练中的显存消耗，同时提到了安装bitsandbytes库的必要性。

摘要由CSDN通过智能技术生成

#用8bit adam优化器节省显存

# Initialize the optimizer
if cfg.solver.use_8bit_adam:
    try:
        import bitsandbytes as bnb
    except ImportError:
        raise ImportError(
            "Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
        )

    optimizer_cls = bnb.optim.AdamW8bit
else:
    optimizer_cls = torch.optim.AdamW

optimizer = optimizer_cls(
    trainable_params,
    lr=learning_rate,
    betas=(cfg.solver.adam_beta1, cfg.solver.adam_beta2),
    weight_decay=cfg.solver.adam_weight_decay,
    eps=cfg.solver.adam_epsilon,
)