WWDC Quick Look 💓 By SwiftGGTeam
Accelerate machine learning with Metal

Accelerate machine learning with Metal

元の動画を見る

ハイライト

Apple は、Metal Performance Shaders バックエンドを PyTorch 1.12 に追加し、TensorFlow Metal と MPS Graph を拡張し、Mac 上での機械学習トレーニング、カスタム オペレーター、計算グラフの同期で Apple GPU を使用できるようにしました。

主要内容

機械学習トレーニングは最も多くのコンピューティング能力を消費します。モデルは前方更新、後方更新、勾配更新を繰り返し実行する必要があり、CPU がすぐにボトルネックになってしまいます。 GPU は並列作業に適していますが、開発者はトレーニング コードを GPU に接続するためのフレームワークも必要です。

Metal の機械学習のエントリ ポイントは Metal Performance Shaders (MPS) です。 MPS は高性能 GPU プリミティブを提供し、MPS Graph は汎用計算グラフとその上の多次元テンソル サポートを提供します。 Core ML、TensorFlow、PyTorch などの上位層フレームワークは、MPS Graph、MPS、および Metal コマンド キューにオペレーターを配置できます。

このセッションは 3 つの段落でわかりやすく説明されています。PyTorch には公式の MPS バックエンドが追加されました。 TensorFlow Metal は、より大きなバッチ、カスタム操作、分散トレーニングをサポートします。 MPS Graph は、共有イベント、RNN/LSTM/GRU、最大プーリング インデックス、Philox ランダムおよびテンソル変換操作を追加します。

PyTorch: 3 つのステップで MPS デバイスに切り替える

(03:32) PyTorch コミュニティから最も頻繁に要求される機能の 1 つは、Apple シリコンでの GPU アクセラレーションの使用です。 Apple は MPS バックエンドを公式 PyTorch GitHub リポジトリにマージし、PyTorch 1.12 に入りました。これには、PyTorch オペレーション カーネルとランタイム フレームワークが含まれています。オペレータは MPS Graph/MPS を呼び出し、ランタイムは Metal のコマンド キュー、コマンド バッファ、および同期プリミティブを使用します。

開発者はトレーニング フレームワークを変更する必要はありません。 PyTorch をインストールし、MPS デバイスを作成し、モデルと入力をこのデバイスに移動して、後続の操作を GPU で実行できるようにします。

TensorFlow: GPU タイムラインで欠落した演算子を維持する

(08:01) TensorFlow Metal は、TensorFlow 2.5 以降のプラグインを介して GPU アクセラレーションを提供します。 2022 年のアップデートには、より大きなバッチ サイズ、新しい GPU アクセラレーション操作、カスタム操作、RNN の改善、分散トレーニングが含まれます。

この問題は、TensorFlow API でまだサポートされていないカスタム損失または演算子で発生します。この作業が CPU タイムラインに戻ると、同期オーバーヘッドが発生し、GPU が待機することになります。 TensorFlow Metal Stream プロトコルは、コマンド バッファー、ディスパッチ キュー、commit および commitAndWait をカスタム ops に公開し、開発者がカスタム GPU カーネルを同じ Metal ストリームにエンコードできるようにします。

MPS Graph: マルチキュー同期と追加のトレーニング オペレーター

(18:13) 一部のアプリケーションは、より多くの GPU 並列処理を得るために、計算、後処理、表示を異なるコマンド キューに分割します。これを行うとデータ競合が発生し、計算によって結果が生成される前に後処理が開始される可能性があります。 MPS Graph の共有イベント API は、シグナル/待機を使用してキュー間の依存関係を確立します。

(20:43) MPS Graph には、トレーニングと推論のための一般的な操作のバッチも追加されています。RNN、LSTM、GRU。リターンインデックスを使用した最大プーリング。フィロックスはランダム。ハミング距離。 ExpandDims、squeeze、split、stack、coordineAlongAxis 等テンソル変換。

詳細

PyTorch MPS バックエンドの最小パス

(03:44) PyTorch 1.12 以降、基本パッケージは pip 経由で直接インストールできます。

python -m pip install torch

キーポイント:

  • python -m pip別のインタープリター環境へのインストールを回避するには、現在の Python 環境から pip を使用します。 -install torch公式 PyTorch パッケージをインストールします。 MPS バックエンドは PyTorch 1.12 の一部です。

(03:59) デバイスを作成するときは、まず MPS が利用可能かどうかを確認してから、CPU にフォールバックします。

import torch

mpsDevice = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

キーポイント:

  • import torchPyTorch の紹介。 -torch.backends.mps.is_available()現在の環境で MPS バックエンドを使用できるかどうかを確認します。 -torch.device("mps" ...)MPS デバイスを作成します。利用できない場合に使用されるcpu、同じコード部分は引き続き実行できます。

(04:15) モデルはデフォルトで CPU 上で実行されます。中間テンソルが MPS を使用できるようにするには、モデルを MPS デバイスに移動する必要があります。

import torchvision

model = torchvision.models.resnet50()

model_mps = model.to(device=mpsDevice)

キーポイント:

  • torchvision.models.resnet50()ResNet50 モデルを作成します。 -model.to(device=mpsDevice)モデルを MPS デバイスに転送します。 -model_mps後続の実行では、モデル内で生成された中間テンソルも高速化された MPS バックエンドを使用します。

(04:46) 入力テンソルも同じデバイスに配置する必要があります。デフォルトでは、tensor は CPU に割り当てられます。デバイスの指定を忘れた場合、モデルと入力は別の場所に配置されます。

sample_input = torch.randn((32, 3, 254, 254), device=mpsDevice)

prediction = model_mps(sample_input)

キーポイント:

  • torch.randn(...)ランダム入力のバッチを作成します。 -device=mpsDevice入力テンソルを MPS デバイスに直接割り当てます。 -model_mps(sample_input)推論を実行すると、その後のテンソル演算が GPU 上で高速化されます。

(05:15) Apple は、StyleTransfer ネットワークを使用してトレーニング効果を実証しています。 M1 Ultra では、PyTorch ベンチマークは最大 20 倍、平均 8.3 倍の高速化を実現しました。この数値は、セッションで示された PyTorch ベンチマークから取得されます。

TensorFlow Metal custom op

(09:21) Metal にアクセスするように TensorFlow 演算をカスタマイズするには、コアは次のとおりです。TF_MetalStreamプロトコル。これは、現在のコマンド バッファー、CPU 側同期のためのディスパッチ キュー、および GPU 作業を送信するためのメソッドを提供します。

@protocol TF_MetalStream

- (id <MTLCommandBuffer>)currentCommandBuffer;
- (dispatch_queue_t)queue;
- (void)commit;
- (void)commitAndWait;

@end

キーポイント:

  • currentCommandBuffer現在 GPU カーネルをエンコードしている Metal コマンド バッファを返します。 -queueディスパッチキューを返します。複数のスレッドが作業を送信するときにエンコード プロセスをシリアル化するために使用されます。 -commit現在のコマンド バッファを GPU に送信します。 -commitAndWait現在のコマンド バッファが完了するまで待機します。送信シーケンスのデバッグに適しています。

(10:04) カスタム op を実装するには 3 つの手順があります。op を登録し、Metal ストリームを使用して op を実装し、トレーニング スクリプトにダイナミック ライブラリをインポートします。

// Register the operation
REGISTER_OP("ZeroOut")
    .Input("to_zero: int32")
    .Output("zeroed: int32")
    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      return Status::OK();
    });

キーポイント:

  • REGISTER_OP("ZeroOut")TensorFlowに名前を登録するZeroOut手術。 -.Input("to_zero: int32")入力テンソルの名前と型を宣言します。 -.Output("zeroed: int32")出力テンソルの名前と型を宣言します。 -SetShapeFn出力形状が入力形状と同じであることを示します。

(10:41) 計算関数では、まず入力と出力を取得し、次に渡します。TF_GetStreamMetal ストリームを取得し、ストリームのキュー内の GPU カーネルをエンコードします。

// Define Compute function
void MetalZeroOut::Compute(TF_OpKernelContext *ctx) {
    // Get input and allocate outputs
    TF_Tensor* input = nullptr;
    TF_GetInput(ctx, 0, &input, status);
    TF_Tensor* output;
    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, input.shape(), &output));

    // Use TF_MetalStream to encode the custom op
    id<TF_MetalStream> metalStream = (id<TF_MetalStream>)(TF_GetStream(ctx, status));
    dispatch_sync(metalStream.queue, ^() {
        id<MTLCommandBuffer> commandBuffer = metalStream.currentCommandBuffer;
        // Create encoder and encode GPU kernel
        [metalStream commit];
    });

    // Delete the TF_Tensors
    TF_DeleteTensor(input);
    TF_DeleteTensor(output);
}

キーポイント:

  • TF_GetInputTensorFlow によって渡された入力テンソルを取得します。 -allocate_output出力テンソル用のスペースを割り当てます。 -TF_GetStream現在の op コンテキストから Metal ストリームを取得します。 -dispatch_sync(metalStream.queue, ^{ ... })マルチスレッドの送信が互いにステップアップするのを避けるために、ストリームによって提供されるキューにエンコード プロセスを置きます。 -metalStream.currentCommandBufferこれは、その後のエンコーダの作成と GPU カーネルのエンコードのためのエントリ ポイントです。 -[metalStream commit]この GPU 作品を送信してください。 -TF_DeleteTensor入力および出力のテンソル参照を解放します。

(11:30) ダイナミック ライブラリが構築されたら、Python トレーニング スクリプトにロードできます。

# Import operation in python script for training
import tensorflow as tf
zero_out_module = tf.load_op_library('./zero_out.so')
print(zero_out_module.zero_out([[1, 2], [3, 4]]).numpy())

キーポイント:

  • import tensorflow as tfTensorFlow の紹介。 -tf.load_op_library('./zero_out.so')カスタム操作の共有動的ライブラリをロードします。 -zero_out_module.zero_out(...)通常の Python 関数と同様に、登録された op を呼び出します。 -.numpy()出力を簡単に検査できるように、結果を NumPy 表現に変換します。

(12:00) Apple は、NeRF の例を使ってカスタム操作の価値を説明しています。オリジナルの 2 層 MLP バージョンはエポックごとに約 100 秒かかり、30 分のトレーニング後でもまだぼやけています。カスタム ハッシュ テーブル オペレーションを追加すると、各エポックに約 10 秒かかり、明確なモデルをより速く取得できるようになります。この例は、TensorFlow でネイティブにサポートされていないハッシュ テーブルも Metal プラグインで実装できることを証明しています。

MPS Graph shared events

(19:21) 計算グラフとポストプロセス グラフが異なるコマンド キューで実行される場合、明示的な同期が必要です。共有イベントの使用法は、最初の記述子信号、2 番目の記述子待機です。

// Using shared events
let executionDescriptor = MPSGraphExecutionDescriptor()
let event = MTLCreateSystemDefaultDevice()!.makeSharedEvent()!
executionDescriptor.signal(event, atExecutionEvent: .completed, value: 1)

let fetch = computeGraph.runAsync(with: commandQueue1,
                                  feeds: [input0Tensor: input0,
                                          input1Tensor: input1],
                                  targetTensors: [finalTensor],
                                  targetOperations: nil,
                                  executionDescriptor: executionDescriptor)

let executionDescriptor2 = MPSGraphExecutionDescriptor()
executionDescriptor2.wait(for: event, value: 1)

let fetch2 = postProcessGraph.runAsync(with: commandQueue2,
                                       feeds: [input0Tensor: fetch[finalTensor]!,
                                               input1Tensor: input1],
                                       targetTensors: [finalTensor],
                                       targetOperations: nil,
                                       executionDescriptor: executionDescriptor2)

キーポイント:

  • MPSGraphExecutionDescriptor()実行グラフを作成するときに使用される実行記述子。 -makeSharedEvent()Metal デバイスで共有イベントを作成します。 -signal(... .completed, value: 1)計算グラフが完了したときにシグナルを発行する必要があります。 -computeGraph.runAsync(...)存在するcommandQueue1計算グラフを非同期で実行します。 -executionDescriptor2.wait(for: event, value: 1)2 番目のグラフに、同じイベントが指定された値に達するまで待機するように要求します。 -postProcessGraph.runAsync(...)存在するcommandQueue2後処理グラフを実行しますが、依存関係はイベント制約を超えています。

MPS Graph の新しい操作: LSTM、プーリング インデックス、ランダムおよびテンソル操作

(22:03) RNN、LSTM、および GRU はリカレント ニューラル ネットワークで一般的に使用される層です。以前は、複雑なサブグラフを手動で作成できましたが、現在は直接使用できるようになりました。MPSGraphLSTMDescriptorそしてgraph.LSTM(...)

let descriptor = MPSGraphLSTMDescriptor()

descriptor.inputGateActivation = .sigmoid
descriptor.forgetGateActivation = .sigmoid
descriptor.cellGateActivation = .tanh
descriptor.outputGateActivation = .sigmoid
descriptor.activation = .tanh
descriptor.bidirectional = false
descriptor.training = true

let lstm = graph.LSTM(inputTensor,
                      recurrentWeight: recurrentWeightsTensor,
                      inputWeight: weightsTensor,
                      bias: nil,
                      initState: nil,
                      initCell: nil,
                      descriptor: descriptor,
                      name: nil)

キーポイント:

  • MPSGraphLSTMDescriptor()LSTM 構成を保存します。
  • 4 つのゲート アクティベーションは、それぞれ入力ゲート、忘却ゲート、セルゲート、出力ゲートのアクティベーション機能を設定します。 -descriptor.activationセル出力で使用されるアクティベーション関数を設定します。 -bidirectional = false一方向の LSTM を表します。 -training = trueこの LSTM によって使用されるトレーニング パスを示します。 -graph.LSTM(...)LSTM ユニットを計算グラフに追加し、入力、反復重み、入力重み、および記述子を渡します。

(23:35) Max pooling API は、最大値が存在するインデックスを返すことができます。トレーニング中に、バックプロパゲーションによって勾配が最大位置に戻され、インデックスの再利用は PyTorch と TensorFlow で最大 6 倍高速になります。

// Forward pass
let descriptor = MPSGraphPooling4DOpDescriptor(kernelSizes: @[1,1,3,3],
                                               paddingStyle: .TF_SAME)
descriptor.returnIndicesMode = .globalFlatten4D

let [poolingTensor, indicesTensor] = graph.maxPooling4DReturnIndices(sourceTensor,
                                                                     descriptor: descriptor,
                                                                     name: nil)

// Backward pass
let outputShape = graph.shapeOf(destination, name: nil)
let gradientTensor = graph.maxPooling4DGradient(gradient: gradientTensor,
                                                indices: indicesTensor,
                                                outputShape: outputShape,
                                                descriptor: descriptor,
                                                name: nil)

キーポイント:

  • MPSGraphPooling4DOpDescriptorプーリング ウィンドウとパディングについて説明します。 -returnIndicesMode = .globalFlatten4Dインデックスの戻りモードを指定します。 -maxPooling4DReturnIndicesプーリング結果とインデックスの両方を返します。 -indicesTensorトレーニング パイプラインの逆フェーズにキャッシュできます。 -maxPooling4DGradientインデックスを使用して、前方パスで最大値が取得された位置に勾配を戻します。

(24:42) 新しいランダム操作は Philox アルゴリズムを使用します。同じシードを与えると、TensorFlow と一致する結果が返され、トレーニング グラフの重みの初期化に使用できます。

// Declare Philox state tensor
let stateTensor = graph.randomPhiloxStateTensor(seed: 2022, name: nil)

// Declare RandomOp descriptor
let descriptor = MPSGraphRandomOpDescriptor(distribution: .truncatedNormal,
                                            dataType: .float32)
descriptor.mean = -1.0
 descriptor.standardDeviation = 2.5
descriptor.min = descriptor.mean - 2 * descriptor.standardDeviation
descriptor.max = descriptor.mean + 2 * descriptor.standardDeviation

let [randomTensor, stateTensor] = graph.randomTensor(shapeTensor: shapeTensor,
                                                     descriptor: descriptor,
                                                     stateTensor: stateTensor,
                                                     name: nil)

キーポイント:

  • randomPhiloxStateTensor(seed:)シードされた Philox 状態テンソルを作成します。 -MPSGraphRandomOpDescriptor分散タイプとデータタイプを指定します。 -.truncatedNormal切り捨てられた正規分布を表します。 -meanstandardDeviationminmax配布範囲を定義します。 -randomTensor(...)ランダム テンソルと新しい状態テンソルを出力します。新しい状態は引き続き次のランダム操作に渡されます。

(26:18) MPS Graph は、次元の拡張、次元の圧縮、セグメント、スタック、座標の生成を行う一連のテンソル変換操作も追加します。

// Expand the input tensor dimensions, 4x2 -> 4x1x2
let expandedTensor = graph.expandDims(inputTensor,
                                      axis: 1,
                                      name: nil)

// Squeeze the input tensor dimensions, 4x1x2 -> 4x2
let squeezedTensor = graph.squeeze(expandedTensor,
                                   axis: 1,
                                   name: nil)

// Split the tensor in two, 4x2 -> (4x1, 4x1)
let [split1, split2] = graph.split(squeezedTensor,
                                   numSplits: 2,
                                   axis: 0,
                                   name: nil)

// Stack the tensor back together, (4x1, 4x1) -> 2x4x1
let stackedTensor = graph.stack([split1, split2],
                                axis: 0,
                                name: nil)

キーポイント:

  • expandDims(... axis: 1)次元 1 に新しい次元を挿入します。 -squeeze(... axis: 1)次元 1 の単一要素次元を削除します。 -split(... numSplits: 2, axis: 0)テンソルを 0 次元に沿って 2 つの等しい部分に切ります。 -stack([split1, split2], axis: 0)複数のテンソルを 0 次元に沿って再スタックします。

重要ポイント

  • やるべきこと: MPS デバイス スイッチを PyTorch トレーニング スクリプトに追加します

    • 実行する価値がある理由: このセッションは、PyTorch 1.12 の MPS バックエンドが公式パッケージに入ったことを示しています。モデルと入力が MPS デバイスに移動された後、後続のテンソル演算は GPU に送られます。
    • 開始方法: 最初に使用しますtorch.backends.mps.is_available()デバイスを作成してから、model.to(device=mpsDevice)そして入力テンソルdevice=mpsDevice完了。
  • やるべきこと: GPU 上の TensorFlow トレーニングでカスタム損失を残す

    • 実行する価値がある理由: CPU への損失をカスタマイズすると、同期オーバーヘッドが発生し、GPU タイムラインに穴が残ります。
    • 開始方法: を使用しますREGISTER_OP登録操作、Compute合格TF_GetStream得るTF_MetalStream、 またmetalStream.queue中程度のエンコーディングのメタル カーネル。
  • やるべきこと: 共有イベントを使用して ML 後処理パイプラインをアンパックする

    • 実行する価値がある理由: コンピューティングと後処理を別のコマンド キューに配置すると、GPU の並列処理をさらに活用できますが、データの依存関係を処理する必要があります。
    • 開始方法: 計算グラフ内MPSGraphExecutionDescriptorシグナル共有イベントで、ポストプロセス グラフの記述子で同じイベントを待ちます。
  • やるべきこと: トレーニング グラフの手書きの LSTM サブグラフを置き換えます

    • 実行する価値がある理由: MPS Graph は新しい LSTM 操作を追加し、リカレント ユニットに必要な GPU 作業を効率的にエンコードできることがセッションで明確に指摘されています。
    • 開始方法: 作成するMPSGraphLSTMDescriptor、ゲートのアクティブ化、トレーニング、その他の属性を設定してから呼び出します。graph.LSTM(...)図にユニットを追加します。
  • やるべきこと: プールバックプロパゲーション中にインデックスをキャッシュする

    • 実行する価値がある理由: 最大プーリング リターン インデックス API を使用すると、リバース フェーズでフォワードの最大値の位置を再利用できます。セッションによって与えられる PyTorch および TensorFlow トレーニングの最大加速は 6 倍です。
    • 開始方法: フォワードを使用しますgraph.maxPooling4DReturnIndices(...)テイクアウトindicesTensor、逆方向に使用しますgraph.maxPooling4DGradient(...)同じインデックスを渡します。

関連セッション

コメント

GitHub Issues · utterances