Open2024/05/15にコメント追加1

MPS backend

Pytorch MPSバックエンド

MPSデバイスは、Metalプログラミングフレームワークを用いてMacOSデバイスでの高性能なトレーニングを可能にします。これは、機械学習の計算グラフおよびプリミティブを非常に効率的なMetal Performance Shaders GraphフレームワークおよびMetal Performance Shadersフレームワークによって提供される調整されたカーネルにマッピングする新しいデバイスを導入します。

新しいMPSバックエンドはPyTorchエコシステムを拡張し、既存のスクリプトに対してGPU上での操作を設定および実行する機能を提供します。

始めるには、単にTensorとModuleをmpsデバイスに移動します。

# MPS (Metal Performance Shaders)を利用可能か確認
if not torch.backends.mps.is_available():
  # MPSがビルド時に無効化されていた場合
  if not torch.backends.mps.is_built():
    print("MPSが無効です。現在のPyTorchはMPS有効でのビルドが行われていません。")
  # MPSが有効でも利用できない場合
  else:
    print("MPSが無効です。macOSバージョンが12.3以降でないか、MPS対応デバイスがこのマシン上にありません。")
else:
  # MPSデバイスを定義
  mps_device = torch.device("mps")

  # MPSデバイス上に直接テンソルを作成 (2通りの方法)
  x = torch.ones(5, device=mps_device)
  # または
  x = torch.ones(5, device="mps")

  # 演算はMPSデバイス上で行われる
  y = x * 2

  # モデルを他のデバイスへ送るようにMPSデバイスへ送る
  model = YourFavoriteNet()
  model.to(mps_device)

  # 以降のモデル呼び出しはMPSデバイス上で実行
  pred = model(x)

MPS backend — PyTorch 2.3 documentation