Sanic框架下部署Pytorch模型

前言

本文针对业余范围的Pytorch模型部署，类似各位想把自己开发的深度学习模型上线web端demo等等。

大家比较熟悉的Python框架主要有flask，使用flask部署上线深度学习模型过程简单，只需要在应用初始化时构造模型（model），在视图函数里调用模型前馈推理（inference）即可。这种做法的弊端在于，服务器将每个请求独立加载到GPU执行，而熟悉深度学习的朋友们都知道将输入张量堆叠成batch才能更有效利用GPU资源。

因此本文介绍的Sanic框架能便利地将多个用户请求集合成一个batch执行，执行完后再拆分开将结果返回给用户。本文将以cycleGAN的执行为例，其他model处理过程相似。

本文引用的代码来自《DeepLearning with Pytorch》，讲解为博主原创。

Sanic框架简介

Sanic框架利用了Python 3.6版本后新增的异步功能，是一个支持 async/await 语法的异步无阻塞框架。

Sanic框架与Flask框架同为Python环境下轻量级web框架，功能用法类似，已有flask基础的读者会很快适应本文，因此本文将重点介绍使用Sanic框架部署pytorch程序的思路，具体Web框架的执行细节可以参阅其他资料。Sanic框架入门指南可以在这里获得，这里也有相关函数介绍，官方文档在这里。

初始化Sanic应用

# 这里与flask框架类似，初始化一个Sanic app对象
app = Sanic(__name__)
# 定义设备
device = torch.device('cuda:0')
# 定义一些全局参数
MAX_QUEUE_SIZE = 3  # 队列最大长度，即容许等待的最大队列长度，超过则返回“too busy”
MAX_BATCH_SIZE = 2  # 单个batch内最大容量，一般取决于显存
MAX_WAIT = 1  # 从第一个batch被加入队列起的最大等待时间

构造模型

这部分我们初始化一个ModelRunner类，类中将会构造pytorch模型，构造栈等等

class ModelRunner:
    def __init__(self, model_name):
        self.model_name = model_name
        self.queue = []

        self.queue_lock = None
        # 将这里修改为构造自己的model，确保output=self.model(input)
        self.model = get_pretrained_model(self.model_name, map_location=device)
        # 这是一个事件，顾名思义，代表当下是否需要执行一批（一个batch）的推理
        self.needs_processing = None
        # 为needs_processing事件准备的计时器
        self.needs_processing_timer = None

构造一个循环执行的函数model_runner

ModelRunner内的函数，这个函数创建了协程锁和协程事件（触发后就执行推理的事件）后，就进入无限循环，事件每触发一次，就会执行一次推理。

    async def model_runner(self):
        """
        在Sanic主循环(app.loop)启动之后执行一些后台任务
        """
        # 初始化了协程锁和协程事件
        # asyncio.Lock 参见 https://docs.python.org/zh-cn/3/library/asyncio-sync.html#lock
        # asyncio.Event 参见 https://docs.python.org/zh-cn/3/library/asyncio-sync.html#asyncio.Event
        self.queue_lock = asyncio.Lock(loop=app.loop)  # 实现一个用于 asyncio 任务的互斥锁，保证对共享资源的独占访问（即访问时保证独占app.loop）
        self.needs_processing = asyncio.Event(loop=app.loop)
        logger.info("started model runner for {}".format(self.model_name))
        # 以下代码将一直执行
        while True:
            await self.needs_processing.wait()  # 等待，直至needs_processing事件被触发，即开始处理一个batch
            # 清空needs_processing事件和计时器，等待下一次用
            self.needs_processing.clear()
            if self.needs_processing_timer is not None:
                self.needs_processing_timer.cancel()
                self.needs_processing_timer = None
            async with self.queue_lock:  # 相当于等待，直到保证独占app.loop，并上协程锁，执行完后释放锁
                if self.queue:  # 队列非空，计算最大等待时间
                    longest_wait = app.loop.time() - self.queue[0]["time"]
                else:  # 队列空
                    longest_wait = None
                logger.debug(
                    "launching processing. queue size: {}. longest wait: {}".format(len(self.queue), longest_wait))
                # 确保队列长度未溢出
                to_process = self.queue[:MAX_BATCH_SIZE]
                del self.queue[:len(to_process)]
                self.schedule_processing_if_needed()
            # 将等待队列中的输入张量堆叠成一个Batch
            batch = torch.stack([t["input"] for t in to_process], dim=0)
            # 在新线程中执行模型推理
            result = await app.loop.run_in_executor(
                None, functools.partial(self.run_model, batch)
            )
            # 将结果拆开，放到队列中各个请求（字典）对应的output键值中
            for t, r in zip(to_process, result):
                t["output"] = r
                t["done_event"].set()
            del to_process

    def run_model(self, batch):
        """
        执行模型的函数，这里只需要一步，根据实际需要调整
        """
        return self.model(batch.to(device)).to('cpu')

其实这个函数是在Sanic应用被创建后就被调用的，在外层代码中有这样的安排：

app.add_task(style_transfer_runner.model_runner())

构造视图函数

为处理用户请求准备的路由函数，HTTP请求与响应的基础知识大家可以查阅其他资料。

路由函数代表了从用户发出HTTP请求后的操作，即包括接收用户以字节形式上传的图片、图片预处理、放到model中执行、返回给用户执行结果。

@app.route('/image', methods=['PUT'], stream=True)
async def image(request):
    try:
        print(request.headers)
        content_length = int(request.headers.get('content-length', '0'))
        MAX_SIZE = 2 ** 22  # 接收图片的最大大小，这里设置为：10MB
        if content_length:
            if content_length > MAX_SIZE:
                raise HandlingError("Too large")
            data = bytearray(content_length)
        else:
            data = bytearray(MAX_SIZE)
        pos = 0
        while True:
            # so this still copies too much stuff.
            data_part = await request.stream.read()
            if data_part is None:
                break
            data[pos: len(data_part) + pos] = data_part
            pos += len(data_part)
            if pos > MAX_SIZE:
                raise HandlingError("Too large")

        # 对图片流数据使用PIL打开并做必要预处理，预处理过程要根据自己的模型调整
        im = PIL.Image.open(io.BytesIO(data))
        im = torchvision.transforms.functional.resize(im, (228, 228))
        im = torchvision.transforms.functional.to_tensor(im)
        if im.dim() != 3 or im.size(0) < 3 or im.size(0) > 4:
            raise HandlingError("need rgb image")
        # 真正的核心代码，使用runner函数处理输入得到结果
        out_im = await style_transfer_runner.process_input(im)
        # 将结果使用IO流输出给用户
        out_im = torchvision.transforms.functional.to_pil_image(out_im)
        imgByteArr = io.BytesIO()
        out_im.save(imgByteArr, format='JPEG')
        return sanic.response.raw(imgByteArr.getvalue(), status=200,
                                  content_type='image/jpeg')
    except HandlingError as e:
        # 错误处理
        return sanic.response.text(e.handling_msg, status=e.handling_code)

沟通路由函数与model runner的process_input

定义在ModelRunner类中，将单个用户请求堆叠成batch，判断是否执行一个batch，等待执行后返回结果分发给用户

    async def process_input(self, input):
        """
        路由函数将请求导引至这里，这里将会把请求集合成batch（或者报告正忙
        """
        our_task = {"done_event": asyncio.Event(loop=app.loop),  # 代表这批任务是否被处理完的事件
                    "input": input,
                    "time": app.loop.time()}
        async with self.queue_lock:
            # 若队列已满则报告正忙
            if len(self.queue) >= MAX_QUEUE_SIZE:
                raise HandlingError("I'm too busy", code=503)
            # 加入队列
            self.queue.append(our_task)
            logger.debug("enqueued task. new queue size {}".format(len(self.queue)))
            # 决定是否需要安排一波process
            self.schedule_processing_if_needed()
        # 等待处理完
        await our_task["done_event"].wait()
        # 返回结果
        return our_task["output"]

判断是否需要推理一个batch的schedule_processing_if_needed

显然，当队列长度已满或达到最大等待时间，需要送一个batch进模型推理，这里的self.needs_processing.set()就代表了设置当下需要处理。

    def schedule_processing_if_needed(self):
        """
        这个函数决定是否需要安排一波模型推理
        """
        if len(self.queue) >= MAX_BATCH_SIZE:
            # 若队列长度已满，则安排一波推理（将needs_processing设为触发，即为设置需要处理这个batch）
            logger.debug("next batch ready when processing a batch")
            self.needs_processing.set()
        elif self.queue:
            # 若队列长度未满但队列非空，当队列中第一个请求达到最大等待时间时，触发needs_processing事件执行
            logger.debug("queue nonempty when processing a batch, setting next timer")
            self.needs_processing_timer = app.loop.call_at(self.queue[0]["time"] + MAX_WAIT, self.needs_processing.set)

以上是各部分的分别介绍，整体代码如下

import sys
import asyncio
import itertools
import functools
from sanic import Sanic
from sanic.response import json, text
from sanic.log import logger
from sanic.exceptions import ServerError

import sanic
import threading
import PIL.Image
import io
import torch
import torchvision
from .cyclegan import get_pretrained_model

# 这里与flask框架类似，初始化一个Sanic app对象
app = Sanic(__name__)

device = torch.device('cuda:0')
# 定义一些全局参数
MAX_QUEUE_SIZE = 3  # 队列最大长度，即容许等待的最大队列长度，超过则返回“too busy”
MAX_BATCH_SIZE = 2  # 单个batch内最大容量，一般取决于显存
MAX_WAIT = 1  # 从第一个batch被加入队列起的最大等待时间


class HandlingError(Exception):
    def __init__(self, msg, code=500):
        super().__init__()
        self.handling_code = code
        self.handling_msg = msg


class ModelRunner:
    def __init__(self, model_name):
        self.model_name = model_name
        self.queue = []

        self.queue_lock = None
        # 将这里修改为构造自己的model，确保output=self.model(input)
        self.model = get_pretrained_model(self.model_name, map_location=device)
        # 这是一个事件，顾名思义，代表当下是否需要执行一批（一个batch）的推理
        self.needs_processing = None
        # 为needs_processing事件准备的计时器
        self.needs_processing_timer = None

    def schedule_processing_if_needed(self):
        """
        这个函数决定是否需要安排一波模型推理
        """
        if len(self.queue) >= MAX_BATCH_SIZE:
            # 若队列长度已满，则安排一波推理（将needs_processing设为触发，即为设置需要处理这个batch）
            logger.debug("next batch ready when processing a batch")
            self.needs_processing.set()
        elif self.queue:
            # 若队列长度未满但队列非空，当队列中第一个请求达到最大等待时间时，触发needs_processing事件执行
            logger.debug("queue nonempty when processing a batch, setting next timer")
            self.needs_processing_timer = app.loop.call_at(self.queue[0]["time"] + MAX_WAIT, self.needs_processing.set)

    async def process_input(self, input):
        """
        路由函数将请求导引至这里，这里将会把请求集合成batch（或者报告正忙
        """
        our_task = {"done_event": asyncio.Event(loop=app.loop),  # 代表这批任务是否被处理完的事件
                    "input": input,
                    "time": app.loop.time()}
        async with self.queue_lock:
            # 若队列已满则报告正忙
            if len(self.queue) >= MAX_QUEUE_SIZE:
                raise HandlingError("I'm too busy", code=503)
            # 加入队列
            self.queue.append(our_task)
            logger.debug("enqueued task. new queue size {}".format(len(self.queue)))
            # 决定是否需要安排一波process
            self.schedule_processing_if_needed()
        # 等待处理完
        await our_task["done_event"].wait()
        # 返回结果
        return our_task["output"]

    def run_model(self, batch):
        """
        执行模型的函数，这里只需要一步，根据实际需要调整
        """
        return self.model(batch.to(device)).to('cpu')

    async def model_runner(self):
        """
        在Sanic主循环(app.loop)启动之后执行一些后台任务
        """
        # 初始化了协程锁和协程事件
        # asyncio.Lock 参见 https://docs.python.org/zh-cn/3/library/asyncio-sync.html#lock
        # asyncio.Event 参见 https://docs.python.org/zh-cn/3/library/asyncio-sync.html#asyncio.Event
        self.queue_lock = asyncio.Lock(loop=app.loop)  # 实现一个用于 asyncio 任务的互斥锁，保证对共享资源的独占访问（即访问时保证独占app.loop）
        self.needs_processing = asyncio.Event(loop=app.loop)
        logger.info("started model runner for {}".format(self.model_name))
        # 以下代码将一直执行
        while True:
            await self.needs_processing.wait()  # 等待，直至needs_processing事件被触发，即开始处理一个batch
            # 清空needs_processing事件和计时器，等待下一次用
            self.needs_processing.clear()
            if self.needs_processing_timer is not None:
                self.needs_processing_timer.cancel()
                self.needs_processing_timer = None
            async with self.queue_lock:  # 相当于等待，直到保证独占app.loop，并上协程锁，执行完后释放锁
                if self.queue:  # 队列非空，计算最大等待时间
                    longest_wait = app.loop.time() - self.queue[0]["time"]
                else:  # 队列空
                    longest_wait = None
                logger.debug(
                    "launching processing. queue size: {}. longest wait: {}".format(len(self.queue), longest_wait))
                # 确保队列长度未溢出
                to_process = self.queue[:MAX_BATCH_SIZE]
                del self.queue[:len(to_process)]
                self.schedule_processing_if_needed()
            # 将等待队列中的输入张量堆叠成一个Batch
            batch = torch.stack([t["input"] for t in to_process], dim=0)
            # 在新线程中执行模型推理
            result = await app.loop.run_in_executor(
                None, functools.partial(self.run_model, batch)
            )
            # 将结果拆开，放到队列中各个请求（字典）对应的output键值中
            for t, r in zip(to_process, result):
                t["output"] = r
                t["done_event"].set()
            del to_process


style_transfer_runner = ModelRunner(sys.argv[1])


@app.route('/image', methods=['PUT'], stream=True)
async def image(request):
    try:
        print(request.headers)
        content_length = int(request.headers.get('content-length', '0'))
        MAX_SIZE = 2 ** 22  # 接收图片的最大大小，这里设置为：10MB
        if content_length:
            if content_length > MAX_SIZE:
                raise HandlingError("Too large")
            data = bytearray(content_length)
        else:
            data = bytearray(MAX_SIZE)
        pos = 0
        while True:
            # so this still copies too much stuff.
            data_part = await request.stream.read()
            if data_part is None:
                break
            data[pos: len(data_part) + pos] = data_part
            pos += len(data_part)
            if pos > MAX_SIZE:
                raise HandlingError("Too large")

        # 对图片流数据使用PIL打开并做必要预处理，预处理过程要根据自己的模型调整
        im = PIL.Image.open(io.BytesIO(data))
        im = torchvision.transforms.functional.resize(im, (228, 228))
        im = torchvision.transforms.functional.to_tensor(im)
        if im.dim() != 3 or im.size(0) < 3 or im.size(0) > 4:
            raise HandlingError("need rgb image")
        # 真正的核心代码，使用runner函数处理输入得到结果
        out_im = await style_transfer_runner.process_input(im)
        # 将结果使用IO流输出给用户
        out_im = torchvision.transforms.functional.to_pil_image(out_im)
        imgByteArr = io.BytesIO()
        out_im.save(imgByteArr, format='JPEG')
        return sanic.response.raw(imgByteArr.getvalue(), status=200,
                                  content_type='image/jpeg')
    except HandlingError as e:
        # 错误处理
        return sanic.response.text(e.handling_msg, status=e.handling_code)

# 在Sanic主循环(app.loop)启动之后执行一些后台任务
app.add_task(style_transfer_runner.model_runner())
# 启动Sanic应用
app.run(host="0.0.0.0", port=8000, debug=True)