Torchserve打包和部署你训练的深度学习模型

一、Torchserve介绍

Torchserve是Facebooke公司开发的在线深度学习模型部署框架，它可以很方便的部署pytorch的深度学习模型，读者可以访问Github地址获取最新功能和详细说明：官方地址https://github.com/pytorch/serve/blob/master/docs/README.md。

我们已经在文章Ubuntu配置Torchserve环境，并在线发布你的深度学习模型中描述了怎么一步步安装-部署深度学习模型。读者可以自行查阅。在上文中，我们利用的是Torchserve提供的官方模型mnist例子作为演示。在本文中，我会描述怎么一步步打包和部署自己的深度学习模型。下面，开始正文！

二、打包你的深度学习模型

主要采用的是pytorch框架来训练你的深度学习模型。看这个文章的人，肯定已经知道什么pytorch了，我就不废话了。我主要是利用pytorch进行CT图像的分类，采用了github上大佬编写的框架(github地址),喜欢可以给他个star。我同时也针对这个框架进行了增添，新的代码我也会贡献在Github上，敬请期待吧！

我采用的深度学习模型是resnet101，使用预训练模型进行迁移学习。训练完成之后，得到了模型参数resnet101.pth，这个文件中包含了模型的全部权重参数，简言之，这就训练后针对CT图像分类的模型。

接下来，你要想部署到Tochserve上，你必须把模型文件resnet101.pth转换为torch脚本，然后将上文转换完成的脚本序列化为pt模型文件。具体做法如下所示：
1.定义一个python文件，输入resnet101.pth，输出resnet101.pt

import torch
import cfg # cfg是参数的预定于文件

def load_checkpoint(filepath):
    checkpoint = torch.load(filepath, map_location='cpu') 
    model = checkpoint['model']  # 提取网络结构
    model.load_state_dict(checkpoint['model_state_dict'])  # 加载网络权重参数
    for parameter in model.parameters():
        parameter.requires_grad = False
    model.eval()
    return model
    
if __name__ == "__main__":
    #利用trace把模型转化为pt
    trained_model = cfg.TRAINED_MODEL #cfg.TRAINED_MODEL表示resnet101.pth所在的位置
    model = load_checkpoint(trained_model)
    example = torch.rand(1, 3, 224, 224)
    traced_script_module = torch.jit.trace(model, example)
    traced_script_module.save('resnet101.pt')
    output = traced_script_module(torch.ones(1, 3, 224, 224))
    print(output)

Torch脚本是一种从PyTorch代码创建可序列化和可优化模型的方法。用Torch脚本编写的代码可以从Python进程中保存，并在没有Python依赖的进程中加载。我们提供了一些工具帮助我们将模型从纯Python程序逐步转换为可以独立于Python运行的Torch脚本程序。Torch脚本程序可以在其他语言的程序中运行（例如，在独立的C ++程序中）。这使得我们可以使用熟悉的工具在PyTorch中训练模型，而将模型导出到出于性能和多线程原因不能将模型作为Python程序运行的生产环境中去，参见文章PyTorch 1.0 中文文档：Torch 脚本。

运行上述文件，得到resnet101.pt文件。

*2.准备打包文件
我们已经在文章Ubuntu配置Torchserve环境，并在线发布你的深度学习模型讲述了打包模型的方法，但是，当我们打包自己模型的时候，其中有几个最重要的文件需要我们自己定义。如下所示：

torch-model-archiver --model-name resnet101  --version 1.0  --serialized-file resnet101.pt  --extra-files index_to_name.json,MyHandler.py  --handler my_handler.py

详细解释如下：

• –model-name: 模型的名称，这是我们自己定义的，就是为你的模型起个名字，我们就称为resnet101。
• –model-file: 模型采用的类别，即模型的网络结果，我们采用的是resnet101，所以可以编写一个python文件，定义一个模型类，命名为model.py：

from torchvision.models.resnet import ResNet, Bottleneck
class ResNet101ImageClassifier(ResNet):
    def __init__(self):
        super(ResNet101ImageClassifier, self).__init__(Bottleneck, [3, 4, 23, 3])

• –serialized-file: 直接调用我们训练好的resnet101.pt文件。
• –handler: 打包模型后模型的执行流程，主要包括预处理、预测和后处理的流程，建立一个my_handler.py文件，来调用我们自定义的MyHandler.py，读者也可参看文章PyTorch最新工具torchserve用于0.部署模型：

 from MyHandler import MyHandler
_service = MyHandler()
def handle(data, context):
    if not _service.initialized:
        _service.initialize(context)
    if data is None:
        return None
    data = _service.preprocess(data)
    data = _service.inference(data)
    data = _service.postprocess(data)
    return data

• –extra-files：表示额外文件，我们把index_to_name.json文件（torchserve内置）和MyHandler.py定义到里面。其中MyHandler.py是我们自定义的handler文件，如下所示：

"""
Module for image classification default handler
"""
import logging
import torch
import torch.nn.functional as F
import io
from PIL import Image
from torchvision import transforms
from ts.torch_handler.base_handler import BaseHandler

class MyHandler(BaseHandler):
    """
    ImageClassifier handler class. This handler takes an image
    and returns the name of object in that image.
    """
    def __init__(self, *args, **kwargs):
        super().__init__()   
        self.transform = transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225])
        ])
    
    def preprocess_one_image(self, req):
        #对输入图片进行预处理操作（这表示单个图像）
        image = req.get("data")
        if image is None:
            image = req.get("body")
        image = Image.open(io.BytesIO(image))
        image = self.transform(image) #这个就是预处理操作
        image =image.unsqueeze(0)
        return image

    def preprocess(self, requests):
        #为应对同时处理多个图片，我们设计成多图片的预处理
        images = [self.preprocess_one_image(req) for req in requests]
        images = torch.cat(images)
        return images
    
    def inference(self, images):
        #推理，调用模型对输入进行预测，预测之后得到判断结果
        outs = self.model(images)
        probs = F.softmax(outs, dim=1)
        preds = torch.argmax(probs, dim=1)
        return preds
    
    def postprocess(self, preds):
       #输出，定义输出展示方式
        preds = preds.cpu().tolist()
        return preds

最后我们的文件夹中的文件结构是这样子的：

*3.执行打包命令
通过电脑CMD进入，CD到存放文件的文件夹中，进入conda的torchserve环境，可以参看我的文章Ubuntu配置Torchserve环境，并在线发布你的深度学习模型获取详细步骤；执行以下命令：

torch-model-archiver --model-name resnet101  --version 1.0  --serialized-file resnet101.pt  --extra-files index_to_name.json,MyHandler.py  --handler my_handler.py

这时候，我们的文件夹中会生成一个resnet101.mar文件。

三、部署resnet101.mar文件到Ubuntu云服务器

恭喜你已经完成了模型的打包，这个时候只要按照文章Ubuntu配置Torchserve环境，并在线发布你的深度学习模型中的步骤，就能成功把模型进行在线部署了！