pytorch-accelerated
pytorch-accelerated 是一个轻量级库,旨在通过提供一个最小但可扩展的训练循环来加速 PyTorch 模型的训练过程
该训练循环封装在一个单独的 Trainer 对象中,灵活性足以满足大多数使用场景,并能在无需修改代码的情况下利用不同的硬件
选项。
pytorch-accelerated 提供精简的功能集,并极为强调 简单性 和 透明性,
让用户能够准确了解底层发生了什么,但无需自己编写和维护模板代码!
主要特点包括:
- 一个简单且独立,但易于定制的训练循环,能够在简单场景下开箱即用;
- 处理设备分配、混合精度、DeepSpeed 集成、多 GPU 和分布式训练,无需更改代码。
- 使用纯 PyTorch 组件,无需额外修改或包装,并能轻松与其他流行库如 timm、
- 精简的小型 API 确保现有 PyTorch 用户的学习曲线最小化。
pytorch-accelerated 自豪而透明地构建在
Hugging Face Accelerate 之上,后者负责
数据在设备间的移动和训练配置的启动。当自定义训练器或启动
训练时,建议用户查阅 Accelerate 文档
以了解所有可用选项;Accelerate 提供了诸如收集张量和梯度裁剪等便捷函数,
其使用示例可在 pytorch-accelerated 的
examples 文件夹中找到!
想要了解本库背后的动机以及详细的入门指南,请参阅这篇博客文章。
安装
可以通过以下命令从 pip 安装 pytorch-accelerated:
pip install pytorch-accelerated
为了使软件包尽可能精简,运行示例所需的软件包默认不包括在内。要包含这些软件包,您可以使用以下命令:pip install pytorch-accelerated[examples]快速开始
要开始,只需导入并使用pytorch加速的Trainer,如下面的代码片段所示,
然后使用下面描述的
accelerate CLI
启动训练。
# examples/core/train_mnist.py
import osfrom torch import nn, optim
from torch.utils.data import random_split
from torchvision import transforms
from torchvision.datasets import MNIST
from pytorch_accelerated import Trainer
class MNISTModel(nn.Module):
def __init__(self):
super().__init__()
self.main = nn.Sequential(
nn.Linear(in_features=784, out_features=128),
nn.ReLU(),
nn.Linear(in_features=128, out_features=64),
nn.ReLU(),
nn.Linear(in_features=64, out_features=10),
)
def forward(self, input):
return self.main(input.view(input.shape[0], -1))
def main():
dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())
train_dataset, validation_dataset, test_dataset = random_split(dataset, [50000, 5000, 5000])
model = MNISTModel()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
loss_func = nn.CrossEntropyLoss()
trainer = Trainer(
model,
loss_func=loss_func,
optimizer=optimizer,
)
trainer.train(
train_dataset=train_dataset,
eval_dataset=validation_dataset,
num_epochs=8,
per_device_batch_size=32,
)
trainer.evaluate(
dataset=test_dataset,
per_device_batch_size=64,
)
if __name__ == "__main__":
main()
要使用accelerate CLI启动训练,
在你的机器上运行: accelerate config --config_file accelerate_config.yaml
并回答所提出的问题。这将生成一个配置文件,用于在执行
accelerate launch --config_file accelerate_config.yaml train.py [--training-args]
时正确设置默认选项。
注意:使用accelerate CLI是完全可选的,训练也可以通过以下常规方式启动:
python train.py / python -m torch.distributed ...
这取决于你的基础设施配置,适合希望对启动命令进行更细粒度控制的用户。
更复杂的训练示例可以在示例文件夹中查看 这里。
或者,如果你更愿意先了解核心概念,可以在文档中找到。
使用说明
pytorch-accelerated的目标用户是谁?
- 熟悉PyTorch但希望避免编写常见训练循环样板代码,以便专注于训练循环有趣部分的用户。
- 喜欢并且擅长选择和创建自己的模型、损失函数、优化器和数据集的用户。
- 重视简单且精简的功能集,行为易于调试、理解和推理的用户!
什么时候不应该使用pytorch-accelerated?
- 如果你正在寻找一套端到端的解决方案,涵盖从数据加载到推理的所有环节,
pytorch-accelerated仅专注于训练过程,其他所有
方面由用户自行负责。
- 如果你想自己编写整个训练循环,只是不想处理所有设备管理的头疼问题,
Trainer的每个部分,训练循环本质上被拆分为多个
不同的方法你可能需要重写。但是,在你开始之前,编写那些 for 循环真的重要到值得再次从头开始写吗😉。
- 如果你正在处理一个自定义的、高度复杂的用例,不符合通常训练循环的模式,
致谢
pytorch-accelerated 的设计和功能的许多方面都深受一些优秀库和框架的启发,比如 fastai、timm、
PyTorch-lightning 和 Hugging Face Accelerate。这些工具
对本库和机器学习社区都产生了巨大影响,其影响力不可言喻!
pytorch-accelerated 只是从这些工具中汲取灵感,所有包含的功能都是从零开始实现的,以适应本库的需求。唯一的例外是
examples
文件夹中的一些脚本,这些脚本采用并修改了现有资源以展示 pytorch-accelerated 的功能;
这些情况都有明确标记,并对原作者给予了致谢。
--- Tranlated By Open Ai Tx | Last indexed: 2026-02-28 ---