首页
/ PyTorch Serve项目:解决Docker容器中GPU无法使用的问题

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

2025-06-14 14:27:14作者:胡易黎Nicole

问题背景

在使用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资源,提高模型推理效率。

登录后查看全文
热门项目推荐
相关项目推荐