Presentazione dell'X10

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


Visualizza su TensorFlow.org Visualizza la fonte su GitHub

Per impostazione predefinita, Swift For TensorFlow esegue operazioni tensoriali utilizzando l'invio desideroso. Ciò consente un'iterazione rapida, ma non è l'opzione più performante per l'addestramento dei modelli di machine learning.

La libreria tensore X10 aggiunge un backend ad alte prestazioni a Swift per TensorFlow, sfruttando la tracciatura del tensore e il compilatore XLA . Questo tutorial introdurrà X10 e ti guiderà attraverso il processo di aggiornamento di un ciclo di training da eseguire su GPU o TPU.

Tensori desiderosi contro X10

I calcoli accelerati in Swift per TensorFlow vengono eseguiti tramite il tipo Tensor. I tensori possono partecipare a un'ampia varietà di operazioni e sono gli elementi fondamentali dei modelli di apprendimento automatico.

Per impostazione predefinita, un tensore utilizza l'esecuzione entusiasta per eseguire calcoli operazione per operazione. Ogni tensore ha un dispositivo associato che descrive a quale hardware è collegato e quale backend viene utilizzato per esso.

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)

Se stai eseguendo questo notebook su un'istanza abilitata per GPU, dovresti vedere l'hardware riflesso nella descrizione del dispositivo sopra. Il runtime Entusiast non supporta i TPU, quindi se ne utilizzi uno come acceleratore vedrai la CPU utilizzata come destinazione hardware.

Quando si crea un tensore, il dispositivo in modalità desideroso predefinito può essere sovrascritto specificando un'alternativa. In questo modo puoi attivare l'esecuzione dei calcoli utilizzando il 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)

Se lo stai eseguendo in un'istanza abilitata per GPU, dovresti vedere l'acceleratore elencato nel dispositivo del tensore X10. A differenza dell'esecuzione impaziente, se la esegui in un'istanza abilitata per TPU, ora dovresti vedere che i calcoli utilizzano quel dispositivo. X10 è il modo in cui sfrutti i TPU all'interno di Swift per TensorFlow.

I dispositivi desiderosi e X10 predefiniti tenteranno di utilizzare il primo acceleratore sul sistema. Se hai GPU collegate, utilizzerà la prima GPU disponibile. Se sono presenti TPU, X10 utilizzerà il primo core TPU per impostazione predefinita. Se non viene trovato o supportato alcun acceleratore, il dispositivo predefinito tornerà alla CPU.

Oltre ai dispositivi desiderosi e XLA predefiniti, puoi fornire hardware e target backend specifici in un dispositivo:

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

Addestramento di un modello in modalità desideroso

Diamo un'occhiata a come configureresti e addestreresti un modello utilizzando la modalità di esecuzione impaziente predefinita. In questo esempio, utilizzeremo il semplice modello LeNet-5 dal repository dei modelli swift e il set di dati di classificazione delle cifre scritte a mano MNIST.

Innanzitutto, configureremo e scaricheremo il set di dati 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

Successivamente, configureremo il modello e l'ottimizzatore.

import ImageClassificationModels

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

Ora implementeremo il monitoraggio e il reporting di base dei progressi. Tutte le statistiche intermedie vengono conservate come tensori sullo stesso dispositivo su cui viene eseguito l'addestramento e scalarized() viene chiamato solo durante il reporting. Ciò sarà particolarmente importante in seguito quando si utilizzerà X10, poiché evita la materializzazione non necessaria di tensori pigri.

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

Infine, eseguiremo il modello attraverso un ciclo di addestramento per cinque epoche.

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

Come puoi vedere, il modello si è addestrato come previsto e la sua precisione rispetto al set di convalida è aumentata ad ogni epoca. Ecco come vengono definiti ed eseguiti i modelli Swift per TensorFlow utilizzando l'esecuzione entusiasta, ora vediamo quali modifiche è necessario apportare per sfruttare X10.

Addestramento di un modello X10

I set di dati, i modelli e gli ottimizzatori contengono tensori inizializzati sul dispositivo di esecuzione impaziente predefinito. Per funzionare con X10, dovremo spostare questi tensori su un dispositivo X10.

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

Per i set di dati, lo faremo nel punto in cui i batch vengono elaborati nel ciclo di addestramento, in modo da poter riutilizzare il set di dati dal modello di esecuzione impaziente.

Nel caso del modello e dell'ottimizzatore, li inizializzeremo con i relativi tensori interni sul dispositivo di esecuzione impaziente, quindi li sposteremo sul dispositivo X10.

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

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

Le modifiche necessarie per il ciclo di formazione arrivano in alcuni punti specifici. Per prima cosa dovremo spostare i batch di dati di addestramento sul dispositivo X10. Questa operazione viene eseguita tramite Tensor(copying:to:) quando viene recuperato ciascun batch.

La modifica successiva consiste nell'indicare dove tagliare le tracce durante il ciclo di allenamento. X10 funziona tracciando i calcoli tensoriali necessari nel codice e compilando just-in-time una rappresentazione ottimizzata di quella traccia. Nel caso di un ciclo di addestramento, stai ripetendo la stessa operazione più e più volte, una sezione ideale da tracciare, compilare e riutilizzare.

In assenza di codice che richiede esplicitamente un valore da un tensore (di solito si distinguono come chiamate .scalars o .scalarized() ), X10 tenterà di compilare insieme tutte le iterazioni del ciclo. Per evitare ciò e tagliare la traccia in un punto specifico, inseriamo un LazyTensorBarrier() esplicito dopo che l'ottimizzatore ha aggiornato i pesi del modello e dopo che la perdita e l'accuratezza sono state ottenute durante la convalida. Ciò crea due tracce riutilizzate: ogni passaggio nel ciclo di training e ogni batch di inferenza durante la convalida.

Queste modifiche danno luogo al seguente ciclo di addestramento.

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

L'addestramento del modello utilizzando il backend X10 avrebbe dovuto procedere nello stesso modo del precedente modello di esecuzione impaziente. Potresti aver notato un ritardo prima del primo batch e alla fine della prima epoca, dovuto alla compilazione just-in-time delle tracce uniche in quei punti. Se lo stai eseguendo con un acceleratore collegato, dovresti aver visto l'allenamento da quel momento in poi procedere più velocemente che con la modalità desiderosa.

Esiste un compromesso tra il tempo di compilazione della traccia iniziale e una velocità effettiva più rapida, ma nella maggior parte dei modelli di machine learning l'aumento della velocità effettiva derivante da operazioni ripetute dovrebbe più che compensare il sovraccarico di compilazione. In pratica, in alcuni casi di formazione abbiamo riscontrato un miglioramento di oltre 4 volte nel throughput con X10.

Come affermato in precedenza, l'utilizzo di X10 ora rende non solo possibile ma anche facile lavorare con le TPU, sbloccando l'intera classe di acceleratori per i tuoi modelli Swift per TensorFlow.