%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 código fuente en GitHub |
De forma predeterminada, Swift For TensorFlow realiza operaciones tensoriales 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 bucle de entrenamiento para ejecutarlo en GPU o TPU.
Tensores ansiosos frente a X10
Los cálculos acelerados en Swift para TensorFlow se realizan mediante el tipo Tensor. Los tensores pueden participar en una amplia variedad de operaciones y son los componentes fundamentales de los modelos de aprendizaje automático.
De forma predeterminada, un tensor utiliza una ejecución entusiasta 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 utiliza para ello.
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 esta computadora 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 es compatible con TPU, por lo que si está utilizando uno de ellos como acelerador, verá que la CPU se utiliza como objetivo de hardware.
Al crear un tensor, el dispositivo de modo ansioso predeterminado se puede anular especificando una alternativa. Así es como usted opta por realizar cálculos utilizando el backend de 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 utilizando ese dispositivo. X10 es la forma de aprovechar las TPU dentro de Swift para TensorFlow.
Los dispositivos ansiosos y X10 predeterminados intentarán utilizar el primer acelerador del sistema. Si tiene GPU conectadas, utilizará la primera GPU disponible. Si hay TPU presentes, X10 utilizará el primer núcleo de TPU de forma predeterminada. Si no se encuentra ni 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 backend 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 en modo ansioso
Echemos un vistazo a cómo configurarías y entrenarías un modelo usando el modo de ejecución entusiasta predeterminado. En este ejemplo, usaremos el modelo LeNet-5 simple del repositorio de modelos swift y el conjunto de datos de clasificación de dígitos escritos a mano 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 informes y seguimiento de progreso básico. 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 utilice 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 era de esperar y su precisión frente al conjunto de validación aumentó en cada época. Así es como se definen y ejecutan los modelos Swift para TensorFlow mediante ejecución ansiosa, ahora veamos qué modificaciones se deben realizar para aprovechar X10.
Entrenando un modelo X10
Los conjuntos de datos, modelos y optimizadores contienen tensores que se inicializan en el dispositivo de ejecución ansioso 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, para 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 trasladaremos 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 se producen en algunos puntos específicos. Primero, necesitaremos mover los lotes de datos de entrenamiento al dispositivo X10. Esto se hace mediante 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 tensoriales necesarios en su código y compilando justo a tiempo una representación optimizada de ese rastro. En el caso de un bucle de entrenamiento, estás repitiendo la misma operación una y otra vez, una sección ideal para rastrear, compilar y reutilizar.
En ausencia de código que solicite explícitamente un valor de un tensor (estos generalmente se destacan como llamadas .scalars o .scalarized() ), X10 intentará compilar todas las iteraciones del bucle juntas. Para evitar esto y cortar el seguimiento en un punto específico, colocamos un LazyTensorBarrier() explícito 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 del 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 utilizando el backend X10 debería haberse realizado de la misma manera que lo hizo antes el modelo de ejecución entusiasta. 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 los rastros únicos en esos puntos. Si estás ejecutando esto con un acelerador conectado, deberías 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 del seguimiento inicial y un rendimiento más rápido, pero en la mayoría de los modelos de aprendizaje automático el aumento del rendimiento debido a 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 indicó 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 Swift para TensorFlow.