DLOps - ONNX

Una vez hemos entrenado varios modelos, los hemos comparado y hemos decidido cuál usaremos en produccion, tenemos que exportarlo. Para ello tenemos existen alternativas, en función de la aplicación (desde desplegar un modelo en dispositivos móviles o IoT hasta en servidores en la nube accesibles a través de una API). En esta serie de posts asumiermos que nuestro modelo será ejecutado en un servidor en la nube, lo cual es lo más común ya que de esta manera podemos controlar los recursos computacionales disponibles para su ejecución, monitorizarlo, desplegar nuevas versiones fácilmente, etc. En nuestro caso, que hemos entrenado los modelos usando Pytorch y Pytorch Lightning, podríamos usar cualquier framework en Python que nos permita servir las predicciones a través de internet, como por ejemplo Flask o FastAPI. El principal problema de esta opción es que tendremos que cargar todas las librerías (y sus dependencias) en nuestra API, lo cual resultará en una carga muy pesada. Recientemente, Pytorch incluye una solución dedicada para este caso de uso, Torchserve que si bien nos ofrece una solución optimizada para servir modelos en producción, está limitada al uso de modelos en Pytorch.

Es en este punto en el que entra ONNX, un estándar abierto para la representación de redes neuronales que permite la interoperabilidad entre librerías y ofrece una solución optimizada para servir modelos en producción (tanto en la nube como on dispositvos móviles). De esta manera podemos desacoplar el entrenamiento de los modelos de su puesta en producción, utilizando en cada caso las herramientas preferidas para su desarrollo.

Exportar un modelo a ONNX

Vamos a ver como podemos exportar un modelo entrenado a ONNX. En primer lugar, cargaremos el checkpoint deseado (el cual generamos en el post anterior).

from src import *

module = MNISTModule.load_from_checkpoint('checkpoints/006-val_loss=0.14715-epoch=7.ckpt')
module.mlp

Sequential(
  (0): Linear(in_features=784, out_features=100, bias=True)
  (1): ReLU()
  (2): Linear(in_features=100, out_features=1, bias=True)
)

Es una buena práctica evaluar nuestro modelo antes y después de exportarlo para asegurarnos de que todo funciona correctamenete.

import torch

dm = MNISTDataModule(**module.hparams['datamodule'])
dm.setup()

def torch_eval():
    module.eval()
    with torch.no_grad():
        preds, labels = torch.tensor([]), torch.tensor([])
        for imgs, _labels in dm.val_dataloader():
            outputs = module.predict(imgs) > 0.5
            preds = torch.cat([preds, outputs.cpu().long()])
            labels = torch.cat([labels, _labels])

    acc = (preds == labels).float().mean()
    return acc.item()

torch_eval()

0.949999988079071

Pytorch Lightning nos permite exportar un modelo a ONNX de manera muy sencilla con la siguiente línea.

input_sample = torch.randint(0, 255, (1, 28, 28), dtype=torch.uint8)
module.to_onnx(
    'models/binary_classifier_3.onnx', # nombre del modelo a guardar
    input_sample, # ejemplo de entrada
    export_params=True, # exportar los parametros del modelo
    opset_version=11, # en función de las ops en el modelo, se puede cambiar el opset
    input_names = ['input'], # nombre de la entrada	para usar en producción
    output_names = ['output'],  # nombre de la salida para usar en producción
    dynamic_axes={  # para poder tener diferentes batch sizes
        'input' : {0 : 'batch_size'},
        'output' : {0 : 'batch_size'},
    },
)

Como ves le tenemos que indicar el nombre del modelo y la carpeta en la que lo queremos guardar, darle unas entradas de ejemplo (que ONNX usará para ideantificar todas las operaciones que se realizan dentro del modelo y guardarlas), si queremos exportar los parámetros del modelo, la versión del opset (este es el conjunto de operaciones soportadas, que va cambiando a medida que se añaden nuevas funcionalidades) y, opcionalmente, nombres para entradas y salidas (esto es importante si nuestro modelo tiene varias entradas y/o salidas) así como indicar qué ejes son dinámicos (útil para poder usar el modelo en producción con diferentes batch sizes).

ONNXRuntime

Una vez hemos exportado nuestro modelo podemos cargarlo y ejecutarlo usando el ONNXRuntime. Esta es una de las ventajas de ONNX, y es que existen runtimes para multiples entornos y lenguajes. Así pues, podrás entrenar el modelo en Python y luego desplegarlo tanto en Python como en la web con Javascript, en dipositivos moviles con Android o iOS, etc.

En Python, puedes instalarlo con el comando pip install onnxruntime.

Para cargar el modelo instanciamos una InferenceSession con el path al modelo exportado. Luego, definiermos las entradas al modelo usando el un dict con el nombre definido en la fase de exportación. Date cuenta que ahora el modelo usará como entradas array de NumPy, ya que en este entorno Pytorch ya no existe. Si que es importante, sin embargo, que uses el mismo tamaño y tipo de datos que usaste en el entrenamiento. Por último, podemos obtener las salidas del modelo usando el método run().

import onnxruntime as ort
import numpy as np

ort_session = ort.InferenceSession('models/binary_classifier_3.onnx')

ort_inputs = {
    "input": np.random.randint(0, 255, (10, 28, 28), dtype=np.uint8),
}

ort_output = ort_session.run(['output'], ort_inputs)
ort_output[0].shape

(10,)

Y, como comentaba antes, es importante evaluar el modelo para asegurarnos que lo hemos exportado bien.

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

def onnx_eval():
    with torch.no_grad():
        preds, labels = [], torch.tensor([])
        for imgs, _labels in dm.val_dataloader():
            ort_inputs = {
                "input": imgs.numpy(),
            }
            ort_output = ort_session.run(['output'], ort_inputs)[0]
            outputs = sigmoid(ort_output) > 0.5
            preds += outputs.astype(int).tolist()
            labels = torch.cat([labels, _labels])
    acc = (np.array(preds) == labels.numpy()).astype(float).mean()
    return acc

onnx_eval()

0.95

En este caso obtenemos la misma métrica de evaluación que la obtenida con el checkpoint, así que podemos estar tranquilos de que el modelo se comportará bien en producción.

Versionando modelos

Una vez hemos exportado nuestro modelo y hemos verificado que funciona bien podemos versionarlo de la misma manera que hicimos con nuestro dataset.

Para ello, primero añadimos la carpeta models a nuestro repositorio.

dvc add models

De momento solo tenemos un modelo, el que acabamos de exportar, y podemos sincornizarlo con el repositorio remoto (en el que ya viven varias versiones de nuestro dataset) de la siguiente manera

dvc push

De esta manera, cualquier persona con accesso al proyecto en el que estamos trabajando podrá recuperar este modelo con el comando

dvc pull models.dvc

Puedes porbar que funcione eliminando la carpeta models y recuperándola con el comando anterior. Recuerda generar un nuevo tag y sincronizar con el repositorio de git.

git add .
git commit -m "primer modelo"
git push
git tag -a v3 -m "version 3"
git push origin --tags

A partir de ahora, al entrenar nuevos modelos, podemos añadrilos al repositorio con nuevo tag. Usar un modelo diferente en producción será tan sencillo como cambiar al tag adecuado, lo cual veremos más adelante.

Resumen

En este post hemos visto como podemos exportar nuestros modelos entrenados con Pytorch Lightninga ONNX. Esta representación intermedia nos permitirá desplegar nuestros modelos en entornos de producción gracias a la ONNXRuntime, la cual nos permite ejectuar el modelo de manera optimizada en multitud de entornos (Python, Web, Android, etc.). Además, hemos visto como trackear diferentes versiones de nuestro modelo con DVC de la misma manera que lo hacemos con nuestros datos.