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

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

nvidia-smi

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

export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda-11.6/lib64
export PATH=$PATH:/usr/local/cuda-11.6/bin
source ~/.bashrc

nvcc --version

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*

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

conda info

git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

python train.py --img 640 --batch 16 --epochs 100 --data coco128.yaml --weights yolov5s.pt

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)  

pip install Flask
python3 api.py --port 5000