警告:このプロジェクトは廃止されました。 Swift for TensorFlow は、機械学習、コンパイラー、微分可能プログラミング、システム設計などにわたる最新の研究を組み込んだ、機械学習用の次世代プラットフォームの実験でした。 2021年2月にアーカイブされました。

このページは Cloud Translation API によって翻訳されました。

X10 の紹介

%install '.package(url: "https://github.com/tensorflow/swift-models", .branch("tensorflow-0.12"))' Datasets ImageClassificationModels
print("\u{001B}[2J")

TensorFlow.org で見る

Google Colab で実行する

GitHub でソースを表示

デフォルトでは、Swift For TensorFlow は、即時ディスパッチを使用してテンソル操作を実行します。これにより、迅速な反復が可能になりますが、機械学習モデルをトレーニングするための最もパフォーマンスの高いオプションではありません。

X10 テンソルライブラリは、テンソルトレースとXLA コンパイラを活用して、TensorFlow 用の高性能バックエンドを Swift に追加します。このチュートリアルでは、X10 を紹介し、GPU または TPU で実行するトレーニングループを更新するプロセスについて説明します。

Eager 対 X10 テンソル

Swift for TensorFlow の高速計算は、Tensor タイプを通じて実行されます。テンソルはさまざまな操作に参加でき、機械学習モデルの基本的な構成要素です。

デフォルトでは、Tensor は即時実行を使用してオペレーションごとに計算を実行します。各 Tensor には、どのハードウェアに接続されているか、どのバックエンドが使用されているかを記述する関連デバイスがあります。

import TensorFlow
import Foundation

let eagerTensor1 = Tensor([0.0, 1.0, 2.0])
let eagerTensor2 = Tensor([1.5, 2.5, 3.5])
let eagerTensorSum = eagerTensor1 + eagerTensor2
print(eagerTensorSum)

[1.5, 3.5, 5.5]

print(eagerTensor1.device)

Device(kind: .CPU, ordinal: 0, backend: .TF_EAGER)

このノートブックを GPU 対応インスタンスで実行している場合は、そのハードウェアが上記のデバイスの説明に反映されていることがわかります。 Eager ランタイムは TPU をサポートしていないため、TPU のいずれかをアクセラレータとして使用している場合は、CPU がハードウェアターゲットとして使用されていることがわかります。

Tensor を作成するとき、デフォルトの Eager モードデバイスは代替デバイスを指定することでオーバーライドできます。これは、X10 バックエンドを使用した計算の実行をオプトインする方法です。

let x10Tensor1 = Tensor([0.0, 1.0, 2.0], on: Device.defaultXLA)
let x10Tensor2 = Tensor([1.5, 2.5, 3.5], on: Device.defaultXLA)
let x10TensorSum = x10Tensor1 + x10Tensor2
print(x10TensorSum)

[1.5, 3.5, 5.5]

print(x10Tensor1.device)

Device(kind: .CPU, ordinal: 0, backend: .XLA)

これを GPU 対応インスタンスで実行している場合は、そのアクセラレータが X10 tensor のデバイスにリストされているのが表示されるはずです。即時実行とは異なり、これを TPU 対応インスタンスで実行している場合は、計算にそのデバイスが使用されていることがわかります。 X10 は、Swift 内で TensorFlow 用の TPU を活用する方法です。

デフォルトの Eager デバイスと X10 デバイスは、システム上の最初のアクセラレータの使用を試みます。 GPU が接続されている場合、は最初に利用可能な GPU を使用します。 TPU が存在する場合、X10 はデフォルトで最初の TPU コアを使用します。アクセラレータが見つからないかサポートされていない場合は、デフォルトのデバイスが CPU に戻ります。

デフォルトの Eager デバイスと XLA デバイスを超えて、デバイスに特定のハードウェアとバックエンドターゲットを提供できます。

// let tpu1 = Device(kind: .TPU, ordinal: 1, backend: .XLA)
// let tpuTensor1 = Tensor([0.0, 1.0, 2.0], on: tpu1)

イーガーモードモデルのトレーニング

デフォルトの積極的実行モードを使用してモデルをセットアップおよびトレーニングする方法を見てみましょう。この例では、 swift-models リポジトリの単純な LeNet-5 モデルと MNIST 手書き数字分類データセットを使用します。

まず、MNIST データセットをセットアップしてダウンロードします。

import Datasets

let epochCount = 5
let batchSize = 128
let dataset = MNIST(batchSize: batchSize)

Loading resource: train-images-idx3-ubyte
File does not exist locally at expected path: /home/kbuilder/.cache/swift-models/datasets/MNIST/train-images-idx3-ubyte and must be fetched
Fetching URL: https://storage.googleapis.com/cvdf-datasets/mnist/train-images-idx3-ubyte.gz...
Archive saved to: /home/kbuilder/.cache/swift-models/datasets/MNIST
Loading resource: train-labels-idx1-ubyte
File does not exist locally at expected path: /home/kbuilder/.cache/swift-models/datasets/MNIST/train-labels-idx1-ubyte and must be fetched
Fetching URL: https://storage.googleapis.com/cvdf-datasets/mnist/train-labels-idx1-ubyte.gz...
Archive saved to: /home/kbuilder/.cache/swift-models/datasets/MNIST
Loading resource: t10k-images-idx3-ubyte
File does not exist locally at expected path: /home/kbuilder/.cache/swift-models/datasets/MNIST/t10k-images-idx3-ubyte and must be fetched
Fetching URL: https://storage.googleapis.com/cvdf-datasets/mnist/t10k-images-idx3-ubyte.gz...
Archive saved to: /home/kbuilder/.cache/swift-models/datasets/MNIST
Loading resource: t10k-labels-idx1-ubyte
File does not exist locally at expected path: /home/kbuilder/.cache/swift-models/datasets/MNIST/t10k-labels-idx1-ubyte and must be fetched
Fetching URL: https://storage.googleapis.com/cvdf-datasets/mnist/t10k-labels-idx1-ubyte.gz...
Archive saved to: /home/kbuilder/.cache/swift-models/datasets/MNIST

次に、モデルとオプティマイザーを構成します。

import ImageClassificationModels

var eagerModel = LeNet()
var eagerOptimizer = SGD(for: eagerModel, learningRate: 0.1)

ここで、基本的な進捗状況の追跡とレポートを実装します。すべての中間統計は、トレーニングが実行され、レポート中にのみscalarized()が呼び出される同じデバイス上にテンソルとして保持されます。これは、遅延テンソルの不必要な実体化を回避するため、後で X10 を使用するときに特に重要になります。

struct Statistics {
    var correctGuessCount = Tensor<Int32>(0, on: Device.default)
    var totalGuessCount = Tensor<Int32>(0, on: Device.default)
    var totalLoss = Tensor<Float>(0, on: Device.default)
    var batches: Int = 0
    var accuracy: Float { 
        Float(correctGuessCount.scalarized()) / Float(totalGuessCount.scalarized()) * 100 
    } 
    var averageLoss: Float { totalLoss.scalarized() / Float(batches) }

    init(on device: Device = Device.default) {
        correctGuessCount = Tensor<Int32>(0, on: device)
        totalGuessCount = Tensor<Int32>(0, on: device)
        totalLoss = Tensor<Float>(0, on: device)
    }

    mutating func update(logits: Tensor<Float>, labels: Tensor<Int32>, loss: Tensor<Float>) {
        let correct = logits.argmax(squeezingAxis: 1) .== labels
        correctGuessCount += Tensor<Int32>(correct).sum()
        totalGuessCount += Int32(labels.shape[0])
        totalLoss += loss
        batches += 1
    }
}

最後に、5 エポックのトレーニングループを通じてモデルを実行します。

print("Beginning training...")

for (epoch, batches) in dataset.training.prefix(epochCount).enumerated() {
    let start = Date()
    var trainStats = Statistics()
    var testStats = Statistics()

    Context.local.learningPhase = .training
    for batch in batches {
        let (images, labels) = (batch.data, batch.label)
        let 𝛁model = TensorFlow.gradient(at: eagerModel) { eagerModel -> Tensor<Float> in
            let ŷ = eagerModel(images)
            let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)
            trainStats.update(logits: ŷ, labels: labels, loss: loss)
            return loss
        }
        eagerOptimizer.update(&eagerModel, along: 𝛁model)
    }

    Context.local.learningPhase = .inference
    for batch in dataset.validation {
        let (images, labels) = (batch.data, batch.label)
        let ŷ = eagerModel(images)
        let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)
        testStats.update(logits: ŷ, labels: labels, loss: loss)
    }

    print(
        """
        [Epoch \(epoch)] \
        Training Loss: \(String(format: "%.3f", trainStats.averageLoss)), \
        Training Accuracy: \(trainStats.correctGuessCount)/\(trainStats.totalGuessCount) \
        (\(String(format: "%.1f", trainStats.accuracy))%), \
        Test Loss: \(String(format: "%.3f", testStats.averageLoss)), \
        Test Accuracy: \(testStats.correctGuessCount)/\(testStats.totalGuessCount) \
        (\(String(format: "%.1f", testStats.accuracy))%) \
        seconds per epoch: \(String(format: "%.1f", Date().timeIntervalSince(start)))
        """)
}

Beginning training...
[Epoch 0] Training Loss: 0.528, Training Accuracy: 50154/59904 (83.7%), Test Loss: 0.168, Test Accuracy: 9468/10000 (94.7%) seconds per epoch: 11.9
[Epoch 1] Training Loss: 0.133, Training Accuracy: 57488/59904 (96.0%), Test Loss: 0.107, Test Accuracy: 9659/10000 (96.6%) seconds per epoch: 11.7
[Epoch 2] Training Loss: 0.092, Training Accuracy: 58193/59904 (97.1%), Test Loss: 0.069, Test Accuracy: 9782/10000 (97.8%) seconds per epoch: 11.8
[Epoch 3] Training Loss: 0.071, Training Accuracy: 58577/59904 (97.8%), Test Loss: 0.066, Test Accuracy: 9794/10000 (97.9%) seconds per epoch: 11.8
[Epoch 4] Training Loss: 0.059, Training Accuracy: 58800/59904 (98.2%), Test Loss: 0.064, Test Accuracy: 9800/10000 (98.0%) seconds per epoch: 11.8

ご覧のとおり、モデルは予想どおりにトレーニングされ、検証セットに対する精度がエポックごとに向上しました。これは、Swift for TensorFlow モデルがどのように定義され、即時実行を使用して実行されるかです。次に、X10 を利用するためにどのような変更を加える必要があるかを見てみましょう。

X10 モデルのトレーニング

データセット、モデル、オプティマイザーには、デフォルトの積極的実行デバイスで初期化されるテンソルが含まれています。 X10 を使用するには、これらのテンソルを X10 デバイスに移動する必要があります。

let device = Device.defaultXLA
print(device)

Device(kind: .CPU, ordinal: 0, backend: .XLA)

データセットの場合は、トレーニングループでバッチが処理される時点でこれを行うため、熱心な実行モデルからデータセットを再利用できます。

モデルとオプティマイザーの場合は、熱心な実行デバイス上の内部テンソルを使用してそれらを初期化し、X10 デバイスに移動します。

var x10Model = LeNet()
x10Model.move(to: device)

var x10Optimizer = SGD(for: x10Model, learningRate: 0.1)
x10Optimizer = SGD(copying: x10Optimizer, to: device)

トレーニングループに必要な変更は、いくつかの特定のポイントで行われます。まず、トレーニングデータのバッチを X10 デバイスに移動する必要があります。これは、各バッチが取得されるときにTensor(copying:to:)を介して行われます。

次の変更は、トレーニングループ中にトレースを切断する場所を示すことです。 X10 は、コード内で必要なテンソル計算をトレースし、そのトレースの最適化された表現をジャストインタイムでコンパイルすることによって機能します。トレーニングループの場合、同じ操作を何度も繰り返すことになるため、トレース、コンパイル、再利用するには理想的なセクションです。

Tensor からの値を明示的に要求するコードがない場合 (これらは通常.scalarsまたは.scalarized()呼び出しとして目立ちます)、X10 はすべてのループ反復をまとめてコンパイルしようとします。これを防止し、特定のポイントでトレースをカットするには、オプティマイザがモデルの重みを更新し、検証中に損失と精度が取得された後、 LazyTensorBarrier()を配置します。これにより、トレーニングループの各ステップと検証中の推論の各バッチという 2 つの再利用トレースが作成されます。

これらの変更により、次のトレーニングループが作成されます。

print("Beginning training...")

for (epoch, batches) in dataset.training.prefix(epochCount).enumerated() {
    let start = Date()
    var trainStats = Statistics(on: device)
    var testStats = Statistics(on: device)

    Context.local.learningPhase = .training
    for batch in batches {
        let (eagerImages, eagerLabels) = (batch.data, batch.label)
        let images = Tensor(copying: eagerImages, to: device)
        let labels = Tensor(copying: eagerLabels, to: device)
        let 𝛁model = TensorFlow.gradient(at: x10Model) { x10Model -> Tensor<Float> in
            let ŷ = x10Model(images)
            let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)
            trainStats.update(logits: ŷ, labels: labels, loss: loss)
            return loss
        }
        x10Optimizer.update(&x10Model, along: 𝛁model)
        LazyTensorBarrier()
    }

    Context.local.learningPhase = .inference
    for batch in dataset.validation {
        let (eagerImages, eagerLabels) = (batch.data, batch.label)
        let images = Tensor(copying: eagerImages, to: device)
        let labels = Tensor(copying: eagerLabels, to: device)
        let ŷ = x10Model(images)
        let loss = softmaxCrossEntropy(logits: ŷ, labels: labels)
        LazyTensorBarrier()
        testStats.update(logits: ŷ, labels: labels, loss: loss)
    }

    print(
        """
        [Epoch \(epoch)] \
        Training Loss: \(String(format: "%.3f", trainStats.averageLoss)), \
        Training Accuracy: \(trainStats.correctGuessCount)/\(trainStats.totalGuessCount) \
        (\(String(format: "%.1f", trainStats.accuracy))%), \
        Test Loss: \(String(format: "%.3f", testStats.averageLoss)), \
        Test Accuracy: \(testStats.correctGuessCount)/\(testStats.totalGuessCount) \
        (\(String(format: "%.1f", testStats.accuracy))%) \
        seconds per epoch: \(String(format: "%.1f", Date().timeIntervalSince(start)))
        """)
}

Beginning training...
[Epoch 0] Training Loss: 0.421, Training Accuracy: 51888/59904 (86.6%), Test Loss: 0.134, Test Accuracy: 9557/10000 (95.6%) seconds per epoch: 18.6
[Epoch 1] Training Loss: 0.117, Training Accuracy: 57733/59904 (96.4%), Test Loss: 0.085, Test Accuracy: 9735/10000 (97.3%) seconds per epoch: 14.9
[Epoch 2] Training Loss: 0.080, Training Accuracy: 58400/59904 (97.5%), Test Loss: 0.068, Test Accuracy: 9791/10000 (97.9%) seconds per epoch: 13.1
[Epoch 3] Training Loss: 0.064, Training Accuracy: 58684/59904 (98.0%), Test Loss: 0.056, Test Accuracy: 9804/10000 (98.0%) seconds per epoch: 13.5
[Epoch 4] Training Loss: 0.053, Training Accuracy: 58909/59904 (98.3%), Test Loss: 0.063, Test Accuracy: 9779/10000 (97.8%) seconds per epoch: 13.4

X10 バックエンドを使用したモデルのトレーニングは、以前の熱心な実行モデルと同じ方法で進められるはずです。最初のバッチの前と最初のエポックの終了時に、それらの時点で固有のトレースがジャストインタイムでコンパイルされるため、遅延に気づいたかもしれません。アクセラレータを接続してこれを実行している場合、その時点以降のトレーニングが Eager モードよりも速く進行していることがわかるはずです。

初期トレースのコンパイル時間とスループットの高速化にはトレードオフがありますが、ほとんどの機械学習モデルでは、繰り返しの操作によるスループットの増加はコンパイルのオーバーヘッドを相殺する以上の効果があります。実際、一部のトレーニングケースでは、X10 によるスループットの 4 倍以上の向上が確認されています。

前に述べたように、X10 を使用すると、TPU の操作が可能になるだけでなく簡単になり、Swift for TensorFlow モデルのアクセラレータのクラス全体のロックが解除されます。

X10 の紹介

Eager 対 X10 テンソル

イーガーモード モデルのトレーニング

X10 モデルのトレーニング

イーガーモードモデルのトレーニング