发布于: Nov 30, 2022
【概要】在本文中，我们将使用 SageMaker Experiments 跟踪不同的迭代剪枝。您也可以使用 Amazon SageMaker Studio 中的 “Experiments” 视图快速识别并部署在准确性与大小之间取得最佳平衡点的模型。
模型剪枝是在不牺牲深度学习效果准确性的前提下，对学习模型进行“瘦身”。在 Amazon SageMaker Debugger 中，我们可以通过训练任务发出张量，同时使用内置规则自动检测训练过程中出现的问题。另外，大家也可以检索梯度与激活输出，并据此计算过滤器排名。Amazon SageMaker Experiments 则允许您以规模化方式定义、可视化并跟踪机器学习实验。在本文中，我们将使用 SageMaker Experiments 跟踪不同的迭代剪枝。您也可以使用 Amazon SageMaker Studio 中的 “Experiments” 视图快速识别并部署在准确性与大小之间取得最佳平衡点的模型。
为了逐步进行迭代模型剪枝，本文选择使用在 ImageNet 上预训练完毕的 Resnet18 模型，并利用仅包含 101 个类的 Caltech101 数据集对模型进行微调。您也可以在重新训练期间反复调整图像数据集的大小。ResNet 是由多个残差组成的卷积神经网络。其中每个残差块，具体包括卷积层、批处理归一化层以及 ReLu 函数。跳过连接，则意味着允许输入绕开当前块。在本文中，我们将使用 ResNet18，即包含约 1100 万个参数的最精简 ResNet 模型版本。在每一轮剪枝迭代中，我们将删除排名最低的 200 个过滤器。
如下图所示，为这套解决方案的实际工作流程：
其具体操作步骤包括：
- 启动训练作业。
- 获取权重、梯度、偏差与激活输出。
- 计算过滤器排名。
- 对排名最低的过滤器进行剪枝。
- 设置新的权重。
- 使用经过剪枝的模型启动新一轮训练任务。
在实施这套解决方案之前，我们首先需要创建一个模型剪枝 experiment。这里的 experiment 代表一组试验的集合，而每项 trial 又是多个训练步骤的集合。具体请参见以下代码：
from smexperiments.experiment import Experiment
Experiment.create(experiment_name="model_pruning_experiment",
description="Iterative model pruning of ResNet18 trained on Caltech101",
sagemaker_boto_client=sagemaker_boto_client)
在每一轮剪枝迭代中，我们都会启动新的 Amazon SageMaker 训练任务，该任务是 experiment 中新的 trail。具体参见以下代码：
from smexperiments.trial import Trial trial = Trial.create(experiment_name="model_pruning_experiment", sagemaker_boto_client=sagemaker_boto_client)
接下来，需要定义 experiment_config，即传递给 Amazon SageMaker 训练任务的字典。以此为基础，Amazon SageMaker 会将训练作业与 experiement 及 trail 关联起来。具体请参见以下代码：
experiment_config = { "ExperimentName": "model_pruning_experiment",
"TrialName": trial.trial_name,
"TrialComponentDisplayName": "Training"}
在开始训练之前，需要先定义 debugger hook 配置。Amazon SageMaker Debugger 为权重、偏差、梯度以及损失提供多套默认集合。在本文中，我们还需要保存激活输出。要执行检索，应创建一个自定义集合，其中的正则表达式负责指示需要包含的张量名称。由于 ResNet 由批处理归一化层组成，所以我们也可以存储运行均值和方差的批量归一化统计信息。张量每 100 步保存一次，其中每一步代表一次前进与后退传递。具体请参见以下代码：
from sagemaker.debugger import DebuggerHookConfig, CollectionConfig
debugger_hook_config = DebuggerHookConfig(
collection_configs=[
CollectionConfig(
name="custom_collection",
parameters={ "include_regex": ".*relu|.*weight|.*bias|.*running_mean|.*running_var",
"save_interval": "100" })])
现在，我们已经做好利用 Amazon SageMaker 训练 ResNet18 模型的全部准备工作了。训练循环将在 entry_point 文件的 train.py 中进行定义。要了解更多详细信息，请参考 GitHub repo。要发出张量，请将 debugger hook 配置传递至 PyTorch Estimator 当中，具体参见以下代码：
import sagemaker
from sagemaker.pytorch import PyTorch
estimator = PyTorch(role=sagemaker.get_execution_role(),
train_instance_count=1,
train_instance_type='ml.p3.2xlarge',
train_volume_size=400,
source_dir='src',
entry_point='train.py',
framework_version='1.3.1',
py_version='py3',
debugger_hook_config=debugger_hook_config
)
在定义估计器对象之后，即可调用 fit。这步操作将启动 ml.p3.2xlarge 托管实例以运行您的训练脚本。如前文所述，您需要将 experiment_config 传递至训练作业，详见以下代码：
estimator.fit(experiment_config=experiment_config)
在训练任务完成之后，您可以获取其张量，例如梯度、权重与偏差。在这里，大家可以使用 smdebug 库，通过相关函数读取并过滤各张量。首先，创建一个 trial 以访问 debugger 所保存的张量。关于更多详细信息，请参见 GitHub repo。在 Amazon SageMaker Debugger 使用场景下，trail 代表的是一个对象，用户可以借此查询特定训练任务中的张量。在 Amazon SageMaker Experiment 使用场景下，trial 属于 experiment 的一部分，负责表示与单一训练任务相关的训练步骤集合。详见以下代码示例：
from smdebug.trials import create_trial path = estimator.latest_job_debugger_artifacts_path() smdebug_trial = create_trial(path)
要访问张量值，请调用 smdebug_trial.tensor()。例如，要想获取第一个卷积层的激活输出，请使用以下代码：
smdebug_trial.tensor('layer1.0.relu_0_output_0').value(0, mode=modes.TRAIN)
现在您可以访问张量，并计算它们的过滤器排名了。首先迭代当前训练步骤，而后检索激活输出及其梯度。例如，以下代码片段会计算模型中第一个特征层的过滤器排名，而后为每个过滤器计算出对应的排名值：
rank = 0
for step in smdebug_trial.steps(mode=modes.TRAIN):
activation_output = smdebug_trial.tensor( 'layer1.0.relu_0_output_0').value(step, mode=modes.TRAIN)
gradient = smdebug_trial.tensor( 'gradient/layer1.0.relu_ReLU_output').value(step, mode=modes.TRAIN)
product = activation_output * gradient
rank += np.mean(product, axis=(0,2,3))
接下来，我们可以归一化过滤器排名并按大小对结果进行排序。
现在，您可以检索排名最低的过滤器。以下代码显示，您可以对 layer1.0 中的 1、36 与 127 号过滤器进行剪枝，因为它们的排名值为 0.0：
[('layer1.0.relu_ReLU_output', 1, 0.0),
('layer1.0.relu_ReLU_output', 127, 0.0),
('layer1.0.relu_ReLU_output', 36, 0.0)]
要对过滤器执行剪枝，请使用 SageMaker Debugger 获取 layer1.0 中第一个卷积层的权重张量，并删除第二个维度 (axis=1) 中的上述条目。具体参见以下代码：
weight = trial.tensor('ResNet_layer1.0.conv1.weight').value(step, mode=modes.TRAIN)
weight = np.delete(weight, [1,36,127], axis=1)
您还需要调整卷积参数。Resnet18 模型内 layer1.0 当中的卷积层拥有 64 个输出通道。在移除以上三个过滤器后，它只余下 61 个输出通道。您还需要调整后续批量归一化层的权重。为此，我们删除第一维 (axis=0) 中的条目，具体参见以下代码：
weight = trial.tensor('ResNet_layer1.0.bn1.weight').value(step, mode=modes.TRAIN)
bias = trial.tensor('ResNet_layer1.0.bn1.bias').value(step, mode=modes.TRAIN)
weight = np.delete(weight, [1,36,127], axis=0)
bias = np.delete(bias, [1,36,127], axis=0)
在对 200 个最小过滤器进行剪枝之后，使用最新权重保存新的模型定义，而后开始下轮剪枝迭代。您可以重复完成这些步骤，每一轮迭代都将削减模型的实际大小。
您可以在 Amazon SageMaker Studio 当中跟踪并可视化各迭代模型剪枝 experiment。训练脚本会利用 SageMaker Debugger 中的 save_scalar 方法保存模型参数数量及模型准确性。关于更多详细信息，请参见 GitHub repo。其中 save_scalar 的值会被写入至数据存储内，供 Amazon SageMaker Studio 用于创建可视化结果。在以下散点图中，x 轴代表模型参数的数量，而 y 轴则代表模型的验证准确性。
最初，这套模型中包含 1100 万个参数。在经过 11 轮迭代之后，参数数量减少至 706000，准确性提高到 90%，但准确性从第 8 次剪枝迭代后开始下降。以下截图所示为 experiment 视图，其中列出了每一轮 trail 的详细信息。
在之前的示例中，当模型的参数少于 400 万时，准确性开始降低。在达到这一阈值后，我们希望停止 experiment。SageMaker Debugger 中提供内置规则，可以在模型训练发生问题（例如梯度消失或者损失不再降低）时被触发。如果模型的容量不足（参数太少），则说明学习效果不好。其代表的可能性之一就是损失无法进一步降低。关于更多详细信息，请参阅《Amazon SageMaker Deubgger 中的内置规则（Built-in Rules Provided by Amazon SageMaker Debugger）》。
在本文中，我们可以定义一项自定义规则，利用它将当前模型的准确性与前一训练任务的模型准确性进行比较。例如，如果准确性下降超过 10%，则触发规则并返回 True。接下来，您可以设置 Amazon CloudWatch 警报及 Amazon Lambda 函数以停止训练作业，并防止将基础设施资源浪费在对低质量模型的剪枝实验当中。关于更多详细信息，请参见 GitHub repo。
以下代码所示，为一条自定义规则的高级概述：
from smdebug.rules.rule import Rule
class check_accuracy(Rule):
def __init__(self, base_trial,
previous_accuracy=0.0):
self.previous_accuracy = float(previous_accuracy)
def invoke_at_step(self, step):
predictions = np.argmax(self.base_trial.tensor('CrossEntropyLoss_0_input_0').value(step, mode=modes.EVAL), axis=1)
labels = self.base_trial.tensor('CrossEntropyLoss_0_input_1').value(step, mode=modes.EVAL)
current_accuracy = compute_accurcay(predictions, labels)
if self.previous_accuracy - current_accuracy > 0.10:
return True
return False
这条规则用于实现一个继承自 smdebug 规则类的 Python 类。它以前一训练任务的模型准确性为参数，在完成每一步新操作并产生新张量时调用 invoke_at_step 函数。利用 smdebug，您可以将损失以模型预测与标签的形式输入到损失函数当中。接下来，您就可以计算当前模型准确性并将其与训练前的水平进行比较。关于规则机制的完整实现方法与细节信息，请参见 GitHub repo。
要在训练作业中运行该规则，大家需要将其传递至 PyTorch 估计器对象。在此之前，我们还要先创建一个自定义规则配置，详见以下代码示例：
from sagemaker.debugger import Rule, CollectionConfig, rule_configs
check_accuracy_rule = Rule.custom(
name='CheckAccuracy',
image_uri='840043622174.dkr.ecr.us-east-2.amazonaws.com/sagemaker-debugger-rule-evaluator:latest',
instance_type='ml.c4.xlarge',
volume_size_in_gb=400,
source='custom_rule/check_accuracy.py',
rule_to_invoke='check_accuracy',
rule_parameters={"previous_accuracy": "0.0"},
)
这条规则配置负责指定规则定义的位置、对应的输入参数以及规则容器将要运行的实例类型。您还需要为规则容器指定镜像。关于更多详细信息，请参阅《Amazon SageMaker 自定义规则估计器注册表 ID（Amazon SageMaker Custom Rule Evaluator Registry Ids）》。
在规则定义完成后，将参数 rules = [check_accuracy_rule] 传递至该 Pytorch 估计器。
在每一轮剪枝迭代中，我们都需要将上一轮训练作业中的准确性传递至规则当中，并通过 SageMaker Experiments 与 ExperimentAnalytics 模块进行检索。详见以下代码示例：
from sagemaker.analytics import ExperimentAnalytics trial_component_analytics = ExperimentAnalytics(experiment_name="model_pruning_experiment") accuracy = trial_component_analytics.dataframe()['scalar/accuracy_EVAL - Max'][0]
使用以下代码覆盖规则配置中的值：
check_accuracy_rule.rule_parameters["previous_accuracy"] = str(accuracy)
在每一轮迭代中，检查作业状态。如果上一代作业停止，则退出循环。详见以下代码：
job_name = estimator.latest_training_job.name
client = estimator.sagemaker_session.sagemaker_client
description = client.describe_training_job(TrainingJobName=job_name)
if description['TrainingJobStatus'] == 'Stopped':
break
下图所示，为参数数量与模型准确性之间的关系。与之前的实验不同，训练作业会在生成低质量模型时自动停止，而后 experiment 结束。结果就是，模型剪枝只运行了 8 轮迭代。
以下截图所示，为 SageMaker Studio 中的 “Debugger” 视图，可以看到自定义规则发现了一个问题。
关于使用 SageMaker Debugger 定义并运行自定义规则的更多详细信息，请参阅《如何使用自定义规则（How to Use Custom Rules）》。
本文探讨了使用 Amazon SageMaker 进行迭代模型剪枝的方法，同时介绍了如何通过识别对训练过程鲜有帮助的冗余参数来显著降低模型大小并保持模型准确性。我们还在本文中引入了使用预训练模型的应用示例，可以看到该模型通过迭代剪枝成功实现了准确性保障前提下的“瘦身”任务。
相关文章