WWDC Quick Look 💓 By SwiftGGTeam
Accelerate machine learning with Metal Performance Shaders Graph

Accelerate machine learning with Metal Performance Shaders Graph

元の動画を見る

ハイライト

MPSGraph は、TensorFlow Metal プラグイン、制御フロー API (if/else/for/while)、ステンシル オペレーター、および GatherND オペレーターを追加して、GPU アクセラレーションによる機械学習のトレーニングと推論をより柔軟にします。

主な内容

Mac 上で TensorFlow を使用してモデルをトレーニングしますが、デフォルトでは CPU のみが使用されます。 ResNet50 のエポックには 20 分かかります。 GPU アクセラレーションを使用したいが、CUDA は Mac では使用できません。

Apple のソリューションは TensorFlow Metal Plugin です。これは MPSGraph に基づいており、TensorFlow の計算グラフを Metal GPU に変換します。

##詳細

TensorFlow Metal Plugin

02:19

インストール方法:

pip install tensorflow-macos
pip install tensorflow-metal

TensorFlow をインストールすると、Mac の GPU が自動的に認識されます。

import tensorflow as tf

# 列出可用设备
print(tf.config.list_physical_devices())
# 输出包含 GPU 设备

# 定义 ResNet50 模型
model = tf.keras.applications.ResNet50(
    weights=None,
    input_shape=(224, 224, 3),
    classes=1000
)

# 训练 - 自动使用 Metal GPU
model.compile(optimizer='adam', loss='categorical_crossentropy')
model.fit(x_train, y_train, epochs=10)

03:00

キーポイント:

  • pip install tensorflow-metalプラグインのインストール
  • モデルコードを変更する必要はありません
  • ResNet50のトレーニング速度が約4倍に向上
  • 一部のモデルは最大 8 倍まで加速できます
  • M1 MacBook Proをサポート

コア ML 推論の高速化

05:08

MPSGraph の最適化は Core ML にも利益をもたらします。 BERT モデルの推論速度は M1 で 2 倍に向上し、ResNet50 では 16% 向上しました。

これらの改善は、MPS の基礎となるオペレーターの最適化、特に NHWC および NCHW レイアウトでの Convolution2D のパフォーマンスの向上によってもたらされます。

依存関係の制御: 演算子が最適化されないようにする

08:35

MPSGraph は、出力されない中間演算子を最適化する場合があります。ただし、一部の演算子 (バッチ正規化における移動平均更新など) は、出力が直接使用されない場合でも実行する必要があります。

// 问题:assign 算子不被其他算子依赖,可能被优化掉
let assignOp = graph.assign(variable: runningMean, value: newMean)
let exponentOp = graph.exponent(input: x)

// 解决方案:让 exponent 依赖 assign
let dependentExponent = graph.controlDependency(
    dependent: exponentOp,
    dependentOn: assignOp
)

09:01

キーポイント:

  • controlDependency実行順序を明示的に指定する
  • dependentオペレータは次の場所にいる必要がありますdependentOn後に実行
  • targetOperation をグローバルに追跡する必要はありません
  • バッチ正規化やその他のシナリオに適しています

ステンシル オペレーター: スライディング ウィンドウの計算

09:25

Stencil オペレーターは畳み込みを一般化したもので、任意のスライディング ウィンドウの加重削減をサポートします。ラプラシアンやローカル応答正規化などの演算に適しています。

// 5-point 2D Laplacian stencil
let stencil = graph.stencil(
    source: input,
    descriptor: stencilDescriptor,
    weights: weightsTensor
)

10:28

キーポイント:

  • 2D/3D ステンシルをサポート
  • 複数のリダクション モード (sum、argmin、argmax) をサポート
  • 複数のパディング モードをサポート (反射、clampToZero)
  • 単一のカーネル起動に結合可能

GatherND: N 次元インデックス

11:05

GatherND は、座標インデックスによる N 次元テンソルからのデータ抽出をサポートしており、線形インデックスよりも柔軟です。

// 从 3D 张量中按坐标 gather
// indices: [[0, 1], [2, 3]] 表示取第 (0,1) 行和第 (2,3) 行
let result = graph.gatherND(
    updates: inputTensor,
    indices: indicesTensor,
    batchDimensions: 0
)

12:29

キーポイント:

  • あらゆる次元のインデックスをサポート
  • 不特定寸法の自動スライスコピー
  • ルックアップやその他のシーンの埋め込みに適しています
  • 各座標はスライスの開始点を指定します

MPSGraphExecutable: プリコンパイルされたグラフ

14:42

デフォルトでは、MPSGraph は最初の実行時にコンパイルされます。 MPSGraphExecutable を使用すると、事前にコンパイルし、コンパイルのタイミングを制御できます。

// 定义图
let graph = MPSGraph()
let x = graph.placeholder(shape: [1, 224, 224, 3], dataType: .float32, name: "x")
let y = graph.convolution2D(...)

// 编译为 executable
let executable = graph.compile(
    device: device,
    feeds: [x: MPSGraphTensorData],
    targetTensors: [y],
    compilationDescriptor: descriptor
)

// 后续直接执行,跳过编译
let result = executable.run(
    withCommandQueue: commandQueue,
    inputs: [inputData]
)

14:42

キーポイント:

  • compile事前にグラフを作成する
  • run直接実行、再コンパイル不要
  • コンパイルされた実行可能ファイルをキャッシュできます
  • 初回の遅延を回避するための推論シナリオに適しています

型推論を無効にする

16:38

入力サイズが変化するネットワーク (NLP の異なる長さの文など) では、新しいサイズごとに型推論と再コンパイルがトリガーされます。

let descriptor = MPSGraphCompilationDescriptor()
descriptor.disableTypeInference = true

let executable = graph.compile(
    ...,
    compilationDescriptor: descriptor
)

16:38

キーポイント:

  • disableTypeInference = trueスキップ型推論
  • MPSGraph はエンコード中に型を動的に推論します
  • コンパイル時間を大幅に節約します
  • 少数の最適化の機会を犠牲にする

制御フロー: if/else、for、while

19:22

MPSGraph は、条件付き実行とループをサポートする新しい制御フロー API を追加します。

If/else:

let predicate = graph.lessThan(x, y)

let result = graph.ifElse(
    predicate: predicate,
    thenBlock: { graph in
        return graph.addition(x, y, name: nil)
    },
    elseBlock: { graph in
        return graph.subtraction(x, y, name: nil)
    }
)

19:22

For loop:

let result = graph.forLoop(
    iterations: 4,
    initialInputs: [input0],
    body: { graph, index, inputs in
        let result = graph.multiplication(inputs[0], input1, name: nil)
        return [result]
    }
)

20:15

While loop:

let result = graph.whileLoop(
    before: { graph, inputs in
        let condition = graph.lessThan(inputs[0], threshold)
        return [condition, inputs[0]]
    },
    after: { graph, inputs in
        let multiplied = graph.multiplication(inputs[0], multiplier, name: nil)
        return [multiplied]
    },
    initialInputs: [input0]
)

22:09

キーポイント:

  • 制御フローは GPU で実行され、CPU と GPU の同期は必要ありません
  • バッチ正規化、RNN、その他のシナリオに適しています
  • 述語はスカラーではなくテンソルです
  • ループ本体は現在のインデックスと前のラウンドの出力を受け取ります

重要なポイント

  1. Mac 上の TensorFlow Metal プラグインを使用してモデルをトレーニングします。 2 行の pip コマンドをインストールするだけで、コードを変更する必要がなく、ResNet50 のトレーニング速度が 4 倍向上しました。入口API:pip install tensorflow-metal

  2. MPSGraphExecutable を使用して推論グラフをプリコンパイルします。これにより、最初の実行時のコンパイル遅延が回避され、リアルタイム推論シナリオに適しています。入口API:graph.compile() + executable.run()

  3. 可変長入力の型推論を無効にします。 NLP モデルで文の長さが異なると、再コンパイルがトリガーされなくなりました。入口API:compilationDescriptor.disableTypeInference = true

  4. 制御フローを使用してバッチ正規化トレーニングを実装します。グラフはトレーニング モードと推論モードの両方をサポートします。入口API:graph.ifElse(predicate: isTraining, ...)

  5. ステンシル オペレーターを使用してカスタム画像フィルタリングを実装します。ラプラシアンやローカル正規化などの操作を単一のカーネルに組み込むことができます。入口API:graph.stencil(source:descriptor:weights:)

関連セッション

コメント

GitHub Issues · utterances