Amazon Web Services ブログ

Amazon SageMaker Debugger と Amazon SageMaker Experiments による機械学習モデルのプルーニング

過去 10 年間、ディープラーニングは、コンピュータービジョンや自然言語処理など、さまざまな分野を進歩させてきました。最先端のモデルは、画像分類などのタスクで人間に近いパフォーマンスを実現しています。ディープニューラルネットワークでこれが実現できるのは、大規模なトレーニングデータセットでトレーニングする数百万のパラメータで構成されているためです。たとえば、BERT (ラージ) モデルは 3 億 4,000 万のパラメータで構成され、Resnet-152 は 6,000 万のパラメータで構成されています。このようなモデルを最初からトレーニングすることは、計算量が多く、数時間、数日、はたまた数週間かかることもあります。

通常、データサイエンティストは転移学習を実行します。これは、ある問題を解くことによって得られた知識を、関連するが異なる問題に応用するプロセスです。転移学習では、より小さなデータセットで事前学習済みモデルを微調整して、精度を向上させられます。このようなシナリオでは、モデルに多数のパラメータは必要ない場合があります。小さなモデルでも同様に機能する場合もあるかもしれません。

エッジでの機械学習 (ML) のコンテキストでは、小さなモデルを持つことが不可欠です。ハードウェアの制約により、レイテンシー、メモリフットプリント、コンピューティング時間などの要素は、モデルの精度と同様に重要です。たとえば、自動運転には高精度で低レイテンシーのモデルが必要です。このようなシナリオでは、精度は 1% 向上するが、予測に 2 倍の時間がかかるモデルは好ましくありません。

モデルプルーニングは、精度を犠牲にすることなくモデルサイズを大幅に削減できます。考え方は簡単です。トレーニングプロセスにほとんど寄与しないモデル内の冗長パラメータを特定するのです。

この記事は、Amazon SageMaker を使用した反復モデルのプルーニング (枝刈り) を示します。この記事では、事前トレーニング済みのモデルを使用したサンプルアプリケーションについて説明します。そのアプリケーションは、精度を大幅に損なうことなく、繰り返しプルーニングして 3 分の 1 以上削減します。

モデルプルーニング

モデルプルーニングは、トレーニングプロセスにあまり貢献しない重みを削減することを目指しています。重みは学習可能なパラメータです。これは、トレーニングプロセス中にランダムに初期化され、最適化されます。転送パスの間、データはモデルを通過します。損失関数は、ラベルを指定してモデル出力を評価します。逆方向パス中は、重みを更新して損失を最小限に抑えます。そうするために、重みに関する損失の勾配を計算し、それぞれの重みは異なる更新を受け取ります。数回反復した後、特定の重みは通常、他の重みよりも影響力があります。プルーニングの目的は、モデルの精度を大幅に低下させることなく、不要なものを削除することにあります。次の図は、このワークフローを示しています。

次のヒューリスティックを使用して、重みの重要性を測定できます。

  • 重みの大きさ – 絶対値がしきい値より小さい場合は、重みを削除します。重みが小さいほど、出力への影響が少なくなります。
  • 平均活性化 – ニューロンがトレーニング全体を通じてほとんど活性化されていない場合、活性化関数に加わる重みの関連性が低いと推測できます。

非構造化と構造化の重みプルーニングは、次のように区別できます。

  • 非構造化プルーニングは任意の重みを削除します (前の図のように)
  • 構造化されたプルーニングは、畳み込みフィルターと関連するチャネル全体を削除します

構造化されたプルーニングは、多くの畳み込み層で通常構成されるコンピュータビジョンモデルに特に関わっています。フィルターはカーネルのコレクションです (1 つの入力チャネルごとに 1 つのカーネル)。フィルターは、出力チャネルとも呼ばれる 1 つの機能マップを生成します。次の図は、3 つの出力機能マップを生成する 3 つのカーネルを示しています。モデルが学習する必要があるパラメータの数 (その重み) は、3 x input_channels x kernel_width x kernel_height であり、この例の入力チャネルの数は 1 です。簡単にするために、この図ではバイアステンソルがないと想定してもかまいません。フィルターをランク付けして、最も重要度の低いものを識別できます (たとえば、黄色のフィルター)。パラメータを削除して、パラメータの数を 1 x input_channels x kernel_width x kernel_height 分だけ減らします。

各フィルターの重要度をランク付けするには、ホワイトペーパー 「リソース効率の高い推論のために畳み込みニューラルネットワークをプルーニングする」 に記載されているようなランク付け方法を使用します。これらは、プルーニングフィルターの損失への影響を推定します。目標は、損失に影響を与えないものを削除することです。フィルターは、その活性化出力と対応する勾配が小さい場合、低いランクを受け取ります。トレーニング全体を通じて活性化出力と勾配の積を累積することにより、フィルターの重要度の推定値を取得できます。

次に、最低ランクのフィルターを削除し、モデルを微調整して、プルーニングから回復し、精度を取り戻します。これらの手順を数回繰り返してもいいでしょう。

Amazon SageMaker での反復モデルプルーニング

Amazon SageMaker Debugger では、トレーニングジョブからテンソルを放出し、トレーニングの問題を自動的に検出する組み込みルールを実行できます。勾配と活性化出力を取得して、フィルターランクを計算できます。詳細については、「Amazon SageMaker Debugger – Machine Learning モデルのデバッグ」を参照してください。Amazon SageMaker Experiments により、ML 実験を大規模にカスタマイズ、視覚化、追跡できます。詳細については、「Amazon SageMaker Experiments – 機械学習トレーニングを整理、追跡、比較する」を参照してください。この記事では、さまざまなプルーニングの反復を追跡するために Amazon SageMaker Experiments を使用します。Amazon SageMaker Studio の [Experiments ] ビューで、最高の精度とサイズのトレードオフをもたらすモデルをすばやく特定してデプロイできます。

反復モデルのプルーニングを段階的に説明するために、この記事では、ImageNet で事前トレーニングされた Resnet18 モデルを使用し、101 クラスのみで構成される Caltech101 データセットで微調整しています。再トレーニング中に画像のサイズを変更します。ResNet は、複数の残差ブロックで構成される畳み込みニューラルネットワークです。ブロックは畳み込み層で構成され、その後にバッチ正規化層と ReLu 関数が続きます。スキップ接続により、入力はブロックをバイパスできます。この記事では、約 1,100 万のパラメータを持つ ResNet モデルの最小バージョンである ResNet18 を使用します。プルーニングを繰り返すたびに、200 個の最低ランクのフィルターを削除します。

次の図は、ソリューションのワークフローを示しています。

手順は次のとおりです。

  1. トレーニングジョブを開始する
  2. 重み、勾配、バイアス、活性化出力を取得する
  3. フィルターランクを計算する
  4. 低ランクフィルターをプルーニングする
  5. 新しい重みを設定する
  6. プルーニング済みモデルでトレーニングジョブを開始する

実験とデバッガーのフック設定を作成する

ソリューションを実装する前に、モデルプルーニング実験を作成します。実験とは、試行の集まりで、試行とはトレーニングステップの集まりです。次のコードを参照してください。

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 トレーニングジョブを開始します。これは、実験内の新しい試行です。次のコードを参照してください。

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 はトレーニングジョブを実験と試行に関連付けることができます。次のコードを参照してください。

experiment_config = { "ExperimentName": "model_pruning_experiment",
                        "TrialName":  trial.trial_name,
                        "TrialComponentDisplayName": "Training"}

トレーニングジョブを開始する前に、デバッガーフック設定を定義する必要があります。Amazon SageMaker Debugger は、重み、バイアス、勾配、および損失のデフォルトのコレクションを提供します。この記事では、活性化出力も保存する必要があります。それを取得するには、正規表現がテンソル名に含めるように示すカスタムコレクションを作成します。ResNet はバッチノルムレイヤーで構成されているため、平均と分散を実行するバッチノルム統計も保存します。テンソルは 100 ステップごとに保存され、ステップは 1 つの順方向パスと逆方向パスを表します。次のコードを参照してください。

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 リポジトリをご覧ください。テンソルを出力するには、デバッガーのフック設定を 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
        )

estimator オブジェクトを定義したら、fit を呼び出すことができます。これにより、マネージド ml.p3.2xlarge インスタンスが起動し、トレーニングスクリプトが実行されます。前述のとおり、experiment_config をトレーニングジョブに渡します。次のコードを参照してください。

estimator.fit(experiment_config=experiment_config)

勾配、重み、バイアスを取得する

トレーニングジョブが完了したら、勾配、重み、バイアスなどのテンソルを取得します。テンソルを読み取り、フィルタリングする関数を提供する smdebug ライブラリを使用できます。まず、デバッガーが保存したテンソルへのアクセスを許可する trial を作成します。詳細については、GitHub リポジトリをご覧ください。Amazon SageMaker Debugger のコンテキストでは、試行は特定のトレーニングジョブのテンソルをクエリできるオブジェクトです。Amazon SageMaker Experiments のコンテキストでは、試行は実験の一部であり、単一のトレーニングジョブに含まれるトレーニングステップのコレクションを提示します。次のコードを参照してください。

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))

その後、フィルターのランクを正規化し、サイズで並べ替えます。

低ランクフィルターのプルーニング

これで、最小のフィルターを取得できます。次のコードは、ランクが 0.0 であるため、layer1.0 のフィルター 1、36、および 127 をプルーニングできることを示しています。

[('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 デバッガーを使用して、layer1.0 の最初の畳み込み層の重みテンソルを取得し、2 番目の次元 (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 の出力チャネルがあります。3 つのフィルターを削除すると、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 で反復モデルプルーニング実験を追跡して視覚化できます。トレーニングスクリプトは、SageMaker Debugger の save_scalar メソッドを使用して、モデルパラメータの数とモデルの精度を格納します。詳細については、GitHub リポジトリをご覧ください。save_scalar が書き込む値は、Amazon SageMaker Studio が視覚化データを作成するために使用するデータストアに送られます。次のグラフは散布図を示しています。x 軸はモデルパラメータの数を示し、y 軸は検証の精度を示しています。

当初、モデルは 1,100 万個のパラメータで構成されていました。11 回の反復後、パラメータの数は 706,000 に減少し、精度は 90% に増加し、8 回のプルーニング反復後に低下し始めました。次のスクリーンショットは、さまざまな試行と詳細がリストされている実験ビューを示しています。

カスタムルールで反復モデルプルーニングを実行する

前の例では、モデルのパラメータが 400 万個未満の場合、精度が低下します。この時点に達したら実験を停止します。SageMaker Debugger には、モデルのトレーニングに、たとえば勾配の消失や損失の不減少などの問題が発生したときにトリガーされる組み込みのルールが用意されています。モデルに十分な容量がない (パラメータが少なすぎる) 場合、うまく学習できません。含意の 1 つは、損失が減少しない可能性があることです。詳細については、「Amazon SageMaker Debugger が用意する組み込みルール」を参照してください。

この記事では、現在のモデルの精度と以前のトレーニングジョブの精度を比較するカスタムルールを定義できます。たとえば、10% 以上減少すると、ルールがトリガーされ、True を返します。次に、Amazon CloudWatch アラームと AWS Lambda 関数をセットアップして、トレーニングジョブを停止し、モデルプルーニング実験が低品質モデルを生成するジョブのリソースを浪費するのを防ぐことができます。詳細については、GitHub リポジトリをご覧ください。

次のコードは、カスタムルールの概要を示しています。

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 リポジトリを参照してください。

トレーニングジョブでルールを実行するには、ルールを 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 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

次のグラフは、パラメータの数と精度を示しています。前の実験とは対照的に、低品質のモデルが生成されるとトレーニングジョブが停止し、実験が終了します。その結果、8 回の反復しか実行されません。

次のスクリーンショットは、SageMaker Studio の Debugger ビューを示しています。カスタムルールが問題を検出したことを示しています。

SageMaker Debugger を使用したカスタムルールの定義と実行の詳細については、「カスタムルールの使用方法」を参照してください。

まとめ

この記事では、Amazon SageMaker を使った反復モデルプルーニングと、トレーニングプロセスにほとんど寄与しない冗長パラメータを特定することで、モデルのサイズを大幅に削減し、精度を維持する方法について説明しました。事前トレーニング済みモデルを使用するサンプルアプリケーションを説明し、繰り返しプルーニングしてもモデルの精度が失われないことを確認しました。

著者について

Nathalie Rauschmayr は、AWS のアプライドサイエンティストで、お客様によるディープラーニングアプリケーションの開発を支援しています。

 

 

 

Julien Simon は EMEA の人工知能と機械学習のエバンジェリストで、開発者や企業がアイデアを実現するのを支援することに力を入れています。

 

 

 

Satadal Bhattacharjee は、AWS AI のプリンシパルプロダクトマネージャーです。彼は SageMaker などのプロジェクトで機械学習エンジン PM チームを率い、TensorFlow、PyTorch、MXNet などの機械学習フレームワークを最適化しています。