透過 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)
內容解密:
load_model()函式:負責載入已經訓練好的模型引數。透過環境變數CATFISH_MODEL_LOCATION指定模型權重的檔案路徑,並使用torch.load()載入權重,最後將權過載入到模型中。predict()函式:處理對/predict路徑的請求,從請求中取得影像 URL,開啟影像並進行預處理(縮放、轉換為 Tensor、標準化),然後將處理後的影像輸入模型進行預測,最後傳回預測的類別和影像 URL。- 使用
unsqueeze(0):將輸入影像的 Tensor 增加一個 batch 維度,使其形狀從[3, 224, 224]變為[1, 3, 224, 224],以符合模型的輸入要求。 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
內容解密:
- 環境變數
CATFISH_MODEL_LOCATION:指定模型權重檔案的路徑,使得模型更新與服務佈署可以解耦。 torch.load(location):載入模型權重檔案。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"]
內容解密:
FROM pytorch/pytorch:latest:使用官方的 PyTorch 映象作為基礎映象,避免了從頭安裝 PyTorch 的麻煩。WORKDIR /app:設定容器中的工作目錄為/app。COPY和RUN pip install:複製requirements.txt到容器中,並安裝所需的 Python 套件。ENV:設定環境變數CATFISH_MODEL_LOCATION,指定模型權重檔案的路徑。EXPOSE 5000:暴露容器的 5000 埠,供外部存取。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"]
內容解密:
FROM continuumio/miniconda3:latest:指定基礎映象為 continuumio/miniconda3 的最新版本。ARG和ENV:定義變數,ARG用於在構建時傳遞變數,ENV用於設定環境變數。RUN:執行命令,例如安裝必要的套件。COPY:將檔案從主機複製到容器中。EXPOSE:指定容器要開放的埠。ENTRYPOINT:指定容器啟動時要執行的命令。
執行模型服務
構建好 Docker 映象後,我們可以使用以下命令執行模型服務:
docker build -t catfish-service .
docker run -p 5000:5000 catfish-service
內容解密:
docker build:構建 Docker 映象。docker run:執行 Docker 容器。-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
內容解密:
- 使用
urlopen()下載模型引數檔案。 - 使用
copyfileobj()將下載的檔案複製到臨時檔案中。 - 使用
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))
內容解密:
- 使用
tempfile.NamedTemporaryFile()建立一個臨時檔案。 - 將下載的模型引數檔案寫入臨時檔案中。
- 使用
torch.load()載入模型引數。 - 在
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"]
內容解密:
- 使用
continuumio/miniconda3:latest作為基礎映像。 - 安裝必要的套件,包括 Flask、PyTorch 和 Waitress。
- 將模型程式碼和服務程式碼複製到容器中。
- 設定環境變數和入口點。
日誌記錄和遙測
為了監控模型的預測結果,我們需要新增日誌記錄功能。我們可以使用 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)
內容解密:
- 使用
logging模組來記錄預測結果。 - 將預測結果、預測張量、影像張量和預測時間等資訊記錄下來。
- 使用
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
內容解密:
- 使用
gcloud命令建立一個新的專案。 - 將我們的 Docker 映像建置並標記為
gcr.io/ml-k8s/catfish-service:v1。 - 使用
gcloud auth configure-docker命令組態 Docker 以使用 Google Container Registry。 - 將我們的 Docker 映像推播到 Google Container Registry。