PyTorch Serve项目：解决Docker容器中GPU无法使用的问题

问题背景

在使用PyTorch Serve部署深度学习模型时，许多开发者会遇到Docker容器无法正确识别和使用GPU的问题。本文将以一个实际案例为基础，详细介绍如何解决这一问题。

问题现象

开发者在使用自定义Docker镜像部署PyTorch Serve服务时，发现服务虽然能够正常运行，但实际使用的是CPU而非GPU进行计算。从日志中可以看到，系统正确识别了GPU设备（NVIDIA GeForce RTX 3050 Laptop GPU），但模型推理仍然在CPU上执行。

环境配置

硬件配置

• CPU: Ryzen 7 4800H
• GPU: RTX 3050 Mobile GPU
• RAM: 24GB
• 存储: 512GB + 256GB SSD

软件环境

• 操作系统: Windows Sub-System Linux Ubuntu 20.04 (WSL集成)
• Docker: 启用WSL集成
• Python: 3.8.16
• PyTorch Serve: 0.9.0
• Torch: 2.1.2
• Torchvision: 0.16.2

问题分析

从日志中可以观察到几个关键点：

1. 系统正确识别了GPU设备
2. ONNX运行时显示已启用
3. 但实际计算仍发生在CPU上
4. 性能指标显示CPU利用率达到100%

这通常表明CUDA运行时与ONNX版本之间存在兼容性问题。

解决方案

1. 检查CUDA与ONNX版本兼容性

经过排查发现，CUDA 12与最新版ONNX运行时存在兼容性问题。解决方案是：

• 降级到CUDA 11.8
• 使用ONNX Runtime GPU 1.15版本

2. 修改Dockerfile

确保Dockerfile中正确安装了CUDA相关依赖：

FROM python:3.8.16-bullseye

RUN apt-get update && apt-get upgrade -y
RUN apt-get install -y openjdk-11-jre && apt-get clean 

RUN useradd -m model-server

ENV JAVA_HOME /usr/lib/jvm/java-11-openjdk-amd64
ADD requirements.txt .
RUN pip install -r requirements.txt

ADD start.sh config.properties  service/
ADD model-store service/model-store/

WORKDIR /service/
RUN chmod +x /service/start.sh \
    && chown -R model-server /service
RUN chown -R model-server /service/model-store
USER model-server

ENTRYPOINT [ "/service/start.sh"  ]
CMD [ "serve" ]

3. 调整requirements.txt

确保requirements.txt中包含正确的依赖版本：

imutils
torch
torchvision
onnxruntime-gpu==1.15
onnx
pyaml
torchserve
grpcio
opencv-python
scipy
torch-model-archiver
matplotlib
nvgpu

4. 模型处理

使用torch-model-archiver正确打包ONNX模型：

torch-model-archiver --model-name FaceRecognition \
--version 1.0 \
--serialized-file ./resnet_34.onnx \
--extra-files ./MyHandler.py \
--handler my_handler.py  \
--export-path model-store -f 

自定义Handler实现

对于ONNX模型，需要实现自定义Handler。以下是YOLOv8模型的Handler示例：

import logging
import onnxruntime
import torch
import cv2
import numpy as np
from PIL import Image
import os
import io
import urllib.request
from ts.torch_handler.base_handler import BaseHandler
from util_v8 import *

class FaceDetection(BaseHandler):
    def __init__(self):
        self.session = None
        self.providers = ['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']
        self.img = None
        self.blob = None

    def initialize(self, context):
        self._context = context
        self.manifest = context.manifest
        properties = context.system_properties
        model_dir = properties.get("model_dir")
        serialized_file = self.manifest['model']['serializedFile']
        model_file_path = os.path.join(model_dir, serialized_file)
        sess_options = onnxruntime.SessionOptions()
        self.session = onnxruntime.InferenceSession(model_file_path, sess_options=sess_options, providers=['CPUExecutionProvider'])

    def preprocess(self, data):
        data_get = data[0].get("data") or data[0].get("body")
        if isinstance(data_get, str):
            req = urllib.request.urlopen(data_get)
            image = Image.open(io.BytesIO(req.read()))
        else:
            byte_data = io.BytesIO(data_get)
            image = Image.open(byte_data)
        raw_data = np.array(image)
        self.img = raw_data
        img = cv2.cvtColor(raw_data, cv2.COLOR_RGB2BGR)
        im = pre_process(img)
        return im.detach().numpy()

    def inference(self, blob):
        self.blob = blob
        outputs = self.session.run(None, {self.session.get_inputs()[0].name: blob})
        return outputs[0]

    def postprocess(self, preds):
        res = []
        preds = non_max_suppression(torch.from_numpy(np.asarray(preds))[0]
        bbox = scale_boxes([640, 640], preds[:, :4], self.img.shape).round().detach().numpy()
        score = preds[:, 4].detach().numpy()
        cls = preds[:, 5].detach().numpy()
        res.append({"output": preds.tolist(), "bbox": bbox.tolist(), "label": cls.tolist(), "score": score.tolist()})
        return [res]

运行容器

确保使用--gpus参数运行容器：

docker run --rm -it -d --gpus all -p 8090:8080 -p 8091:8081 api_lanc

验证GPU使用

成功配置后，可以在日志中看到以下关键信息：

1. GPU利用率指标出现
2. 推理速度显著提升
3. CPU利用率降低

总结

PyTorch Serve在Docker容器中使用GPU时，需要注意以下几点：

1. CUDA版本与ONNX运行时的兼容性
2. 正确的Docker构建和运行参数
3. 适当的Handler实现
4. 版本匹配的依赖关系

通过上述步骤，可以确保PyTorch Serve服务能够充分利用GPU资源，提高模型推理效率。