Przedstawiamy X10

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


Zobacz na TensorFlow.org Uruchom w Google Colab Zobacz źródło w GitHub

Domyślnie Swift For TensorFlow wykonuje operacje tensorowe przy użyciu szybkiego wysyłania. Pozwala to na szybką iterację, ale nie jest najbardziej wydajną opcją w przypadku uczenia modeli uczenia maszynowego.

Biblioteka tensorów X10 dodaje do Swifta wydajny backend dla TensorFlow, wykorzystując śledzenie tensorów i kompilator XLA . W tym samouczku przedstawimy X10 i poprowadzimy Cię przez proces aktualizacji pętli treningowej do działania na procesorach graficznych lub TPU.

Tensory Eager vs. X10

Przyspieszone obliczenia w Swift dla TensorFlow są wykonywane poprzez typ Tensor. Tensory mogą brać udział w różnorodnych operacjach i stanowią podstawowe elementy składowe modeli uczenia maszynowego.

Domyślnie Tensor używa szybkiego wykonywania do wykonywania obliczeń w oparciu o operację po operacji. Każdy Tensor ma powiązane Urządzenie, które opisuje, do jakiego sprzętu jest podłączony i jaki backend jest do niego używany.

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)

Jeśli używasz tego notebooka na instancji obsługującej procesor graficzny, powinieneś zobaczyć ten sprzęt odzwierciedlony w powyższym opisie urządzenia. Środowisko wykonawcze chętnie nie obsługuje procesorów TPU, więc jeśli użyjesz jednego z nich jako akceleratora, zobaczysz, że procesor jest używany jako cel sprzętowy.

Podczas tworzenia Tensora domyślne urządzenie w trybie gotowości można zastąpić, określając alternatywę. W ten sposób wyrażasz zgodę na wykonywanie obliczeń przy użyciu backendu 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)

Jeśli uruchamiasz to w instancji obsługującej procesor graficzny, powinieneś zobaczyć ten akcelerator na liście urządzenia tensora X10. W przeciwieństwie do szybkiego wykonania, jeśli uruchamiasz to w instancji z włączoną obsługą TPU, powinieneś teraz zobaczyć, że obliczenia korzystają z tego urządzenia. X10 to sposób na wykorzystanie TPU w Swift dla TensorFlow.

Domyślne urządzenia chętnie i X10 będą próbowały użyć pierwszego akceleratora w systemie. Jeśli masz podłączone procesory graficzne, użyje pierwszego dostępnego procesora graficznego. Jeśli obecne są TPU, X10 domyślnie użyje pierwszego rdzenia TPU. Jeśli nie zostanie znaleziony lub obsługiwany żaden akcelerator, urządzenie domyślne powróci do procesora.

Oprócz domyślnych urządzeń chętnych i XLA, możesz zapewnić określone cele sprzętowe i zaplecza w urządzeniu:

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

Trenowanie modelu w trybie chętnym

Przyjrzyjmy się, jak skonfigurować i wytrenować model przy użyciu domyślnego trybu wykonywania zachłannego. W tym przykładzie użyjemy prostego modelu LeNet-5 z repozytorium Swift-models i zbioru danych klasyfikacji cyfr pisanych odręcznie MNIST.

Najpierw skonfigurujemy i pobierzemy zbiór danych 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

Następnie skonfigurujemy model i optymalizator.

import ImageClassificationModels

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

Teraz wdrożymy podstawowe śledzenie postępów i raportowanie. Wszystkie statystyki pośrednie są przechowywane jako tensory na tym samym urządzeniu, na którym uruchamiane jest szkolenie, a scalarized() jest wywoływana tylko podczas raportowania. Będzie to szczególnie ważne później podczas używania X10, ponieważ pozwala uniknąć niepotrzebnej materializacji leniwych tensorów.

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
    }
}

Na koniec przeprowadzimy model przez pętlę treningową przez pięć epok.

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

Jak widać, model wytrenowany zgodnie z oczekiwaniami, a jego dokładność względem zbioru walidacyjnego rosła z każdą epoką. Oto jak modele Swift for TensorFlow są definiowane i uruchamiane przy użyciu szybkiego wykonywania, teraz zobaczmy, jakie modyfikacje należy wprowadzić, aby skorzystać z X10.

Trening modelu X10

Zbiory danych, modele i optymalizatory zawierają tensory, które są inicjowane na domyślnym urządzeniu wykonawczym zachłannym. Aby pracować z X10, będziemy musieli przenieść te tensory do urządzenia X10.

let device = Device.defaultXLA
print(device)
Device(kind: .CPU, ordinal: 0, backend: .XLA)

W przypadku zbiorów danych zrobimy to w momencie przetwarzania partii w pętli szkoleniowej, dzięki czemu będziemy mogli ponownie wykorzystać zbiór danych z modelu wykonania chętnie.

W przypadku modelu i optymalizatora zainicjujemy je ich wewnętrznymi tensorami na urządzeniu wykonawczym zachłannym, a następnie przeniesiemy je na urządzenie X10.

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

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

Modyfikacje potrzebne w pętli treningowej dotyczą kilku konkretnych punktów. Najpierw będziemy musieli przenieść partie danych treningowych na urządzenie X10. Odbywa się to za pomocą Tensor(copying:to:) po pobraniu każdej partii.

Następną zmianą jest wskazanie miejsca, w którym należy odciąć ślady podczas pętli treningowej. X10 działa poprzez śledzenie obliczeń tensora potrzebnych w kodzie i kompilowanie na czas zoptymalizowanej reprezentacji tego śladu. W przypadku pętli szkoleniowej powtarzasz tę samą operację w kółko, co jest idealną sekcją do śledzenia, kompilowania i ponownego wykorzystania.

W przypadku braku kodu, który jawnie żąda wartości od Tensora (zazwyczaj są to wywołania .scalars lub .scalarized() ), X10 spróbuje skompilować razem wszystkie iteracje pętli. Aby temu zapobiec i wyciąć ślad w określonym punkcie, umieszczamy jawną funkcję LazyTensorBarrier() po tym, jak optymalizator zaktualizuje wagi modelu oraz po uzyskaniu strat i dokładności podczas walidacji. Tworzy to dwa ponownie wykorzystane ślady: każdy krok w pętli szkoleniowej i każdą partię wnioskowania podczas walidacji.

Zmiany te skutkują następującą pętlą treningową.

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

Uczenie modelu przy użyciu backendu X10 powinno przebiegać w taki sam sposób, jak wcześniej zrobił to model wykonania zachłannego. Być może zauważyłeś opóźnienie przed pierwszą partią i pod koniec pierwszej epoki, spowodowane kompilacją unikalnych śladów w tych punktach na czas. Jeśli uruchamiasz to z podłączonym akceleratorem, powinieneś zauważyć, że trening po tym momencie przebiegał szybciej niż w trybie entuzjazmu.

Istnieje kompromis między początkowym czasem kompilacji śledzenia a większą przepływnością, ale w większości modeli uczenia maszynowego wzrost przepływności w wyniku powtarzanych operacji powinien z nawiązką zrównoważyć obciążenie kompilacji. W praktyce zaobserwowaliśmy ponad 4-krotną poprawę przepustowości dzięki X10 w niektórych przypadkach szkoleniowych.

Jak wspomniano wcześniej, użycie X10 sprawia, że ​​praca z TPU jest teraz nie tylko możliwa, ale i łatwa, odblokowując całą klasę akceleratorów dla modeli Swift dla TensorFlow.