Memperkenalkan X10

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


Lihat di TensorFlow.org Jalankan di Google Colab Lihat sumber di GitHub

Secara default, Swift For TensorFlow melakukan operasi tensor menggunakan pengiriman bersemangat. Hal ini memungkinkan terjadinya iterasi yang cepat, namun bukan merupakan pilihan yang paling efektif untuk melatih model pembelajaran mesin.

Pustaka tensor X10 menambahkan backend berperforma tinggi ke Swift untuk TensorFlow, memanfaatkan penelusuran tensor dan kompiler XLA . Tutorial ini akan memperkenalkan X10 dan memandu Anda melalui proses memperbarui loop pelatihan untuk dijalankan pada GPU atau TPU.

Tensor Eager vs. X10

Penghitungan yang dipercepat di Swift untuk TensorFlow dilakukan melalui tipe Tensor. Tensor dapat berpartisipasi dalam berbagai macam operasi, dan merupakan elemen dasar model pembelajaran mesin.

Secara default, Tensor menggunakan eksekusi yang bersemangat untuk melakukan penghitungan berdasarkan operasi demi operasi. Setiap Tensor memiliki Perangkat terkait yang menjelaskan perangkat keras apa yang terpasang dan backend apa yang digunakan untuk Tensor tersebut.

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)

Jika Anda menjalankan notebook ini pada instance yang mendukung GPU, Anda akan melihat perangkat keras tersebut tercermin dalam deskripsi perangkat di atas. Runtime yang bersemangat tidak memiliki dukungan untuk TPU, jadi jika Anda menggunakan salah satunya sebagai akselerator, Anda akan melihat CPU digunakan sebagai target perangkat keras.

Saat membuat Tensor, perangkat mode bersemangat default dapat diganti dengan menentukan alternatif. Ini adalah cara Anda memilih untuk melakukan penghitungan menggunakan 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)

Jika Anda menjalankan ini dalam instance berkemampuan GPU, Anda akan melihat akselerator tersebut tercantum di perangkat tensor X10. Berbeda dengan eksekusi bersemangat, jika Anda menjalankan ini dalam instance berkemampuan TPU, Anda sekarang akan melihat bahwa penghitungan menggunakan perangkat tersebut. X10 adalah cara Anda memanfaatkan TPU dalam Swift untuk TensorFlow.

Perangkat bersemangat dan X10 default akan mencoba menggunakan akselerator pertama pada sistem. Jika Anda memiliki GPU yang terpasang, maka akan menggunakan GPU pertama yang tersedia. Jika ada TPU, X10 akan menggunakan inti TPU pertama secara default. Jika tidak ada akselerator yang ditemukan atau didukung, perangkat default akan kembali ke CPU.

Selain perangkat bersemangat dan XLA default, Anda dapat menyediakan target perangkat keras dan backend tertentu di Perangkat:

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

Melatih model mode bersemangat

Mari kita lihat cara Anda menyiapkan dan melatih model menggunakan mode eksekusi bersemangat default. Dalam contoh ini, kita akan menggunakan model LeNet-5 sederhana dari repositori model cepat dan kumpulan data klasifikasi digit tulisan tangan MNIST.

Pertama, kita akan menyiapkan dan mendownload dataset 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

Selanjutnya, kita akan mengkonfigurasi model dan pengoptimal.

import ImageClassificationModels

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

Sekarang, kami akan menerapkan pelacakan dan pelaporan kemajuan dasar. Semua statistik perantara disimpan sebagai tensor pada perangkat yang sama tempat pelatihan dijalankan dan scalarized() dipanggil hanya selama pelaporan. Hal ini akan menjadi sangat penting nantinya ketika menggunakan X10, karena ini menghindari perwujudan tensor malas yang tidak diperlukan.

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

Terakhir, kita akan menjalankan model melalui loop pelatihan selama lima epoch.

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

Seperti yang Anda lihat, model dilatih seperti yang kita harapkan, dan keakuratannya terhadap set validasi meningkat setiap periode. Beginilah cara model Swift untuk TensorFlow didefinisikan dan dijalankan menggunakan eksekusi yang bersemangat, sekarang mari kita lihat modifikasi apa yang perlu dilakukan untuk memanfaatkan X10.

Melatih model X10

Kumpulan data, model, dan pengoptimal berisi tensor yang diinisialisasi pada perangkat eksekusi bersemangat default. Untuk bekerja dengan X10, kita perlu memindahkan tensor ini ke perangkat X10.

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

Untuk kumpulan data, kita akan melakukannya pada titik di mana kumpulan diproses dalam loop pelatihan, sehingga kita dapat menggunakan kembali kumpulan data dari model eksekusi yang bersemangat.

Dalam hal model dan pengoptimal, kami akan menginisialisasinya dengan tensor internalnya pada perangkat eksekusi yang bersemangat, lalu memindahkannya ke perangkat X10.

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

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

Modifikasi yang diperlukan untuk loop pelatihan terjadi pada beberapa poin tertentu. Pertama, kita perlu memindahkan kumpulan data pelatihan ke perangkat X10. Hal ini dilakukan melalui Tensor(copying:to:) saat setiap batch diambil.

Perubahan selanjutnya adalah menunjukkan di mana harus memotong jejak selama loop pelatihan. X10 bekerja dengan menelusuri penghitungan tensor yang diperlukan dalam kode Anda dan mengkompilasi representasi pelacakan tersebut secara tepat waktu. Dalam kasus loop pelatihan, Anda mengulangi operasi yang sama berulang kali, bagian yang ideal untuk dilacak, dikompilasi, dan digunakan kembali.

Jika tidak ada kode yang secara eksplisit meminta nilai dari Tensor (biasanya disebut sebagai panggilan .scalars atau .scalarized() ), X10 akan mencoba mengkompilasi semua iterasi loop secara bersamaan. Untuk mencegah hal ini, dan memotong jejak pada titik tertentu, kami menempatkan LazyTensorBarrier() eksplisit setelah pengoptimal memperbarui bobot model dan setelah kehilangan dan akurasi diperoleh selama validasi. Hal ini menciptakan dua jejak yang digunakan kembali: setiap langkah dalam loop pelatihan dan setiap kumpulan inferensi selama validasi.

Modifikasi ini menghasilkan loop pelatihan berikut.

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

Pelatihan model yang menggunakan backend X10 seharusnya berjalan dengan cara yang sama seperti yang dilakukan model eksekusi bersemangat sebelumnya. Anda mungkin telah melihat penundaan sebelum batch pertama dan pada akhir epoch pertama, karena kompilasi jejak unik yang tepat waktu pada titik-titik tersebut. Jika Anda menjalankan ini dengan akselerator terpasang, Anda seharusnya melihat pelatihan setelah titik tersebut berjalan lebih cepat dibandingkan dengan mode bersemangat.

Ada trade-off antara waktu kompilasi jejak awal vs. throughput yang lebih cepat, namun di sebagian besar model pembelajaran mesin, peningkatan throughput dari operasi berulang seharusnya lebih dari sekadar mengimbangi overhead kompilasi. Dalam praktiknya, kami telah melihat peningkatan throughput lebih dari 4X lipat dengan X10 dalam beberapa kasus pelatihan.

Seperti yang telah dinyatakan sebelumnya, penggunaan X10 kini tidak hanya memungkinkan tetapi juga memudahkan untuk bekerja dengan TPU, membuka seluruh kelas akselerator untuk model Swift untuk TensorFlow Anda.