不像Tensorflow自带了Tensorflow Server，Pytorch官方没有自带部署应用。因此如果想要把Pytorch用于生产环境，需要自己搭建Pytorch服务器。考虑到便利，本文直接利用Flask—一个轻量的Python服务器框架。
利用Flask传递Numpy Array
因为在整个Pytorch使用过程中，最核心的数据结构就是Tensor，基本构成就是多维数组，所以先实现传递Numpy Array。
Client
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 | import requests import numpy as np import json addr = \\'http://localhost:5000\\' test_url = addr \\'/api/test\\' content_type = \\'application/json\\' headers = {\\'content-type\\': content_type} temp = np.zeros((2, 4)) 0.1 temp = temp.tolist() data = {\\'data\\': temp} response = requests.post(test_url, json=json.dumps(data), headers=headers) print(response.text) |
Server
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 | from flask import Flask, request, Response import numpy as np import json app = Flask(__name__) @app.route(\\'/api/test\\', methods=[\\'POST\\']) def test(): r = request.json r_json = json.loads(r) data = r_json[\\'data\\'] numpy_data = np.asarray(data) response = {\\'message\\': \\'Data type:{},Shape:{}\\'.format(type(numpy_data), numpy_data.shape)} response_pickled = json.dumps(response) return Response(response=response_pickled, status=200, mimetype="application/json") # start flask app app.run(host="0.0.0.0", port=5000) |
然后先运行服务端，再运行客户端。可以得到：
1 | {"message": "Data type:,Shape:(2, 4)"} |
说明我们成功通过网络传递了一个Numpy array。
加入Pytorch模型
模型训练与本文无关，故不做阐述。随便拿一个模型举例。该模型的功能为识别图中形状为正方形还是圆形，输入为3通道32*32的图像向量。文件结构如下：
-Model
—model.py
—model.pkl
随机形状生成器
先写一个随机生成正方形/圆形的模块。
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | import cv2 import numpy as np import random import matplotlib.pyplot as plt def imshow(img): fig, ax = plt.subplots() fig.set_size_inches(5, 5) ax.axis("off") plt.imshow(img/255) plt.show() class Shape: def __init__(self): self.colors = [ (0, 0, 255), # r (0, 255, 0), # g (255, 0, 0), # b (0, 156, 255), # o (128, 128, 128), # gray (0, 255, 255) # cyan ] self.canvas_size = 100 def make(self,model): img = np.zeros((self.canvas_size, self.canvas_size, 3), dtype=np.float32) * 255 color = self.colors[random.randint(0, 5)] center = [random.randint(40, 60), random.randint(40, 60)] object_size = random.randint(10, 40) if model == \\'retrangle\\': start = (center[0] - object_size, center[1] - object_size) # painting start point end = (center[0] object_size, center[1] object_size) cv2.rectangle(img, start, end, color, -1) if model == \\'circle\\': cv2.circle(img, (center[0], center[1]), object_size, color, -1) img = img return img/255 creator=Shape() circle=creator.make(\\'circle\\') imshow(circle) retrangle=creator.make(\\'retrangle\\') imshow(retrangle) |
运行可以看到生成了圆形和正方形。
测试下模型
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | import PIL import torch from model import model from shape_maker import Shape from torch.autograd import Variable import torchvision.transforms as transforms # Load structure net = model.ShapeDetectNetwork() # Load para net.load_state_dict(torch.load(\'./model/shapeDetect\', map_location=lambda storage, loc: storage)) # Create two shapes creator = Shape() circle = creator.make(\'circle\') retrangle = creator.make(\'retrangle\') # transform img array to pytorch tensor def array2tensor(img): img = PIL.Image.fromarray(img.astype(\'uint8\')) trans = transforms.Compose([ transforms.Resize(32), transforms.ToTensor(), ]) img = trans(img) return img # 2 classes based on training set classes = [\'circle\', \'retrangle\'] o1 = net(Variable(array2tensor(circle).unsqueeze(0))) o2 = net(Variable(array2tensor(retrangle).unsqueeze(0))) _, idx1 = torch.max(o1.data, 1) _, idx2 = torch.max(o2.data, 1) print(classes[idx1[0]], classes[idx2[0]]) |
可以得到:
1 | circle retrangle |
模型能够正常运作。接下来尝试把模型部署到服务端，客户端向服务端传送图像。
Client
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 | import requests from shape_maker import Shape import json # Set server address addr = \'http://localhost:5000\' test_url = addr \'/api/test\' # Set post header content_type = \'application/json\' headers = {\'content-type\': content_type} # Create 2 shapes creator = Shape() circle = creator.make(\'circle\') retrangle = creator.make(\'retrangle\') # Transform numpy array to list circle = circle.tolist() retrangle = retrangle.tolist() # wrap them into json json_f1 = json.dumps({\'data\': retrangle}) json_f2 = json.dumps({\'data\': circle}) # post request response1 = requests.post(test_url, json=json_f1, headers=headers) print(response1.text) response2 = requests.post(test_url, json=json_f2, headers=headers) print(response2.text) |
Server
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | from flask import Flask, request, Response from utilities import array2tensor from model import model import torch from torch.autograd import Variable import json import numpy as np app = Flask(__name__) @app.route(\'/api/test\', methods=[\'POST\']) def test(): net = model.ShapeDetectNetwork() net.load_state_dict(torch.load(\'./model/shapeDetect\', map_location=lambda storage, loc: storage)) r = request.json r_json = json.loads(r) data = r_json[\'data\'] numpy_data = np.asarray(data) o = net(Variable(array2tensor(numpy_data).unsqueeze(0))) classes = [\'circle\', \'retrangle\'] _, idx = torch.max(o.data, 1) shape = classes[idx[0]] # response response = { \'message\': \'The shape is {}\'.format(shape) } # encode response using jsonpickle response_pickled = json.dumps(response) return Response(response=response_pickled, status=200, mimetype="application/json") # start flask app app.run(host="0.0.0.0", port=5000) |
同样先运行服务端，再运行客户端。可以得到：
1 2 | {"message": "The shape is retrangle"} {"message": "The shape is circle"} |
实验成功！
源码链接：https://github.com/nofacer/Flask_Pytorch_Server
文章来源：Dustni知乎专栏
本站微信群、QQ群(三群号 726282629)：
requests.exceptions.ConnectionError: HTTPConnectionPool(host=’localhost’, port=5000): Max retries exceeded with url: /api/test (Caused by NewConnectionError(‘: Failed to establish a new connection: [WinError 10061] 由于目标计算机积极拒绝，无法连接。’,))___请问这是什么问题呀？