הכירו את 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 מבצעת פעולות טנזור תוך שימוש בשיגור נלהב. זה מאפשר איטרציה מהירה, אך אינו האפשרות הביצועית ביותר לאימון מודלים של למידת מכונה.

ספריית הטנזור X10 מוסיפה קצה אחורי בעל ביצועים גבוהים ל-Swift for TensorFlow, תוך מינוף של מעקב אחר טנזור ו- XLA מהדר . מדריך זה יציג את X10 וידריך אותך בתהליך של עדכון לולאת אימון להפעלה על GPUs או TPUs.

להוט לעומת טנסור X10

חישובים מואצים ב-Swift עבור TensorFlow מבוצעים באמצעות סוג 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, אתה אמור לראות את החומרה הזו משתקפת בתיאור ההתקן שלמעלה. לזמן הריצה הנלהב אין תמיכה ב-TPUs, כך שאם אתה משתמש באחד מהם כמאיץ, תראה שה-CPU משמש כיעד חומרה.

בעת יצירת Tensor, ניתן לעקוף את ברירת המחדל של התקן מצב להוט על ידי ציון חלופה. כך אתה מצטרף לביצוע חישובים באמצעות X10 backend.

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 הוא הדרך שבה אתה מנצל את היתרונות של TPUs בתוך Swift for TensorFlow.

מכשירי ברירת המחדל להוט ו-X10 ינסו להשתמש במאיץ הראשון במערכת. אם יש לך GPUs מחוברים, ה-GPU ישתמש ב-GPU הזמין הראשון. אם קיימים TPUs, X10 ישתמש בליבת ה-TPU הראשונה כברירת מחדל. אם לא נמצא או נתמך מאיץ, התקן ברירת המחדל יחזור למעבד.

מעבר למכשירי ברירת המחדל להוט ו-XLA, אתה יכול לספק יעדי חומרה וקצה עורפי ספציפיים במכשיר:

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

הכשרת דוגמנית במצב להוט

בואו נסתכל על האופן שבו הייתם מגדירים ומאמנים מודל באמצעות מצב ברירת המחדל של ביצוע להוט. בדוגמה זו, נשתמש במודל הפשוט של LeNet-5 ממאגר המודלים של swift ומערך הנתונים של סיווג הספרות בכתב יד של 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
    }
}

לבסוף, נריץ את המודל דרך לולאת אימון למשך חמש תקופות.

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)

עבור מערכי הנתונים, נעשה זאת בנקודה שבה אצווה מעובדת בלולאת האימון, כדי שנוכל לעשות שימוש חוזר במערך הנתונים ממודל הביצוע הנלהב.

במקרה של הדגם והאופטימיזר, נאתחל אותם עם הטנזורים הפנימיים שלהם במכשיר הביצוע הנלהב, ואז נעביר אותם למכשיר 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 האחורי היה צריך להתנהל באותו אופן כפי שעשה מודל הביצוע הנלהב קודם לכן. ייתכן ששמתם לב לעיכוב לפני האצווה הראשונה ובסוף העידן הראשון, עקב הידור בדיוק בזמן של העקבות הייחודיות באותן נקודות. אם אתה מפעיל את זה עם מאיץ מחובר, היית צריך לראות את האימון לאחר נקודה זו מתקדם מהר יותר מאשר במצב להוט.

ישנו פשרה של זמן הידור מעקב ראשוני לעומת תפוקה מהירה יותר, אך ברוב המודלים של למידת מכונה העלייה בתפוקה מפעולות חוזרות צריכה יותר מאשר לקזז את תקרת ההידור. בפועל, ראינו שיפור של מעל פי 4 בתפוקה עם X10 במקרים מסוימים.

כפי שצוין בעבר, השימוש ב-X10 מאפשר כעת לא רק לעבוד עם TPUs, אלא גם קל, ופותח את כל המעמד הזה של מאיצים עבור דגמי Swift for TensorFlow שלך.