Giới thiệu X10

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


Xem trên TensorFlow.org Chạy trong Google Colab Xem nguồn trên GitHub

Theo mặc định, Swift For TensorFlow thực hiện các thao tác tensor bằng cách sử dụng công văn háo hức. Điều này cho phép lặp lại nhanh chóng nhưng không phải là lựa chọn hiệu quả nhất để đào tạo các mô hình học máy.

Thư viện tensor X10 bổ sung phần phụ trợ hiệu suất cao cho Swift dành cho TensorFlow, tận dụng tính năng theo dõi tensor và trình biên dịch XLA . Hướng dẫn này sẽ giới thiệu X10 và hướng dẫn bạn quy trình cập nhật vòng đào tạo để chạy trên GPU hoặc TPU.

Háo hức so với tensor X10

Tính toán tăng tốc trong Swift cho TensorFlow được thực hiện thông qua loại Tensor. Tensor có thể tham gia vào nhiều hoạt động khác nhau và là khối xây dựng cơ bản của các mô hình học máy.

Theo mặc định, Tensor sử dụng khả năng thực thi háo hức để thực hiện các phép tính trên cơ sở từng hoạt động. Mỗi Tensor có một Thiết bị liên quan mô tả phần cứng nào được gắn vào và phần phụ trợ nào được sử dụng cho nó.

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)

Nếu bạn đang chạy sổ ghi chép này trên một phiên bản hỗ trợ GPU, bạn sẽ thấy phần cứng đó được phản ánh trong mô tả thiết bị ở trên. Thời gian chạy háo hức không hỗ trợ TPU, vì vậy nếu bạn đang sử dụng một trong số chúng làm bộ tăng tốc, bạn sẽ thấy CPU được sử dụng làm mục tiêu phần cứng.

Khi tạo Tensor, thiết bị ở chế độ háo hức mặc định có thể bị ghi đè bằng cách chỉ định một thiết bị thay thế. Đây là cách bạn chọn tham gia thực hiện các phép tính bằng chương trình phụ trợ 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)

Nếu bạn đang chạy ứng dụng này trong phiên bản hỗ trợ GPU, bạn sẽ thấy bộ tăng tốc đó được liệt kê trong thiết bị của tensor X10. Không giống như thực thi háo hức, nếu bạn đang chạy ứng dụng này trong phiên bản hỗ trợ TPU, giờ đây bạn sẽ thấy các phép tính đang sử dụng thiết bị đó. X10 là cách bạn tận dụng TPU trong Swift cho TensorFlow.

Các thiết bị háo hức và X10 mặc định sẽ cố gắng sử dụng bộ tăng tốc đầu tiên trên hệ thống. Nếu bạn có gắn GPU, GPU sẽ sử dụng GPU đầu tiên có sẵn. Nếu có TPU, X10 sẽ mặc định sử dụng lõi TPU đầu tiên. Nếu không tìm thấy hoặc hỗ trợ bộ tăng tốc, thiết bị mặc định sẽ quay trở lại CPU.

Ngoài các thiết bị XLA và háo hức mặc định, bạn có thể cung cấp các mục tiêu phần cứng và phụ trợ cụ thể trong Thiết bị:

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

Đào tạo mô hình chế độ háo hức

Chúng ta hãy xem cách bạn thiết lập và huấn luyện một mô hình bằng chế độ thực thi háo hức mặc định. Trong ví dụ này, chúng tôi sẽ sử dụng mô hình LeNet-5 đơn giản từ kho lưu trữ mô hình swift và tập dữ liệu phân loại chữ số viết tay MNIST.

Đầu tiên, chúng ta sẽ thiết lập và tải xuống tập dữ liệu 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

Tiếp theo, chúng ta sẽ cấu hình mô hình và trình tối ưu hóa.

import ImageClassificationModels

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

Bây giờ, chúng tôi sẽ triển khai theo dõi và báo cáo tiến độ cơ bản. Tất cả số liệu thống kê trung gian được lưu giữ dưới dạng tensor trên cùng một thiết bị nơi quá trình đào tạo được thực hiện và scalarized() chỉ được gọi trong khi báo cáo. Điều này sẽ đặc biệt quan trọng sau này khi sử dụng X10, vì nó tránh việc hiện thực hóa các tensor lười biếng không cần thiết.

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

Cuối cùng, chúng ta sẽ chạy mô hình thông qua vòng huấn luyện trong năm kỷ nguyên.

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

Như bạn có thể thấy, mô hình được đào tạo như chúng ta mong đợi và độ chính xác của nó so với tập hợp xác thực đã tăng lên sau mỗi kỷ nguyên. Đây là cách Swift cho các mô hình TensorFlow được xác định và chạy bằng cách sử dụng thực thi háo hức, bây giờ hãy xem những sửa đổi nào cần được thực hiện để tận dụng X10.

Đào tạo mô hình X10

Bộ dữ liệu, mô hình và trình tối ưu hóa chứa các tensor được khởi tạo trên thiết bị thực thi háo hức mặc định. Để hoạt động với X10, chúng ta cần di chuyển các tensor này sang thiết bị X10.

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

Đối với các tập dữ liệu, chúng tôi sẽ thực hiện điều đó tại thời điểm các lô được xử lý trong vòng đào tạo, để chúng tôi có thể sử dụng lại tập dữ liệu từ mô hình thực thi háo hức.

Trong trường hợp mô hình và trình tối ưu hóa, chúng tôi sẽ khởi tạo chúng bằng các tensor bên trong trên thiết bị thực thi háo hức, sau đó chuyển chúng sang thiết bị X10.

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

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

Những sửa đổi cần thiết cho vòng đào tạo xuất hiện ở một số điểm cụ thể. Trước tiên, chúng ta cần chuyển các lô dữ liệu huấn luyện sang thiết bị X10. Việc này được thực hiện thông qua Tensor(copying:to:) khi mỗi lô được truy xuất.

Thay đổi tiếp theo là chỉ ra nơi cần cắt bỏ dấu vết trong vòng huấn luyện. X10 hoạt động bằng cách truy tìm các phép tính tensor cần thiết trong mã của bạn và biên dịch kịp thời một biểu diễn được tối ưu hóa cho dấu vết đó. Trong trường hợp vòng lặp đào tạo, bạn lặp đi lặp lại cùng một thao tác, đây là phần lý tưởng để theo dõi, biên dịch và sử dụng lại.

Trong trường hợp không có mã yêu cầu rõ ràng một giá trị từ Tensor (những giá trị này thường nổi bật dưới dạng lệnh gọi .scalars hoặc .scalarized() ), X10 sẽ cố gắng biên dịch tất cả các lần lặp vòng lặp lại với nhau. Để ngăn chặn điều này và cắt dấu vết tại một điểm cụ thể, chúng tôi đặt LazyTensorBarrier() rõ ràng sau khi trình tối ưu hóa cập nhật trọng số mô hình cũng như sau khi đạt được độ mất mát và độ chính xác trong quá trình xác thực. Điều này tạo ra hai dấu vết được sử dụng lại: mỗi bước trong vòng đào tạo và mỗi đợt suy luận trong quá trình xác thực.

Những sửa đổi này dẫn đến vòng lặp đào tạo sau.

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

Việc đào tạo mô hình bằng phần phụ trợ X10 lẽ ra phải được tiến hành theo cách tương tự như mô hình thực thi háo hức đã làm trước đây. Bạn có thể nhận thấy sự chậm trễ trước đợt đầu tiên và ở cuối kỷ nguyên đầu tiên, do việc tổng hợp kịp thời các dấu vết duy nhất tại những điểm đó. Nếu bạn đang chạy chương trình này có gắn máy gia tốc, bạn sẽ thấy quá trình đào tạo sau thời điểm đó diễn ra nhanh hơn so với chế độ háo hức.

Có sự cân bằng giữa thời gian biên dịch dấu vết ban đầu so với thông lượng nhanh hơn, nhưng trong hầu hết các mô hình học máy, mức tăng thông lượng từ các hoạt động lặp lại sẽ bù đắp nhiều hơn chi phí biên dịch. Trong thực tế, chúng tôi đã thấy thông lượng cải thiện hơn 4 lần với X10 trong một số trường hợp đào tạo.

Như đã nêu trước đây, việc sử dụng X10 giờ đây giúp bạn không chỉ có thể làm việc với TPU mà còn dễ dàng, mở khóa toàn bộ loại bộ tăng tốc đó cho các mẫu Swift for TensorFlow của bạn.