From: https://zhuanlan.zhihu.com/p/35879835
随着深度学习越来越火，各种框架也层出不穷，如何训练一个深度学习模型变得越来越简单。然而在实际的工业场景中，我们往往更加关注如何部署一个已经训练好的模型。
在这一点上，tensorflow做得非常好，提供了tensorflow serving来帮助我们非常方便地部署到工业场景下。众所周知，PyTorch的一个非常大的劣势就是没有办法很方便地部署模型，facebook和Microsoft一起搞了一个神经交换机，ONNX，可以将pytorch model转换到Caffe2 model，这样一是麻烦，二是Caffe2目前还在测试，一堆bug，用的人也不多，三是还要多学一个框架Caffe2。所以这并不是一个非常好的选择。
目前的最新消息，Caffe2的源码已经并到pytorch中了，或许这是Facebook准备对付TensorFlow的大招，我们拭目以待。
本文受到keras的一篇博文的启发，会教大家如何使用flask来部署训练好的pytorch模型。首先声明一下，flask我也不太会用，因为看到了keras的文章，希望分享这种思路和想法，使用一种web框架实现深度学习模型的部署。
环境配置
首先确保安装了pytorch，因为需要使用flask这个web框架，所以当然需要安装flask，非常简单，使用下面的命令进行安装。
pip install flask
配置REST API
我们知道每次启动模型，load参数是一件非常费时间的事情，而每次做前向传播的时候模型其实都是一样的，所以我们最好的办法就是load一次模型，然后做完前向传播之后仍然保留这个load好的模型，下一次有新的数据进来，我们就可以不用重新load模型，可以直接做前向传播得到结果，这样无疑节约了很多load模型的时间。所以我们需要建立一个类似于服务器的机制，将模型在服务器上load好，方便我们不断去调用模型做前向传播，那么怎么能够达到这个目的呢？我们可以使用flask来建立一个REST API来达到这一目的。
REST API 是什么呢？REST 是Representational State Transfer的缩写，这是一种架构风格，这里就不再过多描述，感兴趣的同学可以自己去google一下。
那么如何用flask启动这个服务呢？
载入模型
app = flask.Flask(__name__)
model = None
use_gpu = True
def load_model():
"""Load the pre-trained model, you can use your model just as easily.
"""
global model
model = resnet50(pretrained=True)
model.eval()
if use_gpu:
model.cuda()
首先我们需要使用上面的代码来载入模型，前面三句话的非常简单，第一句话表示调用flask初始化一个app，接着定义一个变量model来表示模型，use_gpu表示是否使用gpu。
接着定义 load_model 这个函数，在函数中将模型的参数load到前面定义的model中，这里使用了resnet50，记得使用 .eval将model转换成eval模式，如果要使用GPU，则加上 .cuda()
数据预处理
def prepare_image(image, target_size):
"""Do image preprocessing before prediction on any data.
:param image: original image
:param target_size: target image size
:return:
preprocessed image
"""
if image.mode != 'RGB':
image = image.convert("RGB")
# Resize the input image nad preprocess it.
image = T.Resize(target_size)(image)
image = T.ToTensor()(image)
# Convert to Torch.Tensor and normalize.
image = T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(image)
# Add batch_size axis.
image = image[None]
if use_gpu:
image = image.cuda()
return torch.autograd.Variable(image, volatile=True)
这里就是 pytorch 中标准的预处理流程，首先将图片resize到固定的大小，然后转换成 tensor，接着做标准化。
启动REST API
定义好了模型和数据预处理，接下来我们就需要开始启动 flask 服务了。
@app.route("/predict", methods=["POST"])
def predict():
# Initialize the data dictionary that will be returned from the view.
data = {"success": False}
# Ensure an image was properly uploaded to our endpoint.
if flask.request.method == 'POST':
if flask.request.files.get("image"):
# Read the image in PIL format
image = flask.request.files["image"].read()
image = Image.open(io.BytesIO(image))
# Preprocess the image and prepare it for classification.
image = prepare_image(image, target_size=(224, 224))
# Classify the input image and then initialize the list of predictions to return to the client.
preds = F.softmax(model(image), dim=1)
results = torch.topk(preds.cpu().data, k=3, dim=1)
data['predictions'] = list()
# Loop over the results and add them to the list of returned predictions
for prob, label in zip(results[0][0], results[1][0]):
label_name = idx2label[label]
r = {"label": label_name, "probability": float(prob)}
data['predictions'].append(r)
# Indicate that the request was a success.
data["success"] = True
# Return the data dictionary as a JSON response.
return flask.jsonify(data)
首先定义请求方式为 POST，表示向服务器传输数据，接着定义一个 predict 函数来进行模型的前向传播。
在 prediect 中，首先建立一个字典 data 来存储请求状态，初始化为 false。接着通过 flask.request.method 来判断是否是 POST 请求，如果是的话，我们就通过 flask.request.files.get("image") 来判断是否能够得到从远端传过来的数据，如果确实有数据传过来，我们就可以通过 flask.request.files["image"].read() 来得到从远方 POST 上来的数据。
为了传输的速度考虑，一般都会传二进制的文件，所以通过 io.BytesIO(image) 将二进制的文件读取出来，再通过 PIL.Image.open 来读取这个图片，这样我们就解码了一张从远端传过来的图片了。
然后下面的操作就很简单了，首先通过 prepare_image 将图片做预处理，接着传入到网络当中，这里需要注意我们会使用 F.softmax 将模型的输出得分转换成一个概率分布，因为我们想要输出 top3 的结果和置信概率。最后我们就是将结果存到 data 中，返回成 json 的文件。
最后我们在 main 函数中调用
load_model()
app.run()
就可以启动 flask 服务了。通过上面的代码，我们知道了如何处理传过来的图片并输出预测的结果，那么我们如何传图片呢？这就是下面会讲的如何发起数据请求。
发送数据请求
发送数据请求并不难，首先我们需要知道上面定义好的 flask server 的地址，因为这就是我们在本地定义的，所以地址是
PyTorch_REST_API_URL = 'http://127.0.0.1:5000/predict'
上面的 /predict 是因为我们前面使用了 @app.route("/predict", methods=["POST"])。
接着我们定义一个函数来发送数据请求
def predict_result(image_path):
# Initialize image path
image = open(image_path, 'rb').read()
payload = {'image': image}
# Submit the request.
r = requests.post(PyTorch_REST_API_URL, files=payload).json()
# Ensure the request was successful.
if r['success']:
# Loop over the predictions and display them.
for (i, result) in enumerate(r['predictions']):
print('{}. {}: {:.4f}'.format(i + 1, result['label'],
result['probability']))
# Otherwise, the request failed.
else:
print('Request failed')
传入的参数 image_path 是图片路径，然后使用 requests.post(PyTorch_REST_API_URL, files=payload).json() 向服务器传入数据，同时得到服务器计算的结果，最后将结果 print 出来就可以了。
实验结果
我们使用 ResNet50 作为预训练的模型，传入下面这张图片作为测试
首先在一个终端中运行
python run_pytorch_server.py
来启动 flask server，等待一会儿，可以得到下面的结果
然后我们重新打开一个新的终端，运行下面的代码
python simple_request.py --file='./dog.jpg'
这里的 dog.jpg 可以改成你自己的文件路径，然后我们可以得到下面的结果
讨论
最后我们实现了一个简单的深度学习服务器，当然这个模型是在本地建立的，我们当然可以将模型建立到远端的服务器上，本地向远端发送请求。当然这只是一个 toy model，我们可以基于这种思想设计更加复杂的结构。
本文内容参考自 Building a simple Keras + deep learning REST API
本文的完整code
欢迎关注原作者的知乎专栏深度炼丹
欢迎访问原作者的博客