Présentation de X10

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


Voir sur TensorFlow.org Exécuter dans Google Colab Afficher la source sur GitHub

Par défaut, Swift For TensorFlow effectue des opérations tensorielles à l'aide d'une répartition hâtive. Cela permet une itération rapide, mais ne constitue pas l’option la plus performante pour former des modèles d’apprentissage automatique.

La bibliothèque de tenseurs X10 ajoute un backend hautes performances à Swift pour TensorFlow, en tirant parti du traçage de tenseurs et du compilateur XLA . Ce didacticiel présentera X10 et vous guidera tout au long du processus de mise à jour d'une boucle de formation à exécuter sur des GPU ou des TPU.

Tenseurs impatients vs X10

Les calculs accélérés dans Swift pour TensorFlow sont effectués via le type Tensor. Les tenseurs peuvent participer à une grande variété d’opérations et constituent les éléments fondamentaux des modèles d’apprentissage automatique.

Par défaut, un Tensor utilise une exécution rapide pour effectuer des calculs opération par opération. Chaque Tensor est associé à un périphérique qui décrit le matériel auquel il est connecté et le backend utilisé.

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)

Si vous exécutez ce notebook sur une instance compatible GPU, vous devriez voir ce matériel reflété dans la description de l'appareil ci-dessus. Le moteur d'exécution impatient ne prend pas en charge les TPU, donc si vous utilisez l'un d'entre eux comme accélérateur, vous verrez le processeur utilisé comme cible matérielle.

Lors de la création d'un Tensor, le périphérique en mode impatient par défaut peut être remplacé en spécifiant une alternative. C'est ainsi que vous choisissez d'effectuer des calculs à l'aide du backend 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)

Si vous l'exécutez dans une instance compatible GPU, vous devriez voir cet accélérateur répertorié dans le périphérique du tenseur X10. Contrairement à une exécution rapide, si vous l'exécutez dans une instance compatible TPU, vous devriez maintenant voir que les calculs utilisent cet appareil. X10 vous permet de profiter des TPU dans Swift pour TensorFlow.

Les appareils impatients et X10 par défaut tenteront d'utiliser le premier accélérateur du système. Si des GPU sont connectés, ils utiliseront le premier GPU disponible. Si des TPU sont présents, X10 utilisera le premier cœur TPU par défaut. Si aucun accélérateur n'est trouvé ou pris en charge, le périphérique par défaut reviendra au processeur.

Au-delà des appareils impatients et XLA par défaut, vous pouvez fournir des cibles matérielles et back-end spécifiques dans un appareil :

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

Entraîner un modèle en mode impatient

Voyons comment configurer et entraîner un modèle à l'aide du mode d'exécution rapide par défaut. Dans cet exemple, nous utiliserons le modèle simple LeNet-5 du référentiel Swift-Models et l'ensemble de données de classification des chiffres manuscrits du MNIST.

Tout d’abord, nous allons configurer et télécharger l’ensemble de données 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

Ensuite, nous configurerons le modèle et l'optimiseur.

import ImageClassificationModels

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

Nous allons désormais mettre en œuvre un suivi et des rapports de base sur les progrès. Toutes les statistiques intermédiaires sont conservées sous forme de tenseurs sur le même appareil sur lequel la formation est exécutée et scalarized() n'est appelé que lors du reporting. Cela sera particulièrement important plus tard lors de l'utilisation de X10, car cela évite la matérialisation inutile de tenseurs paresseux.

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

Enfin, nous exécuterons le modèle dans une boucle de formation pendant cinq époques.

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

Comme vous pouvez le constater, le modèle s'est entraîné comme prévu et sa précision par rapport à l'ensemble de validation a augmenté à chaque époque. C'est ainsi que les modèles Swift pour TensorFlow sont définis et exécutés en utilisant une exécution rapide. Voyons maintenant quelles modifications doivent être apportées pour tirer parti de X10.

Entraîner un modèle X10

Les ensembles de données, les modèles et les optimiseurs contiennent des tenseurs qui sont initialisés sur le périphérique d'exécution rapide par défaut. Pour travailler avec X10, nous devrons déplacer ces tenseurs vers un appareil X10.

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

Pour les ensembles de données, nous le ferons au moment où les lots sont traités dans la boucle de formation, afin de pouvoir réutiliser l'ensemble de données du modèle d'exécution impatient.

Dans le cas du modèle et de l'optimiseur, nous les initialiserons avec leurs tenseurs internes sur le périphérique d'exécution impatient, puis les déplacerons vers le périphérique X10.

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

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

Les modifications nécessaires à la boucle d'entraînement interviennent à quelques points précis. Tout d’abord, nous devrons déplacer les lots de données d’entraînement vers l’appareil X10. Cela se fait via Tensor(copying:to:) lorsque chaque lot est récupéré.

Le changement suivant est d'indiquer où couper les traces lors de la boucle d'entraînement. X10 fonctionne en traçant les calculs tensoriels nécessaires dans votre code et en compilant juste à temps une représentation optimisée de cette trace. Dans le cas d'une boucle d'entraînement, vous répétez encore et encore la même opération, une section idéale à tracer, compiler et réutiliser.

En l'absence de code qui demande explicitement une valeur à un Tensor (ceux-ci se distinguent généralement par des appels .scalars ou .scalarized() ), X10 tentera de compiler toutes les itérations de boucle ensemble. Pour éviter cela et couper la trace à un point spécifique, nous plaçons un LazyTensorBarrier() explicite après que l'optimiseur ait mis à jour les poids du modèle et après que la perte et la précision aient été obtenues lors de la validation. Cela crée deux traces réutilisées : chaque étape de la boucle d'entraînement et chaque lot d'inférence lors de la validation.

Ces modifications aboutissent à la boucle d’entraînement suivante.

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

La formation du modèle à l'aide du backend X10 aurait dû se dérouler de la même manière que le modèle d'exécution impatient le faisait auparavant. Vous avez peut-être remarqué un retard avant le premier lot et à la fin de la première époque, dû à la compilation juste à temps des traces uniques à ces points. Si vous exécutez ceci avec un accélérateur connecté, vous devriez avoir vu l'entraînement se dérouler plus rapidement après ce point qu'avec le mode impatient.

Il existe un compromis entre le temps de compilation initial des traces et un débit plus rapide, mais dans la plupart des modèles d'apprentissage automatique, l'augmentation du débit due aux opérations répétées devrait plus que compenser la surcharge de compilation. En pratique, nous avons constaté une amélioration du débit plus de 4 fois avec X10 dans certains cas de formation.

Comme cela a été indiqué précédemment, l'utilisation de X10 rend désormais non seulement possible mais facile le travail avec les TPU, déverrouillant ainsi toute cette classe d'accélérateurs pour vos modèles Swift pour TensorFlow.