%install '.package(url: "https://github.com/tensorflow/swift-models", .branch("tensorflow-0.12"))' Datasets ImageClassificationModels
print("\u{001B}[2J")
 Veja no TensorFlow.org |  Ver fonte no GitHub |
Por padrão, o Swift For TensorFlow executa operações de tensor usando despacho antecipado. Isso permite uma iteração rápida, mas não é a opção de melhor desempenho para treinar modelos de aprendizado de máquina.
A biblioteca de tensores X10 adiciona um back-end de alto desempenho ao Swift for TensorFlow, aproveitando o rastreamento de tensor e o compilador XLA . Este tutorial apresentará o X10 e o guiará pelo processo de atualização de um loop de treinamento para execução em GPUs ou TPUs.
Tensores ansiosos vs. X10
Os cálculos acelerados no Swift for TensorFlow são realizados por meio do tipo Tensor. Os tensores podem participar de uma ampla variedade de operações e são os blocos de construção fundamentais dos modelos de aprendizado de máquina.
Por padrão, um tensor usa a execução antecipada para realizar cálculos operação por operação. Cada tensor tem um dispositivo associado que descreve a qual hardware ele está conectado e qual back-end é usado para ele.
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 você estiver executando este notebook em uma instância habilitada para GPU, deverá ver esse hardware refletido na descrição do dispositivo acima. O tempo de execução ansioso não tem suporte para TPUs, portanto, se você estiver usando um deles como acelerador, verá a CPU sendo usada como destino de hardware.
Ao criar um tensor, o dispositivo de modo ansioso padrão pode ser substituído especificando uma alternativa. É assim que você opta por realizar cálculos usando o back-end 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 você estiver executando isso em uma instância habilitada para GPU, deverá ver esse acelerador listado no dispositivo do tensor X10. Ao contrário da execução antecipada, se você estiver executando isso em uma instância habilitada para TPU, deverá ver agora que os cálculos estão usando esse dispositivo. O X10 é como você aproveita as TPUs no Swift for TensorFlow.
Os dispositivos ansiosos e X10 padrão tentarão usar o primeiro acelerador no sistema. Se você tiver GPUs conectadas, usará a primeira GPU disponível. Se houver TPUs, o X10 usará o primeiro núcleo de TPU por padrão. Se nenhum acelerador for encontrado ou suportado, o dispositivo padrão retornará à CPU.
Além dos dispositivos ansiosos e XLA padrão, você pode fornecer alvos específicos de hardware e back-end em um dispositivo:
// let tpu1 = Device(kind: .TPU, ordinal: 1, backend: .XLA)
// let tpuTensor1 = Tensor([0.0, 1.0, 2.0], on: tpu1)
Treinando um modelo de modo ansioso
Vamos dar uma olhada em como você configuraria e treinaria um modelo usando o modo de execução antecipada padrão. Neste exemplo, usaremos o modelo LeNet-5 simples do repositório swift-models e o conjunto de dados de classificação de dígitos manuscritos MNIST.
Primeiro, vamos configurar e baixar o conjunto de dados 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
Em seguida, configuraremos o modelo e o otimizador.
import ImageClassificationModels
var eagerModel = LeNet()
var eagerOptimizer = SGD(for: eagerModel, learningRate: 0.1)
Agora, implementaremos o rastreamento e relatórios básicos de progresso. Todas as estatísticas intermediárias são mantidas como tensores no mesmo dispositivo em que o treinamento é executado e scalarized() é chamado apenas durante o relatório. Isso será especialmente importante mais tarde ao usar o X10, porque evita a materialização desnecessária de tensores preguiçosos.
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
}
}
Por fim, executaremos o modelo por meio de um loop de treinamento por 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 você pode ver, o modelo treinado como esperávamos e sua precisão em relação ao conjunto de validação aumentou a cada época. É assim que os modelos Swift for TensorFlow são definidos e executados usando execução antecipada, agora vamos ver quais modificações precisam ser feitas para aproveitar o X10.
Treinando um modelo X10
Conjuntos de dados, modelos e otimizadores contêm tensores que são inicializados no dispositivo de execução antecipada padrão. Para trabalhar com o X10, precisaremos mover esses tensores para um dispositivo X10.
let device = Device.defaultXLA
print(device)
Device(kind: .CPU, ordinal: 0, backend: .XLA)
Para os conjuntos de dados, faremos isso no ponto em que os lotes são processados no loop de treinamento, para que possamos reutilizar o conjunto de dados do modelo de execução antecipado.
No caso do modelo e do otimizador, vamos inicializá-los com seus tensores internos no dispositivo de execução antecipada e depois movê-los para o dispositivo X10.
var x10Model = LeNet()
x10Model.move(to: device)
var x10Optimizer = SGD(for: x10Model, learningRate: 0.1)
x10Optimizer = SGD(copying: x10Optimizer, to: device)
As modificações necessárias para o loop de treinamento ocorrem em alguns pontos específicos. Primeiro, precisaremos mover os lotes de dados de treinamento para o dispositivo X10. Isso é feito via Tensor(copying:to:) quando cada lote é recuperado.
A próxima mudança é indicar onde cortar os traços durante o loop de treinamento. O X10 funciona rastreando os cálculos de tensor necessários em seu código e compilando just-in-time uma representação otimizada desse rastreamento. No caso de um loop de treinamento, você está repetindo a mesma operação várias vezes, uma seção ideal para rastrear, compilar e reutilizar.
Na ausência de código que solicite explicitamente um valor de um Tensor (estes geralmente se destacam como .scalars ou .scalarized() ), o X10 tentará compilar todas as iterações de loop juntas. Para evitar isso e cortar o rastreamento em um ponto específico, colocamos um LazyTensorBarrier() explícito após o otimizador atualizar os pesos do modelo e após a perda e a precisão serem obtidas durante a validação. Isso cria dois rastreamentos reutilizados: cada etapa no loop de treinamento e cada lote de inferência durante a validação.
Essas modificações resultam no seguinte loop de treinamento.
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
O treinamento do modelo usando o backend X10 deveria ter prosseguido da mesma maneira que o modelo de execução antecipada fazia antes. Você pode ter notado um atraso antes do primeiro lote e no final da primeira época, devido à compilação just-in-time dos traços únicos nesses pontos. Se você estiver executando isso com um acelerador conectado, deve ter visto o treinamento após esse ponto prosseguir mais rápido do que no modo ansioso.
Há uma compensação entre o tempo de compilação de rastreamento inicial versus a taxa de transferência mais rápida, mas na maioria dos modelos de aprendizado de máquina o aumento na taxa de transferência de operações repetidas deve mais do que compensar a sobrecarga de compilação. Na prática, vimos uma melhoria de mais de 4 vezes no rendimento com o X10 em alguns casos de treinamento.
Como foi dito antes, o uso do X10 agora torna não apenas possível, mas fácil de trabalhar com TPUs, desbloqueando toda essa classe de aceleradores para seus modelos Swift for TensorFlow.