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环境,并在线发布你的深度学习模型中的步骤,就能成功把模型进行在线部署了!

文章出处登录后可见!

已经登录?立即刷新

共计人评分,平均

到目前为止还没有投票!成为第一位评论此文章。

(0)
xiaoxingxing的头像xiaoxingxing管理团队
上一篇 2022年5月11日
下一篇 2022年5月11日

相关推荐