1. Introducción a torchscript y jit
1. Acerca de torchscript
TorchScript es una representación intermedia del modelo Pytorch (heredado de nn.Module). El modelo torchscript guardado se puede ejecutar en un entorno de alto rendimiento como C++.
TorchScript es una forma de crear modelos serializables y optimizables a partir de código PyTorch. Cualquier programa TorchScript puede guardarse desde un proceso de Python y cargarse en un proceso que no tenga dependencias de Python.
En pocas palabras, TorchScript puede convertir un gráfico dinámico en un gráfico estático. Bajo la función de gráfico dinámico flexible de pytorch, torchscript proporciona una herramienta que aún puede obtener la estructura del modelo (definición del modelo).
2. Acerca de torch.jit
¿Qué es JIT?
En primer lugar, debemos saber que JIT es un concepto. El nombre completo es Compilación Just In Time, que se traduce como "Compilación Just In Time" en chino. Es un método de optimización de programas. Un escenario de uso común es "regular". expresión". Por ejemplo, usando expresiones regulares en Python:
prog = re.compile(pattern)
result = prog.match(string)
#或
result = re.match(pattern, string)
Los dos ejemplos anteriores se extraen directamente de la documentación oficial de Python. De la documentación se puede ver que los dos métodos de escritura son "equivalentes" en términos de resultados. Pero preste atención a la primera forma de escribir: la expresión regular se compilará primero y luego se usará. Si continúa leyendo la documentación de Python, puede encontrar el siguiente pasaje:
usar re.compile() y guardar el objeto de expresión regular resultante para su reutilización es más eficiente cuando la expresión se usará varias veces en un solo programa. Y este proceso de compilación puede entenderse como JIT (compilación justo a tiempo).
PyTorch ha sido conocido por su "facilidad de uso" desde su lanzamiento y es más adecuado para el desarrollo de Python nativo, gracias a la estructura de "gráfico dinámico" de PyTorch. Podemos agregar cualquier declaración de control de proceso de Python antes del modelo PyTorch, y no habrá ningún problema incluso para atravesar el punto de interrupción, pero si es TensorFlow, debe usar el proceso desarrollado por TensorFlow, como el control tf.cond. Los modelos de gráficos dinámicos compensan algunas funciones avanzadas para facilitar su uso.
Ventajas de los JIT:
1. Implementación del modelo
Las dos nuevas características principales lanzadas por PyTorch versión 1.0 son JIT y C++ API. No es descabellado lanzar estas dos características juntas. JIT es un puente entre Python y C++. Podemos usar Python para entrenar el modelo y Luego, el modelo se convierte en un módulo independiente del lenguaje a través de JIT, por lo que es muy conveniente llamar a C ++. A partir de entonces, "usar Python para entrenar el modelo y usar C ++ para implementar el modelo en el entorno de producción" se ha convertido en una tarea fácil. tarea para PyTorch. Y gracias al uso de C++, ahora podemos implementar modelos de PyTorch en casi cualquier plataforma y dispositivo: Raspberry Pi, iOS, Android, etc...
- mejora del rendimiento
Dado que es una característica proporcionada para implementación y producción, es inevitable que se haya realizado una gran optimización en el rendimiento. Si la escena inferida tiene requisitos de alto rendimiento, puede considerar convertir el modelo (torch.nn.Module) al módulo TorchScript y luego proceda a inferir.
- visualización del modelo
TensorFlow o Keras son muy amigables con las herramientas de visualización de modelos (TensorBoard, etc.), porque es un modelo de programación de gráficos estáticos. Una vez definido el modelo, la estructura y la lógica directa de todo el modelo ya son claras, pero PyTorch en sí no. lo admite, por lo que los modelos PyTorch siempre han sido malos para visualizar, pero JIT mejora la situación. Ahora puede usar la función de seguimiento de JIT para obtener la lógica directa del modelo PyTorch para una determinada entrada, y puede obtener la estructura aproximada del modelo a través de la lógica directa. (Pero si forwardhay muchas declaraciones de control condicional en el método, este todavía no es un buen método)
3. Dos formas de generar el módulo TorchScript
1. Guiones
Puede usar directamente el lenguaje TorchScript para definir un módulo PyTorch JIT y luego usar torch.jit.script para convertirlo en un módulo TorchScript y guardarlo como un archivo. El lenguaje TorchScript en sí también es código Python, por lo que se puede escribir directamente en un archivo Python.
Usar el lenguaje TorchScript es como usar TensorFlow, es necesario definir un gráfico completo de antemano. Para TensorFlow, sabemos que no podemos usar directamente if y otras declaraciones en Python para el control condicional, pero necesitamos usar tf.cond, pero para TorchScript aún podemos usar directamente declaraciones de control condicional como if y for, por lo que incluso en gráficos estáticos. , PyTorch todavía se adhiere a la función "fácil de usar". El lenguaje TorchScript es un subconjunto de Python escrito estáticamente, que también se implementa usando el módulo de escritura de Python 3, por lo que la experiencia de escribir el lenguaje TorchScript es exactamente la misma que la de Python, excepto que algunas características de Python no se pueden usar (porque es un subconjunto), que se puede pasar Referencia del lenguaje TorchScript para ver las similitudes y diferencias con Python nativo.
En teoría, el módulo TorchScript definido por Scripting es muy amigable para las herramientas de visualización de modelos, porque toda la estructura del gráfico se ha definido de antemano.
- Rastreo
Una forma más sencilla de utilizar el módulo TorchScript es utilizar Tracing, que puede convertir directamente el modelo de PyTorch (torch.nn.Module) en un módulo TorchScript. "Seguimiento", como sugiere el nombre, consiste en proporcionar una "entrada" para permitir que el modelo avance nuevamente, de modo que se obtenga la estructura del gráfico a través de la ruta de flujo de la entrada. Este método es muy práctico para un modelo con una lógica directa simple, pero si hay muchas declaraciones de control de flujo en el proceso directo, puede haber problemas, porque la misma entrada no puede atravesar todas las ramas lógicas.
2. Generar un modelo de antorcha para razonar.
1. Cargue el modelo de puntero de control de antorcha exportado
Cargue archivos de configuración de modelos previamente entrenados y estructura de modelo reescrita
# 【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_cls
self.num_labels_ner = num_labels_ner
self.inner_dim = inner_dim
self.hidden_size = config.hidden_size
self.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 = RoPE
self.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 * indices
embeddings = 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 embeddings
def 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:1
outputs = 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:3
if 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_pos
kw2 = 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 mask
pad_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_h
logits_ner = logits_ner*pad_mask - (1-pad_mask)*1e12 # TODO:6
# 排除下三角
mask = torch.tril(torch.ones_like(logits_ner), -1) # TODO:7
logits_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. Cargar los parámetros del modelo.
Cargar parámetros del modelo ya entrenado
#【训练完成的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. Cree los parámetros de seguimiento de torch.jit
En la lógica de propagación hacia adelante de nuestro modelo multitarea de limpieza de direcciones, no hay múltiples estructuras de condiciones de juicio, por lo que elegimos la forma de seguimiento para registrar el proceso de propagación hacia adelante y la estructura del modelo.
#【定义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_length
segment_ids = [0] * input_length
while 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. Utilice torch.jit para exportar el modelo del módulo TorchScript
jit utiliza la forma de seguimiento (Tracing) para registrar el proceso de propagación hacia adelante y la estructura del modelo.
#【根据输出的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. Verifique que el módulo TorchScript sea correcto
#【查看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】
3. Utilice el servidor Triton para iniciar el modelo de torchscript.
1. Modifique el archivo de configuración config.ptxtx
name: "multitask_model"
platform: "pytorch_libtorch"
max_batch_size: 8
input [
{
name: "input_ids"
data_type: TYPE_INT64
dims: 64
},
{
name: "segment_ids"
data_type: TYPE_INT64
dims: 64
},
{
name: "input_mask"
data_type: TYPE_INT64
dims: 64
}
]
output [
{
name: "cls_logits"
data_type: TYPE_FP32
dims: [1, 1524]
},
{
name: "ner_logits"
data_type: TYPE_FP32
dims: [ -1, 13, 64, 64 ]
}
]
dynamic_batching {
preferred_batch_size: [ 1, 2, 4, 8 ]
max_queue_delay_microseconds: 50
}
instance_group [
{
count: 1
kind: KIND_GPU
gpus: [0]
}
]
2. Estructura del directorio modelo
multitask_test
└── multitask_model
├── 1
│ └── model.pt
├── config.pbtxt
└── label_cls.csv
3. Iniciar tritón
tritonserver --model-store=/root/ljh/space-based/Deep_Learning/Pytorch/multitask_classify_ner/multitask_test --strict-model-config=false --exit-on-error=false
4. La implementación del modelo Torchscript está por resolver
En el proceso de intentar utilizar la implementación de múltiples tarjetas de torchscript, el modelo está vinculado a cuda. Después de iniciar el modelo con triton, solo el dispositivo cuda vinculado al modelo puede ejecutarse normalmente.
Referencia de problemas similares: https://github.com/triton-inference-server/server/issues/2626