返回文章列表

PyTorch 模型生產環境佈署實踐

本文介紹如何將 PyTorch 模型佈署到生產環境,包含使用 Flask 建立 Web 服務,使用 Docker 容器化模型服務,簡化 Docker 映像大小,以及使用 Kubernetes 佈署和管理模型服務。文章涵蓋模型載入、預測函式、Dockerfile

機器學習 Web 開發

透過 Flask 建立 Web 服務介面,將訓練好的 PyTorch 貓狗分類別模型佈署上線。利用環境變數指定模型檔案路徑,實作模型更新與服務佈署解耦。文章詳細說明瞭使用 torch.load() 載入模型權重、unsqueeze(0) 調整輸入資料維度以及 torch.max() 取得預測結果等關鍵步驟。為了方便佈署和管理,使用 Docker 將 Flask 服務容器化,Dockerfile 中包含了安裝 PyTorch、Flask 等依賴函式庫,以及設定環境變數和暴露埠等操作。同時,考慮到模型檔案大小,文章也提供瞭解決方案,透過下載模型引數到臨時檔案,避免 Docker 映像過大,方便後續佈署和維護。

將 PyTorch 模型佈署到生產環境

在前面的章節中,我們已經訓練了一個用於分類別貓魚影像的模型。現在,我們將探討如何將這個模型佈署到生產環境中,使其能夠透過 Web 服務對外提供影像分類別功能。

使用 Flask 建立 Web 服務

在這個例子中,我們使用 Flask 框架建立一個簡單的 Web 服務。Flask 的 @app.route() 註解使得我們可以將普通的 Python 函式繫結到特定的 URL 路徑上,當使用者存取該路徑時,函式就會被執行。

from flask import Flask, request, jsonify
import torch
from PIL import Image
from torchvision import transforms

app = Flask(__name__)

# 定義模型載入函式
def load_model():
    m = CatfishModel()
    location = os.environ["CATFISH_MODEL_LOCATION"]
    m.load_state_dict(torch.load(location))
    return m

# 定義預測函式
@app.route('/predict', methods=['GET', 'POST'])
def predict():
    img_url = request.args.get('img_url') or request.form.get('img_url')
    img = Image.open(img_url)
    transform = transforms.Compose([
        transforms.Resize(224),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    img = transform(img).unsqueeze(0)  # 增加 batch 維度
    model = load_model()
    output = model(img)
    _, predicted = torch.max(output, 1)
    class_name = CatfishClasses[predicted.item()]
    return jsonify({'class': class_name, 'img_url': img_url})

if __name__ == '__main__':
    app.run(debug=True)

內容解密:

  1. load_model() 函式:負責載入已經訓練好的模型引數。透過環境變數 CATFISH_MODEL_LOCATION 指定模型權重的檔案路徑,並使用 torch.load() 載入權重,最後將權過載入到模型中。
  2. predict() 函式:處理對 /predict 路徑的請求,從請求中取得影像 URL,開啟影像並進行預處理(縮放、轉換為 Tensor、標準化),然後將處理後的影像輸入模型進行預測,最後傳回預測的類別和影像 URL。
  3. 使用 unsqueeze(0):將輸入影像的 Tensor 增加一個 batch 維度,使其形狀從 [3, 224, 224] 變為 [1, 3, 224, 224],以符合模型的輸入要求。
  4. torch.max(output, 1):找出輸出 Tensor 中最大值的索引,對應到預測的類別。

設定模型引數

在訓練模型後,我們需要將模型的權重儲存下來,並在生產環境中載入。PyTorch 提供了兩種儲存模型的方式:一種是儲存整個模型,另一種是儲存模型的 state_dict。在這裡,我們選擇儲存 state_dict,因為這種方式更加靈活,不會受到模型結構或目錄結構變化的影響。

def load_model():
    m = CatfishModel()
    location = os.environ["CATFISH_MODEL_LOCATION"]
    m.load_state_dict(torch.load(location))
    return m

內容解密:

  1. 環境變數 CATFISH_MODEL_LOCATION:指定模型權重檔案的路徑,使得模型更新與服務佈署可以解耦。
  2. torch.load(location):載入模型權重檔案。
  3. m.load_state_dict():將載入的權過載入到模型中。

使用 Docker 封裝服務

為了方便佈署,我們使用 Docker 將整個服務封裝成一個容器。Dockerfile 定義了容器的構建過程,包括安裝必要的依賴和複製程式碼。

FROM pytorch/pytorch:latest

# 設定工作目錄
WORKDIR /app

# 複製 requirements.txt
COPY requirements.txt .

# 安裝依賴
RUN pip install --no-cache-dir -r requirements.txt

# 複製程式碼
COPY . .

# 設定環境變數
ENV CATFISH_MODEL_LOCATION=/app/catfishweights.pt

# 暴露埠
EXPOSE 5000

# 執行命令
CMD ["python", "app.py"]

內容解密:

  1. FROM pytorch/pytorch:latest:使用官方的 PyTorch 映象作為基礎映象,避免了從頭安裝 PyTorch 的麻煩。
  2. WORKDIR /app:設定容器中的工作目錄為 /app
  3. COPYRUN pip install:複製 requirements.txt 到容器中,並安裝所需的 Python 套件。
  4. ENV:設定環境變數 CATFISH_MODEL_LOCATION,指定模型權重檔案的路徑。
  5. EXPOSE 5000:暴露容器的 5000 埠,供外部存取。
  6. CMD ["python", "app.py"]:指定容器啟動時執行的命令,執行 Flask 應用。

容器化 PyTorch 模型服務

在將 PyTorch 模型佈署到生產環境時,容器化是一種常見的做法。本文將介紹如何使用 Docker 將 PyTorch 模型服務容器化。

Dockerfile 編寫

首先,我們需要編寫一個 Dockerfile 來構建我們的容器映象。以下是一個範例 Dockerfile:

FROM continuumio/miniconda3:latest
ARG model_parameter_location
ARG model_parameter_name
ARG port
ARG host
ENV CATFISH_PORT=$port
ENV CATFISH_HOST=$host
ENV CATFISH_MODEL_LOCATION=/app/$model_parameter_name
RUN conda install -y flask \
&& conda install -c pytorch torchvision \
&& conda install waitress
RUN mkdir -p /app
COPY ./model.py /app
COPY ./server.py /app
COPY $model_parameter_location/$model_parameter_name /app/
COPY ./run-model-service.sh /
EXPOSE $port
ENTRYPOINT ["/run-model-service.sh"]

內容解密:

  1. FROM continuumio/miniconda3:latest:指定基礎映象為 continuumio/miniconda3 的最新版本。
  2. ARGENV:定義變數,ARG 用於在構建時傳遞變數,ENV 用於設定環境變數。
  3. RUN:執行命令,例如安裝必要的套件。
  4. COPY:將檔案從主機複製到容器中。
  5. EXPOSE:指定容器要開放的埠。
  6. ENTRYPOINT:指定容器啟動時要執行的命令。

執行模型服務

構建好 Docker 映象後,我們可以使用以下命令執行模型服務:

docker build -t catfish-service .
docker run -p 5000:5000 catfish-service

內容解密:

  1. docker build:構建 Docker 映象。
  2. docker run:執行 Docker 容器。
  3. -p 5000:5000:將容器的 5000 埠對映到主機的 5000 埠。

模型引數儲存

在上述範例中,模型引數被硬編碼到 Docker 映象中。這種做法有幾個缺點:首先,模型引數檔案可能很大,導致 Docker 映象過大;其次,如果需要更新模型引數,需要重新構建整個 Docker 映象。

為瞭解決這些問題,我們可以將模型引數儲存在外部儲存中,例如 Azure Blob Storage、Amazon S3 或 Google Cloud Storage。在容器啟動時,下載模型引數檔案。

修改 load_model() 函式

以下是修改後的 load_model() 函式:

from urllib.request import urlopen
from shutil import copyfileobj
from tempfile import NamedTemporaryFile

def load_model():
    m = CatfishModel()
    parameter_url = os.environ["CATFISH_MODEL_LOCATION"]
    with urlopen(parameter_url) as fsrc, NamedTemporaryFile() as fdst:
        copyfileobj(fsrc, fdst)
    m.load_state_dict(torch.load(fdst))
    return m

內容解密:

  1. 使用 urlopen() 下載模型引數檔案。
  2. 使用 copyfileobj() 將下載的檔案複製到臨時檔案中。
  3. 使用 torch.load() 載入模型引數。

透過這種方式,我們可以將模型引數儲存在外部儲存中,避免了 Docker 映象過大的問題,並且可以方便地更新模型引數。

佈署模型服務的最佳實踐:從 Docker 到 Kubernetes

在前面的章節中,我們已經建立了一個簡單的模型服務,並使用 Docker 將其容器化。現在,我們將繼續探討如何簡化 Docker 映像的大小,以及如何使用 Kubernetes 佈署和管理我們的模型服務。

簡化 Docker 映像大小

為了簡化 Docker 映像的大小,我們可以使用 NamedTemporaryFile() 方法來下載模型引數檔案。這樣,我們可以在下載完成後刪除檔案,從而減少 Docker 映像的大小。

import tempfile

with tempfile.NamedTemporaryFile() as tmp_file:
    # 下載模型引數檔案
    response = requests.get(model_url)
    tmp_file.write(response.content)
    # 載入模型引數
    model.load_state_dict(torch.load(tmp_file.name))

內容解密:

  1. 使用 tempfile.NamedTemporaryFile() 建立一個臨時檔案。
  2. 將下載的模型引數檔案寫入臨時檔案中。
  3. 使用 torch.load() 載入模型引數。
  4. with 區塊結束後,臨時檔案將被自動刪除。

簡化後的 Dockerfile 如下:

FROM continuumio/miniconda3:latest
ARG port
ARG host
ENV CATFISH_PORT=$port
RUN conda install -y flask \
&& conda install -c pytorch torch torchvision \
&& conda install waitress
RUN mkdir -p /app
COPY ./model.py /app
COPY ./server.py /app
COPY ./run-model-service.sh /
EXPOSE $port
ENTRYPOINT ["/run-model-service.sh"]

內容解密:

  1. 使用 continuumio/miniconda3:latest 作為基礎映像。
  2. 安裝必要的套件,包括 Flask、PyTorch 和 Waitress。
  3. 將模型程式碼和服務程式碼複製到容器中。
  4. 設定環境變數和入口點。

日誌記錄和遙測

為了監控模型的預測結果,我們需要新增日誌記錄功能。我們可以使用 logging 模組來記錄預測結果,並將其傳送到外部資源,如 Apache Kafka。

import logging
import uuid

logging.basicConfig(level=logging.INFO)

def predict():
    img_url = request.image_url
    img_tensor = open_image(BytesIO(response.content))
    start_time = time.process_time()
    prediction = model(img_tensor)
    end_time = time.process_time()
    predicted_class = CatfishClasses[torch.argmax(prediction)]
    send_to_log({
        "image": img_url,
        "prediction": predicted_class,
        "predict_tensor": prediction,
        "img_tensor": img_tensor,
        "predict_time": end_time - start_time,
        "uuid": uuid.uuid4()
    })
    return jsonify({"image": img_url, "prediction": predicted_class})

def send_to_log(log_line):
    logger.info(log_line)

內容解密:

  1. 使用 logging 模組來記錄預測結果。
  2. 將預測結果、預測張量、影像張量和預測時間等資訊記錄下來。
  3. 使用 uuid 模組生成唯一的識別碼。

佈署到 Kubernetes

Kubernetes 是一個用於自動化佈署、擴充套件和管理容器化應用程式的開源平台。我們可以使用 Kubernetes 來佈署和管理我們的模型服務。

首先,我們需要建立一個 Google Cloud 帳戶,並下載 gcloud SDK。然後,我們可以使用 gcloud 命令來建立一個新的專案,並將我們的 Docker 映像推播到 Google Container Registry。

gcloud projects create ml-k8s --set-as-default
docker build -t gcr.io/ml-k8s/catfish-service:v1 .
gcloud auth configure-docker
docker push gcr.io/ml-k8s/catfish-service:v1

內容解密:

  1. 使用 gcloud 命令建立一個新的專案。
  2. 將我們的 Docker 映像建置並標記為 gcr.io/ml-k8s/catfish-service:v1
  3. 使用 gcloud auth configure-docker 命令組態 Docker 以使用 Google Container Registry。
  4. 將我們的 Docker 映像推播到 Google Container Registry。