%install '.package(url: "https://github.com/tensorflow/swift-models", .branch("tensorflow-0.12"))' Datasets ImageClassificationModels
print("\u{001B}[2J")
 Ver em TensorFlow.org |  Execute no Google Colab |  Ver fonte no GitHub |
Por padrão, o Swift For TensorFlow executa operações de tensor usando despacho rápido. 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 para TensorFlow, aproveitando o rastreamento de tensor e o compilador XLA . Este tutorial apresentará o X10 e guiará você pelo processo de atualização de um loop de treinamento para execução em GPUs ou TPUs.
Tensores ansiosos vs. X10
Cálculos acelerados em Swift para 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 execução antecipada para realizar cálculos operação por operação. Cada Tensor possui 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 alvo 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. X10 é como você aproveita as vantagens das TPUs no Swift para TensorFlow.
Os dispositivos padrão ansioso e X10 tentarão usar o primeiro acelerador do sistema. Se você tiver GPUs conectadas, usará a primeira GPU disponível. Se TPUs estiverem presentes, o X10 usará o primeiro núcleo de TPU por padrão. Se nenhum acelerador for encontrado ou compatível, o dispositivo padrão retornará à CPU.
Além dos dispositivos ansiosos e XLA padrão, você pode fornecer hardware específico e destinos de 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 rápida 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, configuraremos e baixaremos 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
A seguir, configuraremos o modelo e o otimizador.
import ImageClassificationModels
var eagerModel = LeNet()
var eagerOptimizer = SGD(for: eagerModel, learningRate: 0.1)
Agora, implementaremos monitoramento e relatórios básicos de progresso. Todas as estatísticas intermediárias são mantidas como tensores no mesmo dispositivo onde o treinamento é executado e scalarized() é chamado apenas durante o relatório. Isto será especialmente importante posteriormente 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 em 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 foi treinado conforme esperado e sua precisão em relação ao conjunto de validação aumentou a cada época. É assim que os modelos Swift para TensorFlow são definidos e executados usando execução rápida. Agora vamos ver quais modificações precisam ser feitas para aproveitar as vantagens do X10.
Treinando um modelo X10
Conjuntos de dados, modelos e otimizadores contêm tensores que são inicializados no dispositivo de execução rápida 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 antecipada.
No caso do modelo e do otimizador, iremos inicializá-los com seus tensores internos no dispositivo de execução rápida e, em seguida, 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 ciclo 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 ciclo 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ê repete a mesma operação continuamente, uma seção ideal para rastrear, compilar e reutilizar.
Na ausência de código que solicite explicitamente um valor de um Tensor (geralmente se destacam como chamadas .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 do loop de treinamento e cada lote de inferência durante a validação.
Essas modificações resultam no seguinte ciclo 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 back-end X10 deveria ter ocorrido da mesma maneira que o modelo de execução antecipada fez antes. Você deve 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, deverá ter visto o treinamento depois desse ponto prosseguir mais rápido do que no modo ansioso.
Há uma compensação entre o tempo de compilação de rastreamento inicial e 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 4X no rendimento com o X10 em alguns casos de treinamento.
Como foi afirmado antes, o uso do X10 agora torna não apenas possível, mas fácil, trabalhar com TPUs, desbloqueando toda essa classe de aceleradores para seus modelos Swift para TensorFlow.