一.关于torchscript和jit介绍

1.关于torchscript

TorchScript是Pytorch模型（继承自nn.Module）的中间表示，保存后的torchscript模型可以在像C++这种高性能的环境中运行

TorchScript是一种从PyTorch代码创建可序列化和可优化模型的方法。任何TorchScript程序都可以从Python进程中保存，并加载到没有Python依赖的进程中。

简单来说，TorchScript能将动态图转为静态图，在pytorch的灵活的动态图特性下，torchscript提供了依然能够获取模型结构(模型定义)的工具。

2.关于torch.jit

什么是 JIT？
首先要知道 JIT 是一种概念，全称是 Just In Time Compilation，中文译为「即时编译」，是一种程序优化的方法，一种常见的使用场景是「正则表达式」。例如，在 Python 中使用正则表达式：

prog = re.compile(pattern)
result = prog.match(string)
#或
result = re.match(pattern, string)

上面两个例子是直接从 Python 官方文档中摘出来的 ，并且从文档中可知，两种写法从结果上来说是「等价」的。但注意第一种写法种，会先对正则表达式进行 compile，然后再进行使用。如果继续阅读 Python 的文档，可以找到下面这段话：

using re.compile() and saving the resulting regular expression object for reuse is more efficient when the expression will be used several times in a single program.
也就是说，如果多次使用到某一个正则表达式，则建议先对其进行 compile，然后再通过 compile 之后得到的对象来做正则匹配。而这个 compile 的过程，就可以理解为 JIT（即时编译）。

PyTorch 从面世以来一直以「易用性」著称，最贴合原生 Python 的开发方式，这得益于 PyTorch 的「动态图」结构。我们可以在 PyTorch 的模型前向中加任何 Python 的流程控制语句，甚至是下断点单步跟进都不会有任何问题，但是如果是 TensorFlow，则需要使用 tf.cond 等 TensorFlow 自己开发的流程控制。动态图模型通过牺牲一些高级特性来换取易用性。

JIT的优势：

1.模型部署
PyTorch 的 1.0 版本发布的最核心的两个新特性就是 JIT 和 C++ API，这两个特性一起发布不是没有道理的，JIT 是 Python 和 C++ 的桥梁，我们可以使用 Python 训练模型，然后通过 JIT 将模型转为语言无关的模块，从而让 C++ 可以非常方便得调用，从此「使用 Python 训练模型，使用 C++ 将模型部署到生产环境」对 PyTorch 来说成为了一件很容易的事。而因为使用了 C++，我们现在几乎可以把 PyTorch 模型部署到任意平台和设备上：树莓派、iOS、Android 等等…

2. 性能提升

既然是为部署生产所提供的特性，那免不了在性能上面做了极大的优化，如果推断的场景对性能要求高，则可以考虑将模型（torch.nn.Module）转换为 TorchScript Module，再进行推断。

3. 模型可视化

TensorFlow 或 Keras 对模型可视化工具（TensorBoard等）非常友好，因为本身就是静态图的编程模型，在模型定义好后整个模型的结构和正向逻辑就已经清楚了；但 PyTorch 本身是不支持的，所以 PyTorch 模型在可视化上一直表现得不好，但 JIT 改善了这一情况。现在可以使用 JIT 的 trace 功能来得到 PyTorch 模型针对某一输入的正向逻辑，通过正向逻辑可以得到模型大致的结构。（但如果在 forward 方法中有很多条件控制语句，这依然不是一个好的方法）

3.TorchScript Module 的两种生成方式

1.编码（Scripting）

可以直接使用 TorchScript Language 来定义一个 PyTorch JIT Module，然后用 torch.jit.script 来将他转换成 TorchScript Module 并保存成文件。而 TorchScript Language 本身也是 Python 代码，所以可以直接写在 Python 文件中。

使用 TorchScript Language 就如同使用 TensorFlow 一样，需要前定义好完整的图。对于 TensorFlow 我们知道不能直接使用 Python 中的 if 等语句来做条件控制，而是需要用 tf.cond，但对于 TorchScript 我们依然能够直接使用 if 和 for 等条件控制语句，所以即使是在静态图上，PyTorch 依然秉承了「易用」的特性。TorchScript Language 是静态类型的 Python 子集，静态类型也是用了 Python 3 的 typing 模块来实现，所以写 TorchScript Language 的体验也跟 Python 一模一样，只是某些 Python 特性无法使用（因为是子集），可以通过 TorchScript Language Reference 来查看和原生 Python 的异同。

理论上，使用 Scripting 的方式定义的 TorchScript Module 对模型可视化工具非常友好，因为已经提前定义了整个图结构。

2. 追踪（Tracing）

使用 TorchScript Module 的更简单的办法是使用 Tracing，Tracing 可以直接将 PyTorch 模型（torch.nn.Module）转换成 TorchScript Module。「追踪」顾名思义，就是需要提供一个「输入」来让模型 forward 一遍，以通过该输入的流转路径，获得图的结构。这种方式对于 forward 逻辑简单的模型来说非常实用，但如果 forward 里面本身夹杂了很多流程控制语句，则可能会有问题，因为同一个输入不可能遍历到所有的逻辑分枝。

二.生成一个用于推理的torch模型

1.加载已导出的torch checkpointer模型

加载预训练模型配置文件和改写的模型结构

# 【multitask_classify_ner 多任务分类模型代码（包括classify任务和ner任务）】
class BertFourLevelArea(BertPreTrainedModel):"""BERT model for four level area."""def __init__(self, config, num_labels_cls, num_labels_ner, inner_dim, RoPE):super(BertFourLevelArea, self).__init__(config, num_labels_cls, num_labels_ner, inner_dim, RoPE)self.bert = BertModel(config)self.num_labels_cls = num_labels_clsself.num_labels_ner = num_labels_nerself.inner_dim = inner_dimself.hidden_size = config.hidden_sizeself.dense_ner = nn.Linear(self.hidden_size, self.num_labels_ner * self.inner_dim * 2)self.dense_cls = nn.Linear(self.hidden_size, num_labels_cls)self.RoPE = RoPEself.dropout = nn.Dropout(config.hidden_dropout_prob)self.apply(self.init_bert_weights)def sinusoidal_position_embedding(self, batch_size, seq_len, output_dim):position_ids = torch.arange(0, seq_len, dtype=torch.float).unsqueeze(-1)indices = torch.arange(0, output_dim // 2, dtype=torch.float)indices = torch.pow(10000, -2 * indices / output_dim)embeddings = position_ids * indicesembeddings = torch.stack([torch.sin(embeddings), torch.cos(embeddings)], dim=-1)embeddings = embeddings.repeat((batch_size, *([1]*len(embeddings.shape))))embeddings = torch.reshape(embeddings, (batch_size, seq_len, output_dim))embeddings = embeddings.to(self.device)return embeddingsdef forward(self, input_ids, token_type_ids=None, attention_mask=None):# sequence_output: Last Encoder Layer.shape: (batch_size, seq_len, hidden_size)encoded_layers, pooled_output = self.bert(input_ids, token_type_ids, attention_mask)sequence_output = encoded_layers[-1]batch_size = sequence_output.size()[0]seq_len = sequence_output.size()[1]# 【Bert Ner GlobalPointer】:# outputs: (batch_size, seq_len, num_labels_ner*inner_dim*2)outputs = self.dense_ner(sequence_output)# outputs: (batch_size, seq_len, num_labels_ner, inner_dim*2)outputs = torch.split(outputs, self.inner_dim * 2, dim=-1)      # TODO:1outputs = torch.stack(outputs, dim=-2)              # TODO:2# qw,kw: (batch_size, seq_len, num_labels_ner, inner_dim)qw, kw = outputs[...,:self.inner_dim], outputs[...,self.inner_dim:] # TODO:3if self.RoPE:# pos_emb:(batch_size, seq_len, inner_dim)pos_emb = self.sinusoidal_position_embedding(batch_size, seq_len, self.inner_dim)# cos_pos,sin_pos: (batch_size, seq_len, 1, inner_dim)cos_pos = pos_emb[..., None, 1::2].repeat_interleave(2, dim=-1)sin_pos = pos_emb[..., None,::2].repeat_interleave(2, dim=-1)qw2 = torch.stack([-qw[..., 1::2], qw[...,::2]], -1)qw2 = qw2.reshape(qw.shape)qw = qw * cos_pos + qw2 * sin_poskw2 = torch.stack([-kw[..., 1::2], kw[...,::2]], -1)kw2 = kw2.reshape(kw.shape)kw = kw * cos_pos + kw2 * sin_pos# logits_ner:(batch_size, num_labels_ner, seq_len, seq_len)logits_ner = torch.einsum('bmhd,bnhd->bhmn', qw, kw)    # TODO:4# padding maskpad_mask = attention_mask.unsqueeze(1).unsqueeze(1).expand(batch_size, self.num_labels_ner, seq_len, seq_len)   # TODO:5# pad_mask_h = attention_mask.unsqueeze(1).unsqueeze(-1).expand(batch_size, self.num_labels_ner, seq_len, seq_len)# pad_mask = pad_mask_v&pad_mask_hlogits_ner = logits_ner*pad_mask - (1-pad_mask)*1e12    # TODO:6# 排除下三角mask = torch.tril(torch.ones_like(logits_ner), -1)  # TODO:7logits_ner = logits_ner - mask * 1e12   # TODO:8# 【Bert Classify】:pooled_output = self.dropout(pooled_output)logits_cls = self.dense_cls(pooled_output)return logits_cls, logits_ner#【加载预训练模型参数】
config = modeling.BertConfig.from_json_file('/root/ljh/space-based/Deep_Learning/Pytorch/multitask_classify_ner/pretrain_model/bert-base-chinese/config.json')#【加载我们训练的模型  】
#【num_labels_cls 和 num_labels_ner为我们训练的label_counts 这次训练的分类任务标签数为1524 ，NER任务的分类数为13】
num_labels_cls = 1524
num_labels_ner = 13
model = modeling.BertFourLevelArea(config,num_labels_cls=num_labels_cls,num_labels_ner=num_labels_ner,inner_dim=64,RoPE=False
)

2.载入模型模型参数

加载已经训练好的模型参数

#【训练完成的tourch 模型地址】
init_checkpoint='/root/ljh/space-based/Deep_Learning/Pytorch/multitask_classify_ner/outputs.bak/multitask_classify_ner/pytorch_model.bin'
#【载入模型】
checkpoint = torch.load(init_checkpoint, map_location=torch.device("cuda"))
checkpoint = checkpoint["model"] if "model" in checkpoint.keys() else checkpoint
model.load_state_dict(checkpoint)
device = torch.device("cuda")#【将模型导入GPU】
model = model.to(device)
#【模型初始化】
model.eval()

3.构建torch.jit追踪参数

在我们的地址清洗多任务模型的forwad前向传播逻辑中并没有多种判断条件的结构，因此我们选择 追踪（Tracing）的形式来记录模型的前向传播过程和结构。

#【定义tokenizer 】
from transformers import BertTokenizerFast
tokenizer = BertTokenizerFast.from_pretrained('/root/ljh/space-based/Deep_Learning/Pytorch/multitask_classify_ner/pretrain_model/bert-base-chinese', add_special_tokens=True, do_lower_case=False)input_str='上海上海市青浦区华隆路E通世界华新园'
max_seq_length=64#【生成bert模型输出】
def input2feature(input_str, max_seq_length=48):# 预处理字符tokens_a = tokenizer.tokenize(input_str)# 如果超过长度限制，则进行截断if len(input_str) > max_seq_length - 2:tokens_a = tokens_a[0:(max_seq_length - 2)]tokens = ["[CLS]"] + tokens_a + ["[SEP]"]input_ids = tokenizer.convert_tokens_to_ids(tokens)input_length = len(input_ids)input_mask = [1] * input_lengthsegment_ids = [0] * input_lengthwhile len(input_ids) < max_seq_length:input_ids.append(0)input_mask.append(0)segment_ids.append(0)return input_ids, input_mask, segment_ids#【输入地址token化     input_ids --> list()】
input_ids, input_mask, segment_ids = input2feature(input_str, max_seq_length)# 【list -> tensor】
input_ids = torch.tensor(input_ids, dtype=torch.long)
input_mask = torch.tensor(input_mask, dtype=torch.long)
segment_ids = torch.tensor(segment_ids, dtype=torch.long)#【这里stack 是因为模型内部定义的输出参数需要stack 】
input_ids = torch.stack([input_ids], dim=0)
input_mask = torch.stack([input_mask], dim=0)
segment_ids = torch.stack([segment_ids], dim=0)#【将参数推送至cuda设备中】
device = torch.device("cuda")
input_ids = input_ids.to(device)
input_mask = input_mask.to(device)
segment_ids = segment_ids.to(device)#【input_ids.shape --> torch.Size([1, 64])】

4.使用torch.jit 导出TorchScript Module模型

jit 使用追踪（Tracing）的形式来记录模型的前向传播过程和结构

#【根据输出的input_ids, input_mask, segment_id记录前向传播过程】
script_model = torch.jit.trace(model,[input_ids, input_mask, segment_ids],strict=True)
#【保存】
torch.jit.save(script_model, "./multitask_test/multitask_model/1/model.pt")

5.验证TorchScript Module是否正确

#【查看torch模型结果】
cls_res, ner_res = model(input_ids, input_mask, segment_ids)import numpy as np
np.argmax(cls_res.detach().cpu().numpy()) 
#【result:673】#【load torchscript model】
jit_model = torch.jit.load('./multitask_test/multitask_model/1/model.pt')
example_outputs_cls,example_outputs_ner = jit_model(input_ids, input_mask, segment_ids)
np.argmax(example_outputs_cls.detach().cpu().numpy()) 
#【result:673】

三. 使用triton server 启动 torchscript model模型

1. 修改config.ptxtx 配置文件

name: "multitask_model"
platform: "pytorch_libtorch"
max_batch_size: 8
input [{name: "input_ids"data_type: TYPE_INT64dims:  64},{name: "segment_ids"data_type: TYPE_INT64dims:  64 },{name: "input_mask"data_type: TYPE_INT64dims:  64}
]
output [{name: "cls_logits"data_type: TYPE_FP32dims: [1, 1524]},{name: "ner_logits"data_type: TYPE_FP32dims: [ -1, 13, 64, 64 ]}
]dynamic_batching {preferred_batch_size: [ 1, 2, 4, 8 ]max_queue_delay_microseconds: 50}instance_group [
{count: 1kind: KIND_GPUgpus: [0]
}
]

2.模型目录结构

multitask_test
└── multitask_model
├── 1
│ └── model.pt
├── config.pbtxt
└── label_cls.csv

3.启动triton

tritonserver --model-store=/root/ljh/space-based/Deep_Learning/Pytorch/multitask_classify_ner/multitask_test  --strict-model-config=false --exit-on-error=false

四. torchscript model模型部署待解决点

在尝试使用torchscript多卡部署的过程中会出现model模型与cuda绑定的情况，在使用triton启动模型之后只有与模型进行绑定的cuda设备才能够正常执行。

类似问题参考：https://github.com/triton-inference-server/server/issues/2626