在GPU服务器中搭建深度学习应用后端——以YOLOv5为例

2022-04-14 12:02:41 浏览数 (1)

在GPU服务器中搭建深度学习应用后端——以YOLOv5为例

操作步骤

显卡驱动安装, Cuda安装, Cudnn配置

步骤1：显卡驱动安装（对于11.6版本cuda可以跳过此步骤，其他版本建议安装）
  1. 在英伟达驱动下载搜索你所选择的显卡和系统。我的服务器系统是ubuntu 18.04，gpu是T4。所以我如下图选择，系统选linux 64-bit即可。
  1. 点击搜索出来的驱动，再点击新页面中的download。
  1. 在新页面中对AGREE & DOWNLOAD点击右键，复制链接地址。
  1. 用ssh连接到服务器,用wget指令下载链接并安装。
代码语言：javascript复制
wget https://us.download.nvidia.com/tesla/510.47.03/NVIDIA-Linux-x86_64-510.47.03.run
sudo chmod 775 NVIDIA-Linux-x86_64-510.47.03.run
sudo sh NVIDIA-Linux-x86_64-510.47.03.run -no-x-check -no-nouveau-check -no-opengl-files
  1. 验证安装是否成功
代码语言：javascript复制
nvidia-smi
步骤2：安装cuda
  1. 到cuda下载官网按照服务器配置选择系统和安装包类型.
  1. 将页面下方的指令全部复制到ssh运行，如果上一步安装了驱动，这一步要把Driver的X取消
代码语言：javascript复制
wget https://developer.download.nvidia.com/compute/cuda/11.6.2/local_installers/cuda_11.6.2_510.47.03_linux.run
sudo sh cuda_11.6.2_510.47.03_linux.run
  1. 配置cuda的路径
代码语言：javascript复制
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-11.6/lib64
export PATH=$PATH:/usr/local/cuda-11.6/bin
source ~/.bashrc
  1. 测试安装是否成功
代码语言：javascript复制
nvcc --version
步骤3：安装和配置cudnn
  1. 到英伟达cudnn下载官网下载cudnn，过程需要登录会员。
  1. 选择自己的操作系统对应的安装包，右键复制链接地址
  1. 在服务器上下载安装包，并用FTP传到服务器上
  2. 解压安装包，将头文件和so文件拷贝到cuda目录下即完成安装
代码语言：javascript复制
sudo apt-get install zlib1g
tar -xvJf cudnn-linux-x86_64-8.4.0.27_cuda11.6-archive.tar.xz
sudo cp cudnn-linux-x86_64-8.4.0.27_cuda11.6-archive/include/cudnn*.h /usr/local/cuda/include 
sudo cp -P cudnn-linux-x86_64-8.4.0.27_cuda11.6-archive/lib/libcudnn* /usr/local/cuda/lib64 
sudo chmod a r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

安装Anaconda

安装Anaconda
代码语言：javascript复制
wget https://mirrors.tuna.tsinghua.edu.cn/anaconda/archive/Anaconda3-2021.11-Linux-x86_64.sh
sudo bash Anaconda3-2021.11-Linux-x86_64.sh
source ~/.bashrc
测试安装是否成功
代码语言：javascript复制
conda info

训练Yolov5

安装YOLOv5
代码语言：javascript复制
git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt
基于coco128数据集训练YOLOv5s模型
代码语言：javascript复制
python train.py --img 640 --batch 16 --epochs 100 --data coco128.yaml --weights yolov5s.pt

训练速度还可以的哈

搭建服务器后端

笔者搭建了一个简单的flask后端用来测试训练好的模型

代码语言：javascript复制
import argparse
import io
from PIL import Image
import torch
from flask import Flask, request, send_file
from PIL import Image


app = Flask(__name__)


DETECTION_URL = "/v1/object-detection/yolov5s"




@app.route(DETECTION_URL, methods=["POST"])
def predict():
    if not request.method == "POST":
        return


    if request.files.get("image"):
        image_file = request.files["image"]
        image_bytes = image_file.read()
        img = Image.open(io.BytesIO(image_bytes))
        results = model(img, size=640)
        results.render()
        im = Image.fromarray(results.imgs[0])
        imgByteArr = io.BytesIO()
        im.save(imgByteArr,format='PNG')
        imgByteArr = imgByteArr.getvalue()
        return send_file(io.BytesIO(imgByteArr),mimetype='image/png',as_attachment=True,attachment_filename='result.jpg')
	


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Flask API exposing YOLOv5 model")
    parser.add_argument("--port", default=5000, type=int, help="port number")
    opt = parser.parse_args()


    torch.hub._validate_not_a_forked_repo = lambda a, b, c: True


    model = torch.hub.load('/home/ubuntu/yolov5', 'custom', path='/home/ubuntu/yolov5/runs/train/exp3/weights/best.pt', source='local')
    app.run(host="0.0.0.0", port=opt.port)

主要作用是接收到图片之后，进行预测，然后返回渲染好预测框和概率的图片。

将代码复制保存为api.py 然后运行如下代码启动后端

代码语言：javascript复制
pip install Flask
python3 api.py --port 5000

测试结果

测试时将带有图片的POST请求发送至ip:5000/v1/object-detection/yolov5s

响应时间和预测结果都还不错

云服务器 GPU云服务器

0 人点赞