%install '.package(url: "https://github.com/tensorflow/swift-models", .branch("tensorflow-0.12"))' Datasets ImageClassificationModels
print("\u{001B}[2J")
 Ver en TensorFlow.org |  Ejecutar en Google Colab |  Ver fuente en GitHub |
De forma predeterminada, Swift For TensorFlow realiza operaciones de tensor mediante despacho ansioso. Esto permite una iteración rápida, pero no es la opción más eficaz para entrenar modelos de aprendizaje automático.
La biblioteca de tensores X10 agrega un backend de alto rendimiento a Swift para TensorFlow, aprovechando el seguimiento de tensores y el compilador XLA . Este tutorial presentará X10 y lo guiará a través del proceso de actualización de un ciclo de entrenamiento para que se ejecute en GPU o TPU.
Tensores Eager vs. X10
Los cálculos acelerados en Swift for TensorFlow se realizan a través del tipo Tensor. Los tensores pueden participar en una amplia variedad de operaciones y son los componentes básicos de los modelos de aprendizaje automático.
De forma predeterminada, un Tensor utiliza una ejecución ansiosa para realizar cálculos operación por operación. Cada tensor tiene un dispositivo asociado que describe a qué hardware está conectado y qué backend se usa para él.
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 está ejecutando este portátil en una instancia habilitada para GPU, debería ver ese hardware reflejado en la descripción del dispositivo anterior. El tiempo de ejecución ansioso no tiene soporte para TPU, por lo que si está usando uno de ellos como acelerador, verá que la CPU se usa como objetivo de hardware.
Al crear un tensor, el dispositivo de modo entusiasta predeterminado se puede anular especificando una alternativa. Así es como opta por realizar cálculos utilizando el 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 está ejecutando esto en una instancia habilitada para GPU, debería ver ese acelerador en la lista del dispositivo del tensor X10. A diferencia de la ejecución ansiosa, si está ejecutando esto en una instancia habilitada para TPU, ahora debería ver que los cálculos están usando ese dispositivo. X10 es la forma en que aprovecha las TPU dentro de Swift para TensorFlow.
Los dispositivos ansiosos y X10 predeterminados intentarán usar el primer acelerador en el sistema. Si tiene GPU conectadas, usará la primera GPU disponible. Si hay TPU presentes, X10 utilizará el primer núcleo de TPU de forma predeterminada. Si no se encuentra o no se admite ningún acelerador, el dispositivo predeterminado recurrirá a la CPU.
Más allá de los dispositivos ansiosos y XLA predeterminados, puede proporcionar hardware específico y objetivos de back-end en un dispositivo:
// let tpu1 = Device(kind: .TPU, ordinal: 1, backend: .XLA)
// let tpuTensor1 = Tensor([0.0, 1.0, 2.0], on: tpu1)
Entrenando un modelo de modo entusiasta
Echemos un vistazo a cómo configuraría y entrenaría un modelo utilizando el modo de ejecución impaciente predeterminado. En este ejemplo, usaremos el modelo LeNet-5 simple del repositorio de modelos rápidos y el conjunto de datos de clasificación de dígitos escritos a mano del MNIST.
Primero, configuraremos y descargaremos el conjunto de datos 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
A continuación, configuraremos el modelo y el optimizador.
import ImageClassificationModels
var eagerModel = LeNet()
var eagerOptimizer = SGD(for: eagerModel, learningRate: 0.1)
Ahora, implementaremos el seguimiento y la generación de informes de progreso básicos. Todas las estadísticas intermedias se mantienen como tensores en el mismo dispositivo donde se ejecuta el entrenamiento y se llama a scalarized() solo durante la generación de informes. Esto será especialmente importante más adelante cuando use X10, porque evita la materialización innecesaria de tensores perezosos.
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
}
}
Finalmente, ejecutaremos el modelo a través de un ciclo de entrenamiento durante cinco épocas.
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
Como puede ver, el modelo se entrenó como cabría esperar y su precisión frente al conjunto de validación aumentó en cada época. Así es como se definen y ejecutan los modelos de Swift for TensorFlow con una ejecución entusiasta, ahora veamos qué modificaciones se deben realizar para aprovechar X10.
Entrenamiento de un modelo X10
Los conjuntos de datos, modelos y optimizadores contienen tensores que se inicializan en el dispositivo de ejecución impaciente predeterminado. Para trabajar con X10, necesitaremos mover estos tensores a un dispositivo X10.
let device = Device.defaultXLA
print(device)
Device(kind: .CPU, ordinal: 0, backend: .XLA)
Para los conjuntos de datos, lo haremos en el punto en el que se procesan los lotes en el ciclo de entrenamiento, de modo que podamos reutilizar el conjunto de datos del modelo de ejecución entusiasta.
En el caso del modelo y el optimizador, los inicializaremos con sus tensores internos en el dispositivo de ejecución ansioso y luego los moveremos al dispositivo X10.
var x10Model = LeNet()
x10Model.move(to: device)
var x10Optimizer = SGD(for: x10Model, learningRate: 0.1)
x10Optimizer = SGD(copying: x10Optimizer, to: device)
Las modificaciones necesarias para el ciclo de entrenamiento vienen en algunos puntos específicos. Primero, necesitaremos mover los lotes de datos de entrenamiento al dispositivo X10. Esto se hace a través de Tensor(copying:to:) cuando se recupera cada lote.
El siguiente cambio es indicar dónde cortar los rastros durante el ciclo de entrenamiento. X10 funciona rastreando los cálculos de tensor necesarios en su código y compilando justo a tiempo una representación optimizada de ese rastreo. En el caso de un ciclo de entrenamiento, está repitiendo la misma operación una y otra vez, una sección ideal para rastrear, compilar y reutilizar.
En ausencia de un código que solicite explícitamente un valor de un Tensor (estos generalmente se destacan como .scalars o .scalarized() ), X10 intentará compilar todas las iteraciones de bucle juntas. Para evitar esto y cortar el seguimiento en un punto específico, colocamos una LazyTensorBarrier() explícita después de que el optimizador actualice los pesos del modelo y después de que se obtengan la pérdida y la precisión durante la validación. Esto crea dos rastros reutilizados: cada paso en el ciclo de entrenamiento y cada lote de inferencia durante la validación.
Estas modificaciones dan como resultado el siguiente ciclo de entrenamiento.
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
El entrenamiento del modelo con el backend X10 debería haber procedido de la misma manera que lo hizo antes el modelo de ejecución ansiosa. Es posible que haya notado un retraso antes del primer lote y al final de la primera época, debido a la compilación justo a tiempo de las trazas únicas en esos puntos. Si está ejecutando esto con un acelerador adjunto, debería haber visto que el entrenamiento después de ese punto avanza más rápido que con el modo ansioso.
Existe una compensación entre el tiempo de compilación de seguimiento inicial y un rendimiento más rápido, pero en la mayoría de los modelos de aprendizaje automático, el aumento en el rendimiento de las operaciones repetidas debería compensar con creces la sobrecarga de compilación. En la práctica, hemos visto una mejora de más de 4 veces en el rendimiento con X10 en algunos casos de capacitación.
Como se mencionó anteriormente, el uso de X10 ahora hace que no solo sea posible, sino también fácil trabajar con TPU, desbloqueando toda esa clase de aceleradores para sus modelos de Swift for TensorFlow.