X10 소개

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


TensorFlow.org에서 보기 GitHub에서 소스 보기

기본적으로 Swift For TensorFlow는 Eager Dispatch를 사용하여 텐서 작업을 수행합니다. 이를 통해 빠른 반복이 가능하지만 기계 학습 모델을 훈련하는 데 가장 성능이 좋은 옵션은 아닙니다.

X10 텐서 라이브러리는 텐서 추적 및 XLA 컴파일러를 활용하여 Swift for TensorFlow에 고성능 백엔드를 추가합니다. 이 튜토리얼에서는 X10을 소개하고 GPU 또는 TPU에서 실행되도록 훈련 루프를 업데이트하는 과정을 안내합니다.

Eager 대 X10 텐서

Swift for TensorFlow의 가속화된 계산은 Tensor 유형을 통해 수행됩니다. Tensor는 다양한 작업에 참여할 수 있으며 기계 학습 모델의 기본 구성 요소입니다.

기본적으로 Tensor는 즉시 실행을 사용하여 작업별로 계산을 수행합니다. 각 Tensor에는 연결된 하드웨어와 이에 사용되는 백엔드를 설명하는 관련 장치가 있습니다.

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)

GPU 지원 인스턴스에서 이 노트북을 실행하는 경우 위의 장치 설명에 하드웨어가 반영되어 있는 것을 볼 수 있습니다. Eager 런타임은 TPU를 지원하지 않으므로 TPU 중 하나를 가속기로 사용하는 경우 CPU가 하드웨어 대상으로 사용되는 것을 볼 수 있습니다.

Tensor를 생성할 때 기본 Eager 모드 장치는 대안을 지정하여 재정의될 수 있습니다. 이것이 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)

GPU 지원 인스턴스에서 이를 실행하는 경우 X10 텐서의 장치에 해당 가속기가 나열되어 있는 것을 볼 수 있습니다. 즉시 실행과 달리 TPU 지원 인스턴스에서 이를 실행하는 경우 이제 계산이 해당 장치를 사용하고 있음을 확인할 수 있습니다. X10은 Swift for TensorFlow 내에서 TPU를 활용하는 방법입니다.

기본 Eager 및 X10 장치는 시스템의 첫 번째 가속기를 사용하려고 시도합니다. GPU가 연결되어 있는 경우 사용 가능한 첫 번째 GPU가 사용됩니다. TPU가 있는 경우 X10은 기본적으로 첫 번째 TPU 코어를 사용합니다. 가속기가 발견되지 않거나 지원되지 않으면 기본 장치는 CPU로 대체됩니다.

기본 Eager 및 XLA 장치 외에도 장치에 특정 하드웨어 및 백엔드 대상을 제공할 수 있습니다.

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

Eager 모드 모델 학습

기본 즉시 실행 모드를 사용하여 모델을 설정하고 훈련하는 방법을 살펴보겠습니다. 이 예에서는 Swift-models 저장소 의 간단한 LeNet-5 모델과 MNIST 필기 숫자 분류 데이터 세트를 사용합니다.

먼저 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

다음으로 모델과 옵티마이저를 구성하겠습니다.

import ImageClassificationModels

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

이제 기본적인 진행 상황 추적 및 보고를 구현하겠습니다. 모든 중간 통계는 훈련이 실행되고 보고 중에만 scalarized() 호출되는 동일한 장치에 텐서로 유지됩니다. 이는 게으른 텐서의 불필요한 구체화를 방지하므로 나중에 X10을 사용할 때 특히 중요합니다.

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

마지막으로 5개 에포크에 대한 훈련 루프를 통해 모델을 실행합니다.

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

보시다시피 모델은 우리가 예상한 대로 훈련되었으며 검증 세트에 대한 정확도는 각 시대마다 증가했습니다. 이것이 바로 Swift for TensorFlow 모델을 정의하고 즉시 실행을 사용하여 실행하는 방법입니다. 이제 X10을 활용하려면 어떤 수정이 필요한지 살펴보겠습니다.

X10 모델 훈련

데이터 세트, 모델 및 최적화 프로그램에는 기본 즉시 실행 장치에서 초기화되는 텐서가 포함되어 있습니다. X10을 사용하려면 이러한 텐서를 X10 장치로 이동해야 합니다.

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

데이터 세트의 경우 훈련 루프에서 배치가 처리되는 시점에 이를 수행하므로 열정적 실행 모델의 데이터 세트를 재사용할 수 있습니다.

모델과 옵티마이저의 경우 Eager Execution 장치의 내부 텐서로 초기화한 다음 X10 장치로 이동합니다.

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

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

훈련 루프에 필요한 수정은 몇 가지 특정 지점에서 이루어집니다. 먼저 훈련 데이터 배치를 X10 장치로 이동해야 합니다. 이는 각 배치가 검색될 때 Tensor(copying:to:) 통해 수행됩니다.

다음 변경 사항은 훈련 루프 중에 추적을 잘라낼 위치를 나타내는 것입니다. X10은 코드에 필요한 텐서 계산을 추적하고 해당 추적의 최적화된 표현을 적시에 컴파일하는 방식으로 작동합니다. 훈련 루프의 경우 동일한 작업을 계속해서 반복하므로 추적, 컴파일 및 재사용에 이상적인 섹션입니다.

Tensor에서 값을 명시적으로 요청하는 코드(보통 .scalars 또는 .scalarized() 호출로 나타남)가 없는 경우 X10은 모든 루프 반복을 함께 컴파일하려고 시도합니다. 이를 방지하고 특정 지점에서 추적을 잘라내기 위해 최적화 프로그램이 모델 가중치를 업데이트하고 검증 중에 손실과 정확도를 얻은 후 명시적인 LazyTensorBarrier() 배치합니다. 이렇게 하면 훈련 루프의 각 단계와 검증 중 각 추론 배치라는 두 개의 재사용된 추적이 생성됩니다.

이러한 수정으로 인해 다음과 같은 훈련 루프가 생성됩니다.

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

X10 백엔드를 사용한 모델 학습은 이전 Eager Execution 모델과 동일한 방식으로 진행되어야 합니다. 해당 지점에서 고유한 추적이 적시에 컴파일되기 때문에 첫 번째 배치 이전과 첫 번째 에포크가 끝날 때 지연이 발생하는 것을 발견했을 수 있습니다. 가속기가 연결된 상태에서 이를 실행하는 경우 해당 지점 이후의 훈련이 Eager 모드보다 더 빠르게 진행되는 것을 볼 수 있어야 합니다.

초기 추적 컴파일 시간과 더 빠른 처리량 사이에는 장단점이 있지만 대부분의 기계 학습 모델에서는 반복 작업으로 인한 처리량 증가가 컴파일 오버헤드를 상쇄하는 것 이상이어야 합니다. 실제로 일부 교육 사례에서는 X10을 사용하여 처리량이 4배 이상 향상된 것을 확인했습니다.

이전에 언급했듯이 이제 X10을 사용하면 TPU 작업이 가능할 뿐만 아니라 쉬워지며 Swift for TensorFlow 모델에 대한 전체 가속기 클래스를 사용할 수 있습니다.