Kubeflow：為何選擇TF工作實作生產級機器學習佈署

– 13 分鐘閱讀

在將機器學習模型從實驗階段推進到生產環境時，我們經常面臨一系列技術挑戰。這些挑戰往往超出了模型本身的複雜性，涉及基礎設施、資源管理和佈署流程等多個層面。

機器學習佈署的關鍵挑戰

在實際工作中，我發現將訓練程式碼佈署到生產環境時，開發團隊通常會遇到以下難題：

  • 基礎設施差異：雲端與本地環境的設定與管理方式存在顯著差異，導致遷移成本高昂
  • 許可權與協作：如何設計一個系統，讓團隊成員能夠分享訓練任務和資料，同時維持適當的存取控制
  • 資源擴充套件與回收：訓練任務完成後的資源釋放機制，以及根據需求動態擴充套件計算資源的能力
  • 模型匯出與服務化：訓練完成後，如何將模型規範化處理並佈署為可用的推論服務

這些問題通常需要大量的"膠水程式碼"(glue code)來處理底層基礎設施，而這些程式碼往往因環境差異而需要重寫。在我的經驗中，這部分技術開發可能比模型本身的開發更加耗時。

Kubeflow與Kubernetes：容器化機器學習的最佳組合

Kubeflow正是為解決上述問題而設計的。它建立在Kubernetes之上，繼承了Kubernetes的可擴充套件性和可移植性特點。在實際使用中，我發現Kubeflow特別適合以下場景：

  • 需要在不同雲平台間遷移的機器學習工作流程
  • 需要快速擴充套件計算資源的大規模訓練任務
  • 團隊協作開發的機器學習專案

Kubeflow透過自定義資源（Custom Resources）將TensorFlow任務協調到Kubernetes叢集上。這種方式讓開發者只需關注應用的"期望狀態"，底層控制器會自動處理實際的佈署和管理過程。

從Jupyter Notebook到TF工作：容器化你的訓練程式碼

將Jupyter Notebook轉換為生產級TensorFlow訓練任務需要幾個關鍵步驟。以下是我在實踐中總結的最佳流程。

準備訓練程式碼

首先，我們需要將Jupyter Notebook匯出為Python指令碼：

  1. 在Jupyter介面中選擇"File"
  2. 點選"Download as"
  3. 選擇"Python (.py)"

這樣就能獲得一個可執行的Python檔案，它包含了我們在Notebook中開發的所有程式碼。

構建訓練容器

接下來，我們需要將訓練程式碼封裝到容器中。這裡是一個典型的Dockerfile範例：

FROM tensorflow/tensorflow:1.15.2-py3
RUN pip3 install --upgrade pip
RUN pip3 install pandas --upgrade
RUN pip3 install keras --upgrade
RUN pip3 install minio --upgrade
RUN mkdir -p /opt/kubeflow
COPY Recommender_Kubeflow.py /opt/kubeflow/
ENTRYPOINT ["python3", "/opt/kubeflow/Recommender_Kubeflow.py"]

這個Dockerfile做了幾件核心事情：

  1. 使用官方TensorFlow 1.15.2 Python 3基礎映像作為起點
  2. 升級pip並安裝必要的依賴包（pandas、keras和minio）
  3. 建立工作目錄並複製訓練指令碼
  4. 設定容器啟動點為執行訓練指令碼

在實踐中，我發現為每個版本的訓練程式碼建立單獨的容器標籤是個好習慣，這有助於版本控制和回復。

構建並推播容器映像

準備好Dockerfile後，執行以下命令構建並推播容器：

docker build -t kubeflow/recommenderjob:1.0 .
docker push kubeflow/recommenderjob:1.0

這兩條命令完成了容器的構建和發布：

  • docker build命令從當前目錄的Dockerfile建立映像，並標記為kubeflow/recommenderjob:1.0
  • docker push命令將構建好的映像推播到容器倉函式庫，使其可在Kubernetes叢集中使用

需注意的是，你可能需要先登入到你的容器倉函式庫（如Docker Hub、Google Container Registry或私有倉函式庫）。

定義並佈署單容器TF工作

容器準備好後，我們可以建立TF工作設定檔案，告訴Kubeflow如何執行訓練任務。

TF工作規範詳解

以下是一個單容器TF工作的YAML設定範例：

apiVersion: "kubeflow.org/v1"
kind: "TF工作"
metadata:
  name: "recommenderjob"
spec:
  tfReplicaSpecs:
    Worker:
      replicas: 1
      restartPolicy: Never
      template:
        spec:
          containers:
          - name: tensorflow 
            image: kubeflow/recommenderjob:1.0

這個YAML檔案定義了一個TF工作資源，它包含以下關鍵部分：

  • apiVersion: 指定使用的Kubeflow API版本，這裡是kubeflow.org/v1
  • kind: 資源型別，這裡是TF工作
  • metadata: 包含資源的中繼資料，如名稱
  • spec.tfReplicaSpecs: 定義TensorFlow訓練叢集的設定
    • Worker: 指定工作節點的設定
      • replicas: 指定工作節點的數量，這裡是1（單節點訓練）
      • restartPolicy: 指定Pod失敗時的重啟策略，這裡是Never
      • template: 定義Pod的規範，包括使用的容器映像

在實際專案中，我通常會根據訓練任務的特性調整這些設定，尤其是在處理需要長時間執行的訓練任務時。

TF工作的進階設定選項

除了基本設定外，TF工作還支援多種進階選項：

  • activeDeadlineSeconds: 任務的最長執行時間，超過後系統會終止任務
  • backoffLimit: 任務失敗後的重試次數上限
  • cleanPodPolicy: 任務完成後Pod的清理策略，可設為All（清理所有Pod）、Running（只清理執行中的Pod）或None（不清理任何Pod）

在除錯複雜模型時，我通常會設定cleanPodPolicy: None以保留Pod的日誌和狀態訊息，這對分析訓練問題非常有幫助。

佈署和監控TF工作

使用以下命令將TF工作佈署到Kubernetes叢集：

kubectl apply -f recommenderjob.yaml

佈署後，可以透過以下命令監控任務狀態：

kubectl describe tfjob recommenderjob

執行上述命令後，你會看到類別似以下的輸出：

Status:
  Completion Time: 2019-05-18T00:58:27Z
  Conditions:
    Last Transition Time: 2019-05-18T02:34:24Z
    Last Update Time: 2019-05-18T02:34:24Z
    Message: TF工作 recommenderjob is created.
    Reason: TF工作Created
    Status: True
    Type: Created
    Last Transition Time: 2019-05-18T02:38:28Z
    Last Update Time: 2019-05-18T02:38:28Z
    Message: TF工作 recommenderjob is running.
    Reason: TF工作Running
    Status: False
    Type: Running
    Last Transition Time: 2019-05-18T02:38:29Z
    Last Update Time: 2019-05-18T02:38:29Z
    Message: TF工作 recommenderjob successfully completed.
    Reason: TF工作Succeeded
    Status: True
    Type: Succeeded
  Replica Statuses:
    Worker:
      Succeeded: 1

這個輸出提供了TF工作的詳細狀態訊息：

  • Conditions: 顯示任務的狀態變化歷史，包括建立、執行和完成的時間點
  • Replica Statuses: 顯示各型別副本的狀態，這裡顯示一個Worker已成功完成

這些訊息對於除錯非常有價值，特別是當任務失敗時，可以從這裡找到失敗原因。

分散式訓練：突破單機計算瓶頸

隨著模型複雜度和資料量的增長，單機訓練往往會遇到計算資源的瓶頸。在實際工作中，我經常需要處理引數量超過數億的模型或TB級別的訓練資料，這時分散式訓練就成為必然選擇。

分散式訓練的基本概念

分散式訓練主要有兩種方式：

  1. 資料平行性（Data Parallelism）：

    • 將訓練資料分割成多個部分
    • 每個工作節點執行相同的模型，但處理不同的資料子集
    • 定期同步各節點的梯度訊息
    • 適合資料量大但模型相對小的情況

  2. 模型平行性（Model Parallelism）：

    • 將模型本身分割到不同的工作節點
    • 所有節點使用相同的訓練資料
    • 每個節點負責計算模型的一部分
    • 適合模型極其龐大無法裝入單個裝置記憶體的情況

在實踐中，我發現大多數專案會採用資料平行性，因為它實作相對簡單與擴充套件性好。只有在處理超大型模型（如大模型語言）時，才會考慮模型平行性或混合平行策略。

TensorFlow的分散式策略

TensorFlow提供了多種分散式訓練策略，每種策略適用於不同的場景：

  1. 映象策略（Mirrored Strategy）：

    • 同步訓練策略，所有工作節點的訓練步驟和梯度計算保持同步
    • 模型變數在所有裝置和工作節點間複製
    • 適合單機多GPU訓練

  2. TPU策略（TPU Strategy）：

    • 類別似映象策略，但專為Google的TPU加速器最佳化
    • 支援大規模型訓練，利用TPU的高效能計算能力

  3. 多工作節點映象策略（Multiworker Mirrored Strategy）：

    • 根據映象策略的擴充套件，使用CollectiveOps實作多工作節點間的梯度聚合
    • 適合跨多台機器的分散式訓練

  4. 中央儲存策略（Central Storage Strategy）：

    • 不是在所有工作節點間複製變數，而是將變數儲存在中央裝置上
    • 減少記憶體使用，但可能增加通訊開銷

在實際專案中，我通常根據以下因素選擇適當的策略：

  • 可用硬體（GPU、TPU或CPU）
  • 模型大小和複雜度
  • 訓練資料規模
  • 對訓練速度的要求

對於大多數中小型模型，多工作節點映象策略是個不錯的選擇，它提供了良好的擴充套件性和相對簡單的實作。

在Kubeflow中實作分散式TensorFlow訓練

Kubeflow使分散式TensorFlow訓練的設定變得相對簡單。與單機訓練相比，主要區別在於TF工作設定中需要定義多種角色。

分散式TF工作的角色型別

在TensorFlow分散式訓練中，通常會涉及以下幾種角色：

  1. Chief：主工作節點，負責協調訓練過程和模型檢查點的儲存
  2. Worker：普通工作節點，執行實際的訓練計算
  3. PS（Parameter Server）：引數伺服器，在某些架構中用於儲存和更新模型引數
  4. Evaluator：評估器，在訓練過程中定期評估模型效能

根據所選的分散式策略，我們可能會使用不同組合的角色。例如，在使用引數伺服器架構時，我們需要同時設定Worker和PS；而在使用AllReduce架構時，可能只需要Worker和Chief。

分散式TF工作設定範例

以下是一個使用引數伺服器架構的分散式TF工作設定範例：

apiVersion: "kubeflow.org/v1"
kind: "TF工作"
metadata:
  name: "distributed-recommenderjob"
spec:
  tfReplicaSpecs:
    Chief:
      replicas: 1
      restartPolicy: Never
      template:
        spec:
          containers:
          - name: tensorflow
            image: kubeflow/recommenderjob:1.0
            env:
            - name: TF_CONFIG
              value: '{"cluster": {"chief": ["$(CHIEF_SERVICE):2222"], "worker": ["$(WORKER_0_SERVICE):2222", "$(WORKER_1_SERVICE):2222"], "ps": ["$(PS_0_SERVICE):2222", "$(PS_1_SERVICE):2222"]}, "task": {"type": "chief", "index": 0}}'
    Worker:
      replicas: 2
      restartPolicy: Never
      template:
        spec:
          containers:
          - name: tensorflow
            image: kubeflow/recommenderjob:1.0
    PS:
      replicas: 2
      restartPolicy: Never
      template:
        spec:
          containers:
          - name: tensorflow
            image: kubeflow/recommenderjob:1.0

這個設定義了一個包含1個Chief、2個Worker和2個PS的分散式訓練任務：

  • Chief: 主工作節點，只有一個例項
  • Worker: 普通工作節點，有兩個例項，負責大部分計算工作
  • PS: 引數伺服器，有兩個例項，負責儲存和更新模型引數

每個角色使用相同的容器映像，但透過環境變數TF_CONFIG設定不同的角色和任務。在實際佈署中，Kubeflow會自動填充這些環境變數，實作節點間的服務發現。

適配訓練程式碼以支援分散式訓練

要讓訓練程式碼支援分散式訓練，通常需要進行一些修改。以下是一個使用TensorFlow分散式策略API的範例：

import tensorflow as tf
import os

# 取得TF_CONFIG環境變數
tf_config = os.environ.get('TF_CONFIG', '{}')
import json
tf_config = json.loads(tf_config)

# 設定分散式策略
strategy = tf.distribute.experimental.MultiWorkerMirroredStrategy()

# 在策略範圍內建立模型
with strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(64, activation='relu', input_shape=(10,)),
        tf.keras.layers.Dense(32, activation='relu'),
        tf.keras.layers.Dense(1)
    ])
    
    model.compile(
        optimizer=tf.keras.optimizers.Adam(0.001),
        loss='mse',
        metrics=['mae']
    )

# 準備資料
# 這裡假設我們已經有了訓練資料
train_dataset = ...

# 訓練模型
model.fit(train_dataset, epochs=10)

# 只在Chief工作節點儲存模型
if tf_config.get('task', {}).get('type', '') == 'chief':
    model.save('/path/to/model')

這段程式碼展示瞭如何使用TensorFlow的分散式策略API實作分散式訓練：

  1. 從環境變數中取得TF_CONFIG，這是Kubeflow自動注入的設定
  2. 建立MultiWorkerMirroredStrategy分散式策略
  3. 在策略範圍內建立和編譯模型
  4. 使用標準的Keras API訓練模型
  5. 只在Chief工作節點上儲存模型，避免多節點同時寫入造成衝突

在實際專案中，我發現正確處理模型儲存和檢查點是分散式訓練中最容易出錯的部分。建議只在Chief節點執行這些操作，並使用分享儲存（如NFS或雲端儲存）儲存結果。

最佳化分散式訓練效能

在實施分散式訓練時，我們經常會遇到效能問題。以下是一些我在實踐中總結的最佳化技巧。

資料輸入管道最佳化

資料輸入往往是分散式訓練的瓶頸。以下是一些最佳化方法：

  1. 使用TFRecord格式：將訓練資料預處理為TFRecord格式，提高讀取效率
  2. 實施資料預取：使用tf.data.Dataset.prefetch()預取下一批資料
  3. 平行資料處理：使用tf.data.Dataset.map(num_parallel_calls=...)平行處理資料
  4. 資料快取：對於較小的資料集，使用tf.data.Dataset.cache()將資料快取在記憶體中

通訊開銷最佳化

節點間通訊是分散式訓練的另一個瓶頸：

  1. 梯度壓縮：使用梯度壓縮技術減少通訊量
  2. 調整批次大小：增大批次大小可以減少同步次數，但要平衡與模型收斂的關係
  3. 使用高效的集合通訊函式庫：如NCCL或Horovod
  4. 網路拓撲感知：在Kubernetes中使用節點親和性（Node Affinity）將相關Pod排程到網路延遲較低的節點

資源分配最佳化

合理的資源分配可以顯著提高訓練效率：

  1. GPU記憶體最佳化：使用混合精確度訓練（如tf.keras.mixed_precision）減少GPU記憶體使用
  2. Worker與PS比例：調整Worker與PS的數量比例，通常2:1或3:1較為合適
  3. 資源限制設定：在TF工作中為不同角色設定適當的CPU/GPU資源限制
  4. 節點選擇策略：為不同角色選擇適合的節點型別，如為PS選擇高記憶體節點，為Worker選擇高GPU節點

在一個大型推薦系統專案中，我透過最佳化這些引數，將分散式訓練速度提高了近3倍，同時保持了模型效能不變。這些最佳化雖然看似微小，但累積起來效果顯著。

從訓練到服務：模型匯出與佈署

訓練完成後，下一步是將模型匯出並佈署為推論服務。Kubeflow提供了完整的工具鏈支援這一流程。

模型匯出最佳實踐

在分散式訓練中，模型匯出需要特別注意：

  1. 只在Chief節點匯出：避免多節點同時寫入造成衝突
  2. 使用標準格式：匯出為SavedModel格式，確保相容性
  3. 版本管理：為模型增加版本訊息，便於管理和回復
  4. 中繼資料記錄：記錄訓練引數、資料版本等中繼資料

以下是一個模型匯出的程式碼範例：

import tensorflow as tf
import os
import json
import time

# 只在Chief節點匯出模型
tf_config = json.loads(os.environ.get('TF_CONFIG', '{}'))
if tf_config.get('task', {}).get('type', '') == 'chief':
    # 建立版本化路徑
    version = int(time.time())
    export_path = f"/models/recommender/{version}"
    
    # 儲存模型
    tf.keras.models.save_model(
        model,
        export_path,
        overwrite=True,
        include_optimizer=False,
        save_format='tf',
        signatures=None,
        options=None
    )
    
    # 儲存中繼資料
    metadata = {
        "framework": "tensorflow",
        "version": str(version),
        "training_params": {
            "epochs": 10,
            "batch_size": 128,
            "optimizer": "adam"
        },
        "dataset": "movie_ratings_v2",
        "accuracy": float(final_metrics["val_accuracy"])
    }
    
    with open(f"{export_path}/metadata.json", "w") as f:
        json.dump(metadata, f)
    
    print(f"Model exported to: {export_path}")

這段程式碼展示了模型匯出的最佳實踐：

  1. 檢查當前節點是否為Chief，只在Chief節點執行匯出
  2. 使用時間戳建立版本化的模型路徑
  3. 使用標準的SavedModel格式儲存模型，不包含最佳化器狀態以減小體積
  4. 儲存包含訓練引數、資料版本和效能指標的中繼資料
  5. 輸出確認訊息，便於日誌追蹤

使用KFServing佈署模型

Kubeflow生態系統中的KFServing（現已改名為KServe）元件可以輕鬆將訓練好的模型佈署為REST API服務。以下是一個KFServing佈署設定範例：

apiVersion: "serving.kubeflow.org/v1beta1"
kind: "InferenceService"
metadata:
  name: "recommender-service"
spec:
  default:
    predictor:
      tensorflow:
        storageUri: "pvc://models/recommender"
        resources:
          limits:
            cpu: "2"
            memory: "4Gi"

這個設定義了一個根據TensorFlow模型的推論服務：

  • apiVersion/kind: 指定這是一個KFServing InferenceService資源
  • metadata.name: 服務名稱，用於存取和管理
  • spec.default.predictor.tensorflow: 指定使用TensorFlow作為推論引擎
  • storageUri: 指定模型儲存位置，這裡使用了PVC（永續性儲存區宣告）
  • resources: 定義服務使用的資源限制

佈署後，KFServing會自動為模型建立REST API端點，客戶端可以透過HTTP請求進行推論。

整合CI/CD管道

在企業環境中，我通常會將訓練和佈署整合到CI/CD管道中，實作自動化。一個典型的管道包括以下步驟：

  1. 程式碼變更觸發：當模型程式碼或訓練資料發生變更時觸發管道
  2. 環境準備：準備訓練環境和依賴
  3. 模型訓練：使用TF工作執行分散式訓練
  4. 模型評估：評估模型效能，如果達不到閾值則中止管道
  5. 模型註冊：將合格的模型註冊到模型倉函式庫
  6. 金絲雀佈署：先佈署到測試環境進行驗證
  7. 生產佈署：透過驗證後佈署到生產環境
  8. 監控與回復：持續監控模型效能，必要時自動回復

這種自動化管道大提高了模型迭代的效率，同時確保了生產環境的穩定性。

進階主題：混合精確度訓練與模型最佳化

在處理大規模型時，除了分散式訓練外，還有一些進階技術可以進一步提升效能。

混合精確度訓練

混合精確度訓練使用較低精確度（如FP16）進行計算，同時保持模型權重在FP32精確度，這可以顯著提高訓練速度和減少記憶體使用。在TensorFlow中實作混合精確度訓練相對簡單：

# 啟用混合精確度訓練
from tensorflow.keras.mixed_precision import experimental as mixed_precision
policy = mixed_precision.Policy('mixed_float16')
mixed_precision.set_global_policy(policy)

# 在分散式策略範圍內建立模型
with strategy.scope():
    model = tf.keras.Sequential([...])
    
    # 注意：在混合精確度下，最佳化器需要使用loss scaling
    optimizer = tf.keras.optimizers.Adam(0.001)
    optimizer = mixed_precision.LossScaleOptimizer(optimizer)
    
    model.compile(
        optimizer=optimizer,
        loss='mse',
        metrics=['mae']
    )

這段程式碼展示了混合精確度訓練的實作：

  1. 設定全域混合精確度策略為’mixed_float16'
  2. 在分散式策略範圍內建立模型
  3. 使用LossScaleOptimizer包裝原始最佳化器，防止梯度下溢

在支援FP16的硬體（如NVIDIA Volta或更新架構的GPU）上，混合精確度訓練可以將訓練速度提高2-3倍，同時允許使用更大的批次大小。

模型量化與最佳化

對於推論佈署，模型量化是一種有效的最佳化技術：

# 訓練後量化
converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
quantized_model = converter.convert()

# 儲存量化模型
with open('/models/recommender_quantized.tflite', 'wb') as f:
    f.write(quantized_model)

這段程式碼演示瞭如何對SavedModel格式的模型進行訓練後量化：

  1. 使用TFLiteConverter從SavedModel載入模型
  2. 設定DEFAULT最佳化，這會自動應用量化技術
  3. 轉換模型並儲存為TFLite格式

量化後的模型體積通常只有原始模型的1/4左右，推論速度也有顯著提升，特別適合移動裝置和邊緣裝置佈署。

玄貓

玄貓

技術愛好者，專注於分享程式開發、雲端技術與 AI 應用的心得體會。