最新域名网站,页面设计心得体会,网站建设系统设计报告,深圳网站品牌建设一.关于torchscript和jit介绍
1.关于torchscript
TorchScript是Pytorch模型#xff08;继承自nn.Module#xff09;的中间表示#xff0c;保存后的torchscript模型可以在像C这种高性能的环境中运行
TorchScript是一种从PyTorch代码创建可序列化和可优化模型的方法。任何T…一.关于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 等等…
性能提升
既然是为部署生产所提供的特性那免不了在性能上面做了极大的优化如果推断的场景对性能要求高则可以考虑将模型torch.nn.Module转换为 TorchScript Module再进行推断。
模型可视化
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 对模型可视化工具非常友好因为已经提前定义了整个图结构。
追踪Tracing
使用 TorchScript Module 的更简单的办法是使用 TracingTracing 可以直接将 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, dtypetorch.float).unsqueeze(-1)indices torch.arange(0, output_dim // 2, dtypetorch.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_idsNone, attention_maskNone):# 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_vpad_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_clsnum_labels_cls,num_labels_nernum_labels_ner,inner_dim64,RoPEFalse
)
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_locationtorch.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_tokensTrue, do_lower_caseFalse)input_str上海上海市青浦区华隆路E通世界华新园
max_seq_length64#【生成bert模型输出】
def input2feature(input_str, max_seq_length48):# 预处理字符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, dtypetorch.long)
input_mask torch.tensor(input_mask, dtypetorch.long)
segment_ids torch.tensor(segment_ids, dtypetorch.long)#【这里stack 是因为模型内部定义的输出参数需要stack 】
input_ids torch.stack([input_ids], dim0)
input_mask torch.stack([input_mask], dim0)
segment_ids torch.stack([segment_ids], dim0)#【将参数推送至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],strictTrue)
#【保存】
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_idsdata_type: TYPE_INT64dims: 64},{name: segment_idsdata_type: TYPE_INT64dims: 64 },{name: input_maskdata_type: TYPE_INT64dims: 64}
]
output [{name: cls_logitsdata_type: TYPE_FP32dims: [1, 1524]},{name: ner_logitsdata_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-configfalse --exit-on-errorfalse四. torchscript model模型部署待解决点
在尝试使用torchscript多卡部署的过程中会出现model模型与cuda绑定的情况在使用triton启动模型之后只有与模型进行绑定的cuda设备才能够正常执行。
类似问题参考https://github.com/triton-inference-server/server/issues/2626
