文章转载于微信公众号:GiantPandaCV
作者: 阿呆

 

【GiantPandaCV导读】这篇文章包含与PyTorch模型部署相关的两部分内容：

• PyTorch-YOLOv3模型的Web页面展示程序的编写
• 模型的服务接口相关工具的使用

0. 环境依赖:

系统：Ubuntu 18.04
Python版本：3.7
依赖Python包：1. PyTorch==1.3 2. Flask==0.12 3. Gunicorn

需要注意的是Flask 0.12中默认的单进程单线程，而最新的1.0.2则不是（具体是多线程还是多进程尚待考证），而中文博客里面能查到的资料基本都在说Flask默认单进程单线程。

依赖工具 1. nginx 2. apache2-utils

nginx 用于代理转发和负载均衡，apache2-utils用于测试接口

---

1. 制作模型演示界面

图像识别任务的展示这项工程一般是面向客户的，这种场景下不可能把客户拉到你的电脑前面，敲一行命令，等matplotlib弹个结果窗口出来。总归还是要有个图形化界面才显得有点诚意。
为了节约时间，我们选择了Flask框架来开发这个界面。

上传页面和展示页面

做识别演示需要用到两个html页面，代码也比较简单，编写如下：

上传界面

              

使用Flask上传本地图片

                                  

展示界面

              

使用Flask上传本地图片

                                      

上传界面如下图所示，觉得丑的话可以找前端同事美化一下：

 

flask上传图片及展示功能

然后就可以编写flask代码了，为了更好地展示图片，可以向html页面传入图片地址参数。

from flask import Flask, render_template, request, redirect, url_for, make_response, jsonifyfrom werkzeug.utils import secure_filenameimport osimport cv2import timefrom datetime import timedeltafrom main import run, confALLOWED_EXTENSIONS = set([    "png","jpg","JPG","PNG", "bmp"])def is_allowed_file(filename):    return '.' in filename and filename.rsplit('.', 1)[1] in ALLOWED_EXTENSIONSapp = Flask(__name__)# 静态文件缓存过期时间app.send_file_max_age_default = timedelta(seconds=1)@app.route("/upload",methods = ['POST', 'GET'])def upload():    if request.method == "POST":        f = request.files['file']        if not ( f and is_allowed_file(f.filename)):            return jsonify({                "error":1001, "msg":"请检查上传的图片类型，仅限于png、PNG、jpg、JPG、bmp"            })        user_input = request.form.get("name")        basepath = os.path.dirname(__file__)        upload_path = os.path.join(basepath, "static/images",secure_filename(f.filename))        f.save(upload_path)                detected_path = os.path.join(basepath, "static/images", "output" + secure_filename(f.filename))        run(upload_path, conf, detected_path)        # return render_template("upload_ok.html", userinput = user_input, val1=time.time(), path = detected_path)        path = "/images/" + "output" + secure_filename(f.filename)        return render_template("upload_ok.html", path = path, val1 = time.time())    return render_template("upload.html")if __name__ == "__main__":    app.run(host='0.0.0.0', port=8888, debug=True)

目标检测函数

原项目中提供了detection.py来做批量的图片检测，需要稍微修改一下才能用来做flask代码中的接口。

from __future__ import divisionfrom models import *from utils.utils import *from utils.datasets import *import osimport sysimport timeimport datetimeimport argparsefrom PIL import Imageimport torchfrom torchvision import datasetsfrom torch.autograd import Variableimport matplotlib.pyplot as pltimport matplotlib.patches as patchesfrom matplotlib.ticker import NullLocatorclass custom_dict(dict):    def __init__(self, d = None):        if d is not None:            for k,v in d.items():                self[k] = v        return super().__init__()    def __key(self, key):        return "" if key is None else key.lower()    def __str__(self):        import json        return json.dumps(self)    def __setattr__(self, key, value):        self[self.__key(key)] = value    def __getattr__(self, key):        return self.get(self.__key(key))    def __getitem__(self, key):        return super().get(self.__key(key))    def __setitem__(self, key, value):        return super().__setitem__(self.__key(key), value)conf = custom_dict({    "model_def":"config/yolov3.cfg",    "weights_path":"weights/yolov3.weights",    "class_path":"data/coco.names",    "conf_thres":0.8,    "nms_thres":0.4,    "img_size":416})def run(img_path, conf, target_path):    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")    os.makedirs("output", exist_ok=True)    classes = load_classes(conf.class_path)    model = Darknet(conf.model_def, img_size=conf.img_size).to(device)    if conf.weights_path.endswith(".weights"):        # Load darknet weights        model.load_darknet_weights(conf.weights_path)    else:        # Load checkpoint weights        model.load_state_dict(torch.load(conf.weights_path))    model.eval()         img = Image.open(img_path).convert("RGB")    img = img.resize(((img.size[0] // 32) * 32, (img.size[1] // 32) * 32))    img_array = np.array(img)    img_tensor = pad_to_square(transforms.ToTensor()(img),0)[0].unsqueeze(0)    conf.img_size = img_tensor.shape[2]        with torch.no_grad():        detections = model(img_tensor)        detections = non_max_suppression(detections, conf.conf_thres, conf.nms_thres)[0]    cmap = plt.get_cmap("tab20b")    colors = [cmap(i) for i in np.linspace(0, 1, 20)]    plt.figure()    fig, ax = plt.subplots(1)    ax.imshow(img_array)    if detections is not None:        # Rescale boxes to original image        detections = rescale_boxes(detections, conf.img_size, img_array.shape[:2])        unique_labels = detections[:, -1].cpu().unique()        n_cls_preds = len(unique_labels)        bbox_colors = random.sample(colors, n_cls_preds)        for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:            print("/t+ Label: %s, Conf: %.5f" % (classes[int(cls_pred)], cls_conf.item()))            box_w = x2 - x1            box_h = y2 - y1            color = bbox_colors[int(np.where(unique_labels == int(cls_pred))[0])]            # Create a Rectangle patch            bbox = patches.Rectangle((x1, y1), box_w, box_h, linewidth=2, edgecolor=color, facecolo)            # Add the bbox to the plot            ax.add_patch(bbox)            # Add label            plt.text(                x1,                y1,                s=classes[int(cls_pred)],                colo,                verticalalignmen,                bbox={"color": color, "pad": 0},            )    # Save generated image with detections    plt.axis("off")    plt.gca().xaxis.set_major_locator(NullLocator())    plt.gca().yaxis.set_major_locator(NullLocator())    filename = img_path.split("/")[-1].split(".")[0]    plt.savefig(target_path, bbox_inches='tight', pad_inches=0.0)    plt.close()if __name__ == "__main__":    run("data/samples/dog.jpg",conf)

展示效果

编写好了之后，启动server.py，在本地打开localhost:8888/upload就可以看到如下界面了，把图片上传上去，很快就能得到检测结果。
结果如下图所示：

 

2. 深度学习的服务接口编写

接下来介绍的是在生产环境下的部署，使用的是flask+gunicorn+nginx的方式，可以处理较大规模的请求。
下面以图像分类模型为例演示一下深度学习服务接口如何编写。

对于深度学习工程师来说，学习这些内容主要是了解一下自己的模型在生产环境的运行方式，便于在服务出现问题的时候与开发的同事一起进行调试。

flask服务接口

接口不需要有界面显示，当然也可以添加一个API介绍界面，方便调用者查看服务是否已经启动。

from flask import Flask, requestfrom werkzeug.utils import secure_filenameimport uuidfrom PIL import Imageimport osimport timeimport base64import jsonimport torchfrom torchvision.models import resnet18from torchvision.transforms import ToTensorfrom keys import keyapp = Flask(__name__)net = resnet18(pretrained=True)net.eval()@app.route("/",methods=["GET"])def show():    return "classifier api"@app.route("/run",methods = ["GET","POST"])def run():    file = request.files['file']    base_path = os.path.dirname(__file__)    if not os.path.exists(os.path.join(base_path, "temp")):        os.makedirs(os.path.join(base_path, "temp"))    file_name = uuid.uuid4().hex    upload_path = os.path.join(base_path, "temp", file_name)    file.save(upload_path)    img = Image.open(upload_path)    img_tensor = ToTensor()(img).unsqueeze(0)    out = net(img_tensor)    pred = torch.argmax(out,dim = 1)    return "result : {}".format(key[pred])if __name__ == "__main__":    app.run(hos,port=5555,debug=True)

在命令行输入python server.py即可启动服务。

gunicorn启动多个实例

新版的flask已经支持多进程了，不过用在生产环境还是不太稳定，一般生产环境会使用gunicorn来启动多个服务。
使用如下命令即可启动多个图像分类实例

gunicorn -w 4 -b 0.0.0.0:5555 server:app

输出如下内容代表服务创建成功：

[2020-02-11 14:50:24 +0800] [892] [INFO] Starting gunicorn 20.0.4[2020-02-11 14:50:24 +0800] [892] [INFO] Listening at: http://0.0.0.0:5555 (892)[2020-02-11 14:50:24 +0800] [892] [INFO] Using worker: sync[2020-02-11 14:50:24 +0800] [895] [INFO] Booting worker with pid: 895[2020-02-11 14:50:24 +0800] [896] [INFO] Booting worker with pid: 896[2020-02-11 14:50:24 +0800] [898] [INFO] Booting worker with pid: 898[2020-02-11 14:50:24 +0800] [899] [INFO] Booting worker with pid: 899

如果配置比较复杂，也可以将配置写入一个文件中，如：

bind = '0.0.0.0:5555'timeout = 10workers = 4

然后运行：

gunicorn -c gunicorn.conf sim_server:app

nginx负载均衡

如果有多个服务器，可以使用nginx做请求分发与负载均衡。
安装好nginx之后，修改nginx的配置文件

worker_processes auto;error_log /var/log/nginx/error.log;pid /run/nginx.pid;# Load dynamic modules. See /usr/share/nginx/README.dynamic.include /usr/share/nginx/modules/*.conf;events {    worker_connections 1024;}http {    server    {        listen 5556; # nginx端口        server_name localhost;        location / {            proxy_pass http://localhost:5555/run; # gunicorn的url        }    }}

然后按配置文件启动

sudo nginx -c nginx.conf

测试一下服务是否正常

启动了这么多服务之后，可以使用apache2-utils来测试服务的并发性能。
使用apache2-utils进行上传图片的post请求方法参考：
https://gist.github.com/chill...

严格参照，注意一个标点，一个符号都不要错。使用这种方法传输图片的base64编码，在服务端不需要解码也能使用

然后使用下面的方式访问
gunicorn 接口

ab -n 2 -c 2 -T "multipart/form-data; boundary=1234567890" -p turtle.txt http://localhost:5555/run

nginx 接口

ab -n 2 -c 2 -T "multipart/form-data; boundary=1234567890" -p turtle.txt http://localhost:5556/run

- END -

推荐阅读

• Pytorch量化感知训练详解
• PyTorch OCR模型的安卓端部署
• 深度学习量化技术科普

更多嵌入式AI技术干货请关注嵌入式AI专栏。

审核编辑：符乾江