Tensor計算を説明するのは非常に簡単ですが、その計算がいつどのように実行されるかは、 Tensorにどのバックエンドが使用されるか、および結果がホスト CPU でいつ必要になるかによって異なります。
Tensorの操作は舞台裏で GPU やTPU などのアクセラレータにディスパッチされるか、アクセラレータが使用できない場合は CPU 上で実行されます。これは自動的に行われるため、高レベルのインターフェイスを使用して複雑な並列計算を簡単に実行できるようになります。ただし、このディスパッチがどのように発生するかを理解し、最適なパフォーマンスが得られるようにカスタマイズできると便利です。
Swift for TensorFlow には、高速計算を実行するための 2 つのバックエンド (TensorFlow イーガー モードと X10) があります。デフォルトのバックエンドは TensorFlow 熱心モードですが、これはオーバーライドできます。これらのさまざまなバックエンドの使用方法を説明する対話型チュートリアルが利用可能です。
TensorFlow イーガー モード
TensorFlow イーガー モード バックエンドは、 TensorFlow C APIを利用して、各Tensor操作が発生したときに GPU または CPU に送信します。次に、その操作の結果が取得され、次の操作に渡されます。
このオペレーションごとのディスパッチは理解しやすく、コード内で明示的に構成する必要はありません。ただし、多くの場合、多くの小さな操作を送信することによるオーバーヘッドと、操作のグラフが存在する場合に発生する可能性のある操作の融合と最適化の欠如により、最適なパフォーマンスは得られません。最後に、TensorFlow イーガー モードは TPU と互換性がなく、CPU と GPU でのみ使用できます。
X10 (XLA ベースのトレース)
X10 は、遅延テンソル トレースとXLA 最適化コンパイラを使用して、多くの場合、オペレーションごとのディスパッチよりもパフォーマンスを大幅に向上させる Swift for TensorFlow バックエンドの名前です。さらに、機械学習モデル内で見られる種類の計算用に特に最適化されたアクセラレータであるTPUとの互換性が追加されます。
Tensor計算での X10 の使用はデフォルトではないため、このバックエンドをオプトインする必要があります。これは、 Tensorが XLA デバイス上に配置されるように指定することで行われます。
let tensor1 = Tensor<Float>([0.0, 1.0, 2.0], on: Device.defaultXLA)
let tensor2 = Tensor<Float>([1.5, 2.5, 3.5], on: Device.defaultXLA)
それ以降の計算の記述は、TensorFlow 熱心モードの場合とまったく同じです。
let tensor3 = tensor1 + tensor2
Tensorを作成するときに、どのような種類のアクセラレータを使用するか、複数のアクセラレータが利用可能な場合はどのアクセラレータを使用するかなど、さらに詳細を提供できます。たとえば、次のコマンドを使用して、2 番目の TPU デバイス上にTensorを作成できます (プログラムが実行されているホストから認識できると仮定します)。
let tpuTensor = Tensor<Float>([0.0, 1.0, 2.0], on: Device(kind: .TPU, ordinal: 1, backend: .XLA))
デバイス間でのTensorの暗黙的な移動は実行されないため、異なるデバイス上の 2 つのTensorが操作で一緒に使用される場合、ランタイム エラーが発生します。 Tensorの内容を新しいデバイスに手動でコピーするには、 Tensor(copying:to:)イニシャライザを使用できます。モデルやオプティマイザーなど、内部にTensorを含む一部の大規模な構造には、内部のすべてのTensorを 1 ステップで新しいデバイスに移動するためのヘルパー関数があります。
TensorFlow イーガー モードとは異なり、X10 バックエンドを使用する操作は、発生したときに個別にディスパッチされません。代わりに、アクセラレータへのディスパッチは、計算された値をホストに読み戻すか、明示的なバリアを配置することによってのみトリガーされます。これがどのように機能するかというと、ランタイムはホストに読み取られる値 (または手動バリアの前の最後の計算) から開始し、その値をもたらす計算のグラフをトレースします。
このトレースされたグラフは、XLA HLO 中間表現に変換され、XLA コンパイラーに渡されて、アクセラレータでの実行用に最適化およびコンパイルされます。そこから、計算全体がアクセラレータに送信され、最終結果が得られます。
計算は時間のかかるプロセスであるため、X10 は、グラフで表現され、何度も実行される大規模並列計算で使用するのが最適です。ハッシュ値とキャッシュが使用されるため、同一のグラフは一意の構成ごとに 1 回だけコンパイルされます。
機械学習モデルの場合、トレーニング プロセスには、モデルが同じ一連の計算を何度も繰り返すループが含まれることがよくあります。これらの各パスを、内部に繰り返し単位が含まれる 1 つの長いグラフではなく、同じトレースの繰り返しとして見たい場合があります。これは、トレースを終了したいコード内の位置にLazyTensorBarrier()関数の呼び出しを手動で挿入することで有効になります。
X10 での混合精度のサポート
X10 を介した混合精度のトレーニングがサポートされており、それを制御するために低レベル API と高レベル API の両方が提供されています。低レベル API は、完全精度と低減精度の間で変換するtoReducedPrecisionとtoFullPrecision 2 つの計算プロパティと、精度をクエリするisReducedPrecisionを提供します。 Tensorに加えて、この API を使用して、モデルとオプティマイザーを完全精度と低精度の間で変換できます。
低減された精度に変換しても、 Tensorの論理型は変更されないことに注意してください。 tがTensor<Float>の場合、 t.toReducedPrecisionも、基になる表現の精度が低下したTensor<Float>です。
デバイスと同様に、異なる精度のテンソル間の演算は許可されません。これにより、ユーザーが検出するのが難しい、サイレントで不要な 32 ビット浮動小数点への昇格が回避されます。