Hace aproximadamente un año, los desarrolladores de PyTorch presentaron la comunidad TorchScript , una herramienta que le permite crear una solución alienable a partir de una tubería en Python con un par de clics del mouse que pueden integrarse en un sistema C ++. A continuación comparto la experiencia de su uso y trato de describir las trampas encontradas a lo largo de este camino. Prestaré especial atención a la implementación del proyecto en Windows, porque aunque la investigación en ML generalmente se realiza en Ubuntu, la solución final a menudo (¡de repente!) Se requiere bajo las "ventanas".
El código de muestra para exportar un modelo y un proyecto C ++ usando el modelo se puede encontrar en el repositorio en GitHub .

Los desarrolladores de PyTorch no se dejan engañar. La nueva herramienta realmente le permite convertir un proyecto de investigación en PyTorch en código incrustado en un sistema C ++ en un par de días hábiles, y con cierta habilidad más rápido.
TorchScript apareció en PyTorch versión 1.0 y continúa evolucionando y cambiando. Si la primera versión hace un año estaba llena de errores y era más experimental, entonces la versión actual 1.3 al menos en el segundo punto es notablemente diferente: ya no puede llamarse experimental, es bastante adecuada para uso práctico. Me enfocaré en ella.
En el corazón de TorchScript se encuentra su propio compilador independiente (sin Python) de un lenguaje similar a Python, así como herramientas para convertir un programa escrito en Python y PyTorch, métodos para guardar y cargar los módulos resultantes, y una biblioteca para usarlos en C ++. Para trabajar, deberá agregar varias DLL al proyecto con un peso total de aproximadamente 70 MB (para Windows) para trabajar en la CPU y 300 MB para la versión de GPU. TorchScript admite la mayoría de las características de PyTorch y las características principales del lenguaje python. Pero las bibliotecas de terceros, como OpenCV o NumPy, deberán olvidarse. Afortunadamente, muchas funciones de NumPy tienen un análogo en PyTorch.
Convertir tubería a modelo PyTorch en TorchScript
TorchScript ofrece dos formas de convertir el código Python a su formato interno: rastreo y scripting (rastreo y scripting). ¿Por qué dos? No, está claro, por supuesto, que dos son mejores que uno ...

Pero en el caso de estos métodos, resulta, como en el conocido aforismo, sobre las desviaciones izquierda y derecha: ambas son peores. Bueno, el mundo no es perfecto. Solo en una situación específica, debe elegir la que sea más adecuada.
El método de rastreo es muy simple. Se toma una muestra de datos (generalmente inicializados por números aleatorios), se envía a la función o método de la clase que nos interesa, y PyTorch construye y almacena el gráfico de cálculo de la misma manera que lo hace cuando se entrena una red neuronal. Voila: el guión está listo:
import torch import torchvision model = torchvision.models.resnet34(pretrained = True) model.eval() sample = torch.rand(1, 3, 224, 224) scripted_model = torch.jit.trace(model, sample)
El ejemplo anterior produce un objeto de la clase ScriptModule. Se puede guardar
scripted_model.save('my_script.pth')
y luego cárguelo en un programa C ++ (más sobre eso a continuación ) o en el código Python en lugar del objeto original:
Ejemplo de código Python usando un modelo guardado import cv2 from torchvision.transforms import Compose, ToTensor, Normalize transforms = Compose([ToTensor(), Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) img = cv2.resize(cv2.imread('pics/cat.jpg'), (224,224)) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) x = transforms(img).unsqueeze(0)
tensor(282) tensor(12.8130, grad_fn=<SelectBackward>)
El ScriptModule
resultante puede aparecer en cualquier lugar nn.Module
se usa comúnmente.
De la manera descrita, puede rastrear instancias de la clase y funciones nn.Module
(en este último caso, se torch._C.Function
una instancia de la clase torch._C.Function
).
Este método (rastreo) tiene una ventaja importante: de esta manera puede convertir casi cualquier código de Python que no use bibliotecas externas. Pero hay un inconveniente igualmente importante: para cualquier rama, solo se recordará esa rama que se ejecutó en los datos de prueba:
def my_abs(x): if x.max() >= 0: return x else: return -x my_abs_traced = torch.jit.trace(my_abs, torch.tensor(0)) print(my_abs_traced(torch.tensor(1)), my_abs_traced(torch.tensor(-1)))
c:\miniconda3\lib\site-packages\ipykernel_launcher.py:2: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs! tensor(1) tensor(-1)
¡Uy! Esto parece no ser lo que nos gustaría, ¿verdad? Es bueno que al menos se emita un mensaje de advertencia (TracerWarning). Vale la pena prestar atención a tales mensajes.
Aquí el segundo método viene en nuestra ayuda - scripting:
my_abs_script = torch.jit.script(my_abs) print(my_abs_script(torch.tensor(1)), my_abs_script(torch.tensor(-1)))
tensor(1) tensor(1)
¡Hurra, se recibe el resultado esperado! Las secuencias de comandos analizan recursivamente el código de Python y lo convierten en código en su propio lenguaje. En la salida, también obtenemos la clase ScriptModule
(para módulos) o torch._C.Function
(para funciones). Parecería, aquí está, ¡felicidad! Pero surge otro problema: el lenguaje interno de TorchScript está fuertemente tipado, a diferencia de Python. El tipo de cada variable está determinado por la primera asignación, el tipo de los argumentos de la función por defecto es Tensor
. Por lo tanto, por ejemplo, un patrón familiar
def my_func(x): y = None if x.max() > 0: y = x return y my_func = torch.jit.script(my_func)
El seguimiento fallará.
Un error de rastreo se ve así RuntimeError Traceback (most recent call last) <ipython-input-9-25414183a687> in <module>() ----> 1 my_func = torch.jit.script(my_func) d:\programming\3rd_party\pytorch\pytorch_ovod_1.3.0a0_de394b6\torch\jit\__init__.py in script(obj, optimize, _frames_up, _rcb) 1224 if _rcb is None: 1225 _rcb = _gen_rcb(obj, _frames_up) -> 1226 fn = torch._C._jit_script_compile(qualified_name, ast, _rcb, get_default_args(obj)) 1227 # Forward docstrings 1228 fn.__doc__ = obj.__doc__ RuntimeError: Variable 'y' previously has type None but is now being assigned to a value of type Tensor : at <ipython-input-8-75677614fca6>:4:8 def my_func(x): y = None if x.max() > 0: y = x ~ <--- HERE return y
Es de destacar que, aunque se produce un error cuando torch.jit.script
llama a torch.jit.script
, también se indica el lugar que lo causó en el código de torch.jit.script
.
Incluso los puntos después de que las constantes comienzan a jugar un papel:
def my_func(x): if x.max() > 0: y = 1.25 else: y = 0 return y my_func = torch.jit.script(my_func)
dará un error RuntimeError Traceback (most recent call last) <ipython-input-10-0a5f18586763> in <module>() 5 y = 0 6 return y ----> 7 my_func = torch.jit.script(my_func) d:\programming\3rd_party\pytorch\pytorch_ovod_1.3.0a0_de394b6\torch\jit\__init__.py in script(obj, optimize, _frames_up, _rcb) 1224 if _rcb is None: 1225 _rcb = _gen_rcb(obj, _frames_up) -> 1226 fn = torch._C._jit_script_compile(qualified_name, ast, _rcb, get_default_args(obj)) 1227 # Forward docstrings 1228 fn.__doc__ = obj.__doc__ d:\programming\3rd_party\pytorch\pytorch_ovod_1.3.0a0_de394b6\torch\jit\__init__.py in _rcb(name) 1240 # closure rcb fails 1241 result = closure_rcb(name) -> 1242 if result: 1243 return result 1244 return stack_rcb(name) RuntimeError: bool value of Tensor with more than one value is ambiguous
¡Porque es necesario escribir no 0
, sino 0.
para que el tipo en ambas ramas sea el mismo! Echado a perder, ya sabes, con tu pitón!
Este es solo el comienzo de la lista de cambios que debe realizar en el código de Python para que pueda convertirse con éxito en un módulo TorchScript. Enumeraré los casos más típicos con más detalle más adelante . En principio, no hay ciencia de cohetes aquí y su código puede corregirse en consecuencia. Pero la mayoría de las veces no quiero arreglar módulos de terceros, incluidos los estándar de torchvision
, y como de costumbre, generalmente no son adecuados para la creación de scripts.
Afortunadamente, ambas tecnologías se pueden combinar: lo que se está escribiendo se está escribiendo y lo que no se está escribiendo es rastrear:
class MyModule(torch.nn.Module): def __init__(self): super(MyModule, self).__init__() self.resnet = torchvision.models.resnet34(pretrained = True)
En el ejemplo anterior, el rastreo se utiliza para incluir un módulo que no es programable en un módulo donde no hay suficiente rastreo y es necesario realizar secuencias de comandos. Hay una situación inversa. Por ejemplo, si necesitamos cargar un modelo en ONNX, se utiliza el rastreo. Pero el modelo trazado puede incluir funciones TorchScript, por lo que la lógica que requiere ramas y bucles puede implementarse allí. Se da un ejemplo en la documentación oficial de torch.onnx .
Las características proporcionadas por PyTorch para crear módulos TorchScript se describen con más detalle en la documentación oficial y el torch.jit
. En particular, no mencioné una forma conveniente de usar torch.jit.trace
y torch.jit.script
en forma de decoradores, sobre las peculiaridades de la depuración de código con torch.jit.script
. Esto y mucho más está en la documentación.
Incluimos el modelo en un proyecto C ++
Desafortunadamente, la documentación oficial se limita a ejemplos de la forma "agregar 2 tensores generados usando torch.ones
". Preparé un ejemplo de un proyecto más cercano a la realidad que envía una imagen de OpenCV a la red neuronal y recibe los resultados en forma de un tensor de respuesta, una tupla de variables, una imagen con resultados de segmentación.
Para que el ejemplo funcione, necesita guiones de clasificación guardados con ResNet34 y segmentación con DeepLabV3. Para preparar estos scripts, debe ejecutar este bloc de notas jupyter .
Necesitamos la biblioteca torchlib
. Puede obtenerlo de varias maneras:
- Si ya tiene PyTorch instalado utilizando
pip install
, puede encontrarlo en el directorio de Python: <Miniconda3>\Lib\site-packages\torch
; - Si tiene PyTorch compilado desde la fuente, entonces está allí:
<My Pytorch repo>\build\lib.win-amd64-3.6\torch
; - Finalmente, puede descargar la biblioteca por separado desde pytorch.org seleccionando Language = C ++ y descomprimir el archivo.
El código C ++ es bastante simple. Es necesario:
- Incluir archivo de encabezado
#include <torch/script.h>
- Descargar modelo
torch::jit::script::Module module = torch::jit::load("../resnet34_infer.pth");
- Preparar datos
torch::Tensor tensor = torch::from_blob(img.data, { img.rows, img.cols, 3 }, torch::kByte);
- Llama a la función de
forward
y obtén resultados
auto output = module.forward( { tensor } )
- Obtenga datos del resultado. Cómo hacerlo depende de lo que devuelva la red neuronal. Por cierto, en el caso general, también puede aceptar no solo una imagen, por lo tanto, es mejor mirar el código fuente de todo el ejemplo , hay diferentes opciones. Por ejemplo, para obtener datos de un tensor unidimensional de tipo flotante:
float* data = static_cast<float*>(output.toTensor().data_ptr());
- Hay una sutileza más. No olvide insertar el análogo
with torch.no_grad()
en el código para no desperdiciar recursos en el cálculo y almacenamiento de los gradientes que no necesitamos. Desafortunadamente, este comando no se puede incluir en el script, por lo que debe agregarlo al código C ++:
torch::NoGradGuard no_grad;
La forma de construir un proyecto usando CMake se describe en la guía oficial . Pero el tema del proyecto en Visual Studio no se revela allí, por lo que lo describiré con más detalle. Tendrá que ajustar manualmente la configuración del proyecto:
- Probé en Visual Studio 2017. No puedo decir sobre otras versiones.
- El conjunto de herramientas v14.11 v141 debe estar instalado (marque
"VC++ 2017 version 15.4 v14.11 toolset"
en el instalador VS). - La plataforma debe ser
x64
. - En
General → Platform Toolset
v141(Visual Studio 2017)
General → Platform Toolset
seleccione v141(Visual Studio 2017)
- En
C/C++ → General → Additional Include Directories
agregue <libtorch dir>\include
- En
Linker → General → Additional Library Directories
agregue <libtorch dir>\lib
- En
Linker → Input → Additional Dependencies
agregue torch.lib; c10.lib
torch.lib; c10.lib
. En Internet, escriben que caffe2.lib
todavía puede ser necesario, y para la GPU y algo más de <libtorch dir>\lib
, pero en la versión actual, agregar estas dos bibliotecas fue suficiente para mí. Quizás esta es información desactualizada. - También escriben que necesita configurar
C/C++ → Language → Conformance Mode
= No
, pero no vi la diferencia.
Además, la variable __cplusplus
NO debe declararse en el proyecto. Intentando agregar la /Zc:__cplusplus
resultará en errores de compilación en el archivo ivalue.h
.
En el proyecto adjunto, la configuración de ruta (no solo a TorchLib, sino también a OpenCV y CUDA) se transfiere al archivo de accesorios , antes del ensamblaje, debe registrarlos allí de acuerdo con su configuración local. Eso, de hecho, es todo.
¿Qué más tener en cuenta?
Si el proceso descrito le pareció demasiado simple, su intuición no lo engañó. Hay una serie de matices que deben considerarse para convertir un modelo PyTorch escrito en Python a TorchScript. Voy a enumerar a continuación los que tuve que enfrentar. Ya he mencionado algunos, pero repito para recoger todo en un solo lugar.

- El tipo de variables que se pasan a la función es Tensor por defecto. Si en algunos casos (muy frecuentes) esto es inaceptable, tendrá que declarar los tipos manualmente usando anotaciones de tipo estilo MyPy, algo como esto:
def calc_letter_statistics(self, cls_preds: List[Tensor], cls_thresh: float)->Tuple[int, Tuple[Tensor, Tensor, Tensor]]
más o menos:
def calc_letter_statistics(self, cls_preds, cls_thresh):
- Las variables están fuertemente tipadas y el tipo, si no se especifica explícitamente, está determinado por la primera asignación. Construcciones familiares de la forma
x=[]; for ...: x.append(y)
x=[]; for ...: x.append(y)
tendrá que ser editado, porque al momento de asignar []
compilador no puede determinar qué tipo estará en la lista. Por lo tanto, deberá especificar el tipo explícitamente, por ejemplo:
from typing import List x: List[float] = []
o (otro "por ejemplo")
from torch import Tensor from typing import Dict, Tuple, List x: Dict[int: Tuple[float, List[Tensor], List[List[int]]]] = {}
- En el ejemplo anterior, son los nombres los que deben importarse, ya que estos nombres están cosidos en el código TorchScript. Enfoque alternativo, aparentemente legal
import torch import typing x: typing.List[torch.Tensor] = []
dará como resultado un tipeo de constructor de tipo Desconocido. Error de lista al crear scripts
- Otro diseño familiar del que tienes que separarte:
x = None if smth: x = torch.tensor([1,2,3])
Hay dos opciones O asigne Tensor ambas veces (el hecho de que sea de diferentes dimensiones no da miedo):
x = torch.tensor(0) if smth: x = torch.tensor([1,2,3])
y no olvides buscar lo que se romperá después de tal reemplazo. O trata de escribir honestamente:
x: Optional[Tensor] = None if smth: x = torch.tensor([1,2,3])
pero luego con el uso adicional de x
donde se espera el tensor, lo más probable es que obtengamos un error: esperaba un valor de tipo 'Tensor' para el argumento 'x', pero en su lugar encontramos el tipo 'Opcional [Tensor]'.
No olvides escribir, por ejemplo, x=0.
durante la primera asignación x=0.
en lugar del habitual x=0
, etc., si la variable x
debe ser de tipo float
.
Si en algún lugar utilizamos la inicialización antigua del tensor a través de x = torch.Tensor(...)
, tendrá que x = torch.Tensor(...)
él y reemplazarlo con una versión más joven con una letra minúscula x = torch.tensor(...)
. De lo contrario, durante la secuencia de comandos volará: Desconocido incorporado op: aten :: Tensor. Aquí hay algunas sugerencias: aten :: tensor . Parece que incluso explican cuál es el problema, y está claro lo que hay que hacer. Sin embargo, está claro si ya conoce la respuesta correcta.
El código está escrito en el contexto del módulo donde torch.jit.script
llama torch.jit.script
. Por lo tanto, si en algún lugar, en las entrañas de la clase o función con math.pow
, por ejemplo, math.pow
, deberá agregar las import math
al módulo de compilación. Y es mejor hacer un script de la clase donde se declara: ya sea usando el decorador @torch.jit.script
o declarando una función adicional al lado que haga que ScriptModule quede fuera de él. De lo contrario, recibimos un mensaje de error matemático de valor indefinido cuando intentamos compilar una clase de un módulo en el que, aparentemente, se realizó la importación math
.
Si en algún lugar tiene una construcción de la forma my_tensor[my_tensor < 10] = 0
o similar, obtendrá un error críptico al escribir:
*aten::index_put_(Tensor(a!) self, Tensor?[] indices, Tensor values, bool accumulate=False) -> (Tensor(a!)):* *Expected a value of type 'Tensor' for argument 'values' but instead found type 'int'.* *aten::index_put_(Tensor(a!) self, Tensor[] indices, Tensor values, bool accumulate=False) -> (Tensor(a!)):* *Expected a value of type 'List[Tensor]' for argument 'indices' but instead found type 'List[Optional[Tensor]]'.*
Lo que necesita es reemplazar el número con el tensor: my_tensor[my_tensor < 10] = torch.tensor(0.).to(my_tensor.device)
. Y no olvide a) sobre la correspondencia de los tipos my_tensor
y el tensor creado (en este caso, float) yb) sobre .to(my_tensor.device)
. Si olvida el segundo, todo estará programado, pero ya en el proceso de trabajar con la GPU, se molestará, lo que se verá como las palabras crípticas Error CUDA: se encontró un acceso ilegal a la memoria , ¡sin indicar dónde ocurrió el error!
No olvide que, de forma predeterminada, nn.Module
y, en consecuencia, los modelos de torchvision se crean en "modo de tren" (no lo creerá, pero resulta que existe dicho modo ). En este caso, se usan Dropout y otros trucos del modo tren, que rompen el rastro o conducen a resultados inadecuados cuando se ejecutan. Recuerde llamar a model.eval()
antes de escribir o rastrear.
Para funciones y clases ordinarias, debe escribir el tipo de script, para nn.Module, una instancia
Intente en un método con script acceder a una variable global
cls_thresh = 0.3 class MyModule(torch.nn.Module): ... x = r < cls_thresh ...
dará como resultado un error de scripting de la forma en que el valor de python del tipo 'float' no se puede usar como valor . Es necesario hacer de la variable un atributo en el constructor:
cls_thresh = 0.3 class MyModule(torch.nn.Module): def __init__(self): ... self.cls_thresh = cls_thresh ... x = r < self.cls_thresh ...
- Otra sutileza surge si el atributo de clase se usa como un parámetro de corte:
class FPN(nn.Module): def __init__(self, block, num_blocks, num_layers =5): ... self.num_layers = num_layers def forward(self, x): ... return (p3, p4, p5, p6, p7)[:self.num_layers]
provoca errores de secuencia de comandos Los índices de corte de tupla deben ser constantes enteras . Es necesario indicar que el atributo num_layers es constante y no cambiará:
class FPN(nn.Module): num_layers: torch.jit.Final[int] def __init__(self, block, num_blocks, num_layers =5): ...
- En algunos casos, donde el tensor solía encajar normalmente, debe pasar explícitamente el número:
xx1 = x1.clamp(min=x1[i])
arroja un error al Expected a value of type 'Optional[number]' for argument 'min' but instead found type 'Tensor'.
scripts Expected a value of type 'Optional[number]' for argument 'min' but instead found type 'Tensor'.
. Bueno, aquí desde el mensaje de error está claro qué hacer:
xx1 = x1.clamp(min=x1[i].item())
Los problemas anteriores se producen al rastrear. Debido a ellos, generalmente no es posible compilar simplemente soluciones preparadas en TorchScript, y debe masajear el código fuente durante mucho tiempo (si el código fuente es apropiado para editar) o usar el rastreo. Pero el rastro tiene sus propios matices:
- Las construcciones del formulario no funcionan en la traza
tensor_a.to(tensor_b.device)
El dispositivo en el que se carga el tensor se fija al momento del rastreo y no cambia durante la ejecución. Este problema se puede superar parcialmente declarando al tensor miembro de nn.Module
tipo Parameter
. Luego, al cargar el modelo, se iniciará en el dispositivo que se especifica en la función torch.jit.load
.
Epílogo
Todo lo anterior, por supuesto, crea problemas. Pero TorchScript le permite combinar y enviar a la solución como un todo el modelo en sí y el código Python que proporciona el procesamiento previo y posterior. Sí, y el tiempo para preparar la solución para la compilación, incluso a pesar de las dificultades anteriores, es incomparablemente menor que el costo de crear una solución, pero aquí PyTorch ofrece grandes ventajas, por lo que el juego vale la pena.
