数字旗手

电气化、自动化、数字化、智能化、智慧化

0%

轻量级PyTorch通用训练模板pytorch-accelerated解析:0 -- 介绍及入门

介绍

pytorch-accelerated是一个轻量级的库,旨在通过提供一个最小但可扩展的训练循环(封装在一个Trainer对象中)来加速PyTorch模型的训练过程;它足够灵活,可以处理大多数用例,并且能够利用不同的硬件选项而不需要修改代码。更多介绍见原作者的博文
pytorch-accelerated最大的两个特点就是:简单simplicity和透明transparency。怎么理解这两个词呢(个人理解):
(1)简单体现在它是一套可复用的PyTorch的训练代码,每次有新模型和新场景需要训练时不必将原来的代码拷过来拷过去,pytorch-accelerated提供了一套通用但不失灵活性的代码模板;
(2)透明体现在它基于纯正的PyTorch语法和API,而不是对PyTorch进行过度封装(此处PyTorch Lightning中枪),否则用户会感觉新学习了一套框架和语法。
这里可以再扩展说一下,作者的介绍博文后有一条评论:“这个库与PyTorch Lighting对比怎么样,我已经用了Lightning,为啥要用这个库呢?”作者给出了一个长长的回答,解释得挺清晰的:

我使用Lightning很长时间了,但越来越觉得我花在理解和调试Lightning上的时间比花在任务上的时间多。这在一开始可能很简单,但由于选项和指标太多,我发现自己花了很多时间阅读文档,并发现源代码对我来说并不容易一目了然。
另外,从一开始我就发现,Lightning的设计,即模型和训练循环绑在一起,并不适合我想为同一个任务轻松切换不同的模型和优化器的场景;我个人更喜欢训练循环与模型分开。这意味着模型仍然是纯正的PyTorch;我可以插入其他库的模型,如timm,而不需要做任何修改,也不需要在部署时有任何额外的依赖。
从本质上讲,我想要一个简单易懂、易于调试的库,同时保持足够的灵活性以满足我的需要;学习曲线足够浅,我可以把它介绍给其他人,他们也可以很快上手。根据我自己的要求,我创建了PyTorch-accelerated。它在简单的情况下开箱即用,而且很容易定制行为的任何部分。

再再扩展说一下,对于PyTorch Lightning,知乎上有很多真香帖(比如12),也有很多劝退帖(比如12(这篇虽然从标题和全文上看是“真香”,实则在文章最后和评论中作者表示弃坑了,采用了原生的PyTorch))。个人感觉采用PyTorch Lightning有两个很大的弊端(严格来说其实是一个):
(1)与原生PyTorch相比,PyTorch Lightning进行了过度封装,感觉像是在学习另一个框架;
(2)假设使用PyTorch Lightning编写代码,如果是个人使用还好,但如果是一个团队共同维护代码,很难说服别人也采用该写法。
于是,本文的主角pytorch-accelerated就有了用武之地。

目标用户

什么类型的用户/开发者可以尝试使用pytorch-accelerated呢:
(1)熟悉PyTorch的用户,但希望避免编写常见的训练循环模板,以专注于其他更重要的部分。
(2)喜欢并乐于选择和创建自己的模型、损失函数、优化器和数据集的用户。
(3)喜欢简单但高效的功能的用户,这些功能的行为需要很容易调试、理解和推理(PyTorch Lightning:你报我身份证得了)。

另一方面,pytorch-accelerated不适合什么类型的用户/开发者呢?
(1)如果你正在寻找一个端到端的解决方案,包括从加载数据到推理的所有内容,该方案帮助你选择模型、优化器或损失函数,此时可能更适合使用fastaipytorch-accelerated只关注训练过程,其他所有问题都由用户负责。
(2)如果你想自己编写整个训练循环,但是不想涉及恼人的设备管理问题(比如多机多卡并行),你可能最适合使用Hugging Faceaccelerateaccelerate专注于将PyTorch的分布式训练和混合精度训练变得简单,不过整个训练循环得自己写。
(3)如果你正在研究一个定制的、高度复杂的、不符合通常训练循环模式的用例,并且想在你选择的硬件上挤出每一点性能,你可能最好坚持使用原生PyTorch;在高度专业化的情况下,任何高级API都会成为一种开销。

安装

(注意提前安装好pytorch
使用pip安装:

1
pip install pytorch-accelerated

其实看pytorch-acceleratedrequirements.txt,它只依赖于Hugging Faceaccelerate库及tqdm(在终端下显示进度条),因此可以说是它的API就是PyTorch的原生API加上accelerate的分布式训练的API

accelerate库

Hugging Face🤗accelerate库(地址),可以无痛地对Pytorch进行多GPUTPU混合精度训练。见机器之心的报道

多数 PyTorch 高级库都支持分布式训练和混合精度训练,但是它们引入的抽象化往往需要用户学习新的 API 来定制训练循环。许多 PyTorch 用户希望完全控制自己的训练循环,但不想编写和维护训练所需的复杂的样板代码。Hugging Face 最近发布的新库 Accelerate 解决了这个问题。
🤗accelerate提供了一个简单的API,将与多GPUTPUfp16相关的样板代码抽离了出来,保持其余代码不变。PyTorch用户无须使用不便控制和调整的抽象类或编写、维护样板代码,就可以直接上手多GPUTPU

🤗accelerate库有两大特点:
(1)提供了一套简单的API来处理分布式训练和混合精度训练,无需在不同情形下对代码进行大的改动;
(2)提供了一个命令行接口工具来快速配置并并行脚本。

可选

如果是想为了直接运行pytorch-accelerated提供的examples,则可以这样安装其他所依赖的包

1
pip install pytorch-accelerated[examples]

配置和运行

可以使用🤗accelerateCLI工具来生成配置文件:

1
accelerate config --config_file accelerate_config.yaml

然后运行:
1
accelerate launch --config_file accelerate_config.yaml train.py [--training-args]

值得注意的是,也不是必须要用CLI工具。如果想更精细地控制启动命令和参数,仍然可以通过通常的方式来运行脚本:
1
python train.py / python -m torch.distributed ...

MNIST例子

MNIST手写字符分类可以说是深度学习领域的Hello World
这一节将以MNIST来看看pytorch-accelerated是怎么使用的。

数据和模型准备

因为pytorch-accelerated专注于训练模型部分,因此数据的加载和模型的构建、配置都是使用的原生的PyTorch的代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
# examples/train_mnist.py
import os

from torch import nn, optim
from torch.utils.data import random_split
from torchvision import transforms
from torchvision.datasets import MNIST

# 加载数据集
dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())
# 拆分数据集
train_dataset, validation_dataset, test_dataset = random_split(dataset, [50000, 5000, 5000])

# 定义神经网络模型
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))

# 实例化模型
model = MNISTModel()
# 定义优化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# 定义损失函数
loss_func = nn.CrossEntropyLoss()

训练循环

加载了模型和数据后,下一步就是编写训练循环。
这里就是pytorch-accelerated的用武之地,唯一要做的就是导入Trainer

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
from pytorch_accelerated import Trainer

# Trainer被设计成用来封装一个完整的训练循环
# 将模型、损失函数和优化器都传入Trainer中
trainer = Trainer(
model,
loss_func=loss_func,
optimizer=optimizer,
)

# 入口点是train方法,定义了该怎样训练
# 可以设置训练集、验证集(注意是验证集,不是测试集)、迭代次数、批处理大小
# 还可以设置学习率策略、梯度累加等,这里没有用到
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,
)

训练

生成配置文件

1
accelerate config --config_file train_mnist.yaml

开始训练

1
accelerate launch --config_file train_mnist.yaml train_mnist.py

增加指标

上述例子中追踪的指标仅仅是每次迭代的损失loss,为了对训练结果有更深入的洞察,可以增加更多的指标。
对于指标的计算,这里引入了torchmetrics(发现这个库还是PyTorch-Lightning社区开发的,这波PyTorch-Lightning被人摘桃了,血亏🤗),该库兼容分布式训练,因此就不需要手动从不同进程中聚合计算结果。
计算指标有两种方式:
(1)定义一个继承自Trainer的子类,
(2)使用回调callback
具体使用哪种方式极大地取决于用户的喜好。
不过作者有如下建议:因为计算指标实际上是不能影响训练代码的,因此使用callback可能是一个好的方式,因为使用Trainer子类的话,它也会间接地参与训练过程。不过还是具体情况具体分析。注意,因为callbacks都是顺序执行的,必须保证在打印指标之前就调用这些回调。
下面是对比这两种实现方式。

使用Trainer子类

Trainer有很多方法可以被重载,具体的文档在这里。主要的一个特点是有动词前缀(比如createcalculate)的方法都是期望能返回一个数值,其他的方法(比如optimizer.step())则是用来设置内部状态。
示例代码为:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
# Copyright © 2021 Chris Hughes
########################################################################
# This example trains a model on the MNIST Dataset

# This example demonstrates how the default trainer class can be overridden
# so that we can record classification metrics
#
# Note, this example requires installing the torchmetrics package
########################################################################

import os

from torch import nn, optim
from torch.utils.data import random_split
from torchmetrics import MetricCollection, Accuracy, Precision, Recall
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, x):
return self.main(x.view(x.shape[0], -1))


class TrainerWithMetrics(Trainer):
def __init__(self, num_classes, *args, **kwargs):
super().__init__(*args, **kwargs)

# this will be moved to the correct device automatically by the
# MoveModulesToDeviceCallback callback, which is used by default
self.metrics = MetricCollection(
{
"accuracy": Accuracy(num_classes=num_classes),
"precision": Precision(num_classes=num_classes),
"recall": Recall(num_classes=num_classes),
}
)

def calculate_eval_batch_loss(self, batch):
batch_output = super().calculate_eval_batch_loss(batch)
preds = batch_output["model_outputs"].argmax(dim=-1)

self.metrics.update(preds, batch[1])

return batch_output

def eval_epoch_end(self):
metrics = self.metrics.compute()
self.run_history.update_metric("accuracy", metrics["accuracy"].cpu())
self.run_history.update_metric("precision", metrics["precision"].cpu())
self.run_history.update_metric("recall", metrics["recall"].cpu())

self.metrics.reset()


def main():
dataset = MNIST(os.getcwd(), download=True, transform=transforms.ToTensor())
num_classes = len(dataset.class_to_idx)

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 = TrainerWithMetrics(
model=model, loss_func=loss_func, optimizer=optimizer, num_classes=num_classes
)

trainer.train(
train_dataset=train_dataset,
eval_dataset=validation_dataset,
num_epochs=2,
per_device_batch_size=32,
)

trainer.evaluate(
dataset=test_dataset,
per_device_batch_size=64,
)


if __name__ == "__main__":
main()

使用回调

上面的代码可能已经有所显示,对于“增加指标”这种小的微调,如果使用Trainer子类的话会显得有些用力过猛。
此例中,可以保持默认的trainer不变,而使用回调来扩展功能。
为了创建一个新的回调,需要创建TrainerCallback的子类,然后重载相关方法,文档见这里。为了避免与Trainer的方法混淆,所有的回调方法都有on_前缀。
创建新的回调:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
from torchmetrics import MetricCollection, Accuracy, Precision, Recall

from pytorch_accelerated.callbacks import TrainerCallback


class ClassificationMetricsCallback(TrainerCallback):
def __init__(self, num_classes):
self.metrics = MetricCollection(
{
"accuracy": Accuracy(num_classes=num_classes),
"precision": Precision(num_classes=num_classes),
"recall": Recall(num_classes=num_classes),
}
)

def _move_to_device(self, trainer):
self.metrics.to(trainer.device)

def on_training_run_start(self, trainer, **kwargs):
self._move_to_device(trainer)

def on_evaluation_run_start(self, trainer, **kwargs):
self._move_to_device(trainer)

def on_eval_step_end(self, trainer, batch, batch_output, **kwargs):
preds = batch_output["model_outputs"].argmax(dim=-1)
self.metrics.update(preds, batch[1])

def on_eval_epoch_end(self, trainer, **kwargs):
metrics = self.metrics.compute()
trainer.run_history.update_metric("accuracy", metrics["accuracy"].cpu())
trainer.run_history.update_metric("precision", metrics["precision"].cpu())
trainer.run_history.update_metric("recall", metrics["recall"].cpu())

self.metrics.reset()

需要注意的一点是,在训练或验证之前,需要手动地将指标移动到正确地设备上。不过Trainer已经将这一步做了非常简单的处理,可以依据上下文返回正确的设备,即:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from pytorch_accelerated.trainer import DEFAULT_CALLBACKS

# 将自定义的回调传入Trainer
# 因为想保持默认行为,所以将自定义的回调放在所有的默认回调之前。
trainer = Trainer(
model,
loss_func=loss_func,
optimizer=optimizer,
callbacks=(
ClassificationMetricsCallback(
num_classes=num_classes,
),
*DEFAULT_CALLBACKS,
),
)

然后将上面代码复制进入最开始的代码中即可,无需变动其他地方的代码(只是需要在创建数据集后,计算一下其分类数目以传入回调num_classes = len(dataset.class_to_idx))。

后记

实际只从这个MNIST例子还不能看出pytorch-accelerated能在多大程度上提高效率,后面需要多研究一下它的例子,毕竟通用模板的意义在于能适用于多种情形。留坑待填。