Usando Visual Transformers en Nodejs

Hacer que un modelo finetuned de machine learning basado en VITs, corra en un entorno de nodejs no es tarea fácil. Casi todas las herramientas de las que se dispone en el universo de python, no están disponibles, y el soporte en la web es casi nulo.

En lo personal, aunque bien puede que la mayoría de los proyectos no necesiten pasar por un proceso como el que yo indico, pues python puede hacer todo lo que de nodejs podemos esperar, y por tanto es un tanto incensario migrar tus modelos a el último, en este proyecto necesitaba que fuese nodejs quien manejara el acceso al modelo.

Una vez entrenado

Este articulo se basa en la version de google/vit-base-patch16-224 con la tarea de clasificación de imágenes, pero los conceptos pueden que funcionen bien con otros modelos de la misma familias.

El primer paso es exportar a ONNX, y para ello aprovecharemos la función de pytorch, torch.onnx que nos pedirá el modelo, un ejemplo del tipo de dato que debe recibir el modelo y los nombres de las variables de entrada y salida.

torch.onnx.export(modelo, input, input_names=["pixel_values"], output_names=["logics"])

La variable input la puedes construir como siempre, pero al final, si es un tensor de pytorch tendrás que convertirla a numpy. Tambien asegurate de que tenga el shape correcto.

input = input_tensor.numpy()

Con onnx para python puedes probar que la exportación funciono correctamente, pasandole un input de prueba.

Ya exportamos el modelo, ¿ahora que?

En nodejs necesitaremos el runtime de onnx, que podemos instalar con npm con el siguiente código:

npm i onnxjs -s

Luego necesitamos cargar nuestro modelos y crear una sesión:

const url = "./path_to_your_model.onnx";
await session.loadModel(url);

El modelo debe recibir una imagen, en este caso de dimensiones 224X224 de forma/shape [1,3,224,224], con valores que van desde -1 a 1. En python este trabajo ya esta resuelto gracias a la ayuda de ViTFeatureExtractor, pero ya que no contamos con esta librería en nodejs tendremos que crear nuestro propio convertidor.

Para simplificar, nuestro ViTFeatureExtractorJs tendra que hacer lo siguiente:

Cargar la imagen desde el disco
Redimensionar la imagen a 224 pixeles de ancho por 224 pixeles de alto
Hacer una array con 3 valores, uno para cada color, 224 columnas y 224 líneas o [224,224,3]
Normalizar todos los componentes del color del pixel, de un rango entre 0 y 255 a un rango de entre -1 y 1.
Transponer las columnas a 3,2,1, con lo que el shape de la foto debe quedar en [3,224,224].
Convertir a un Float32Array(3*224*224) si es que aun no lo has hecho.

Una vez hemos hecho la anterior, nos queda crear el parámetro input para nuestro modelo usando la clase Tensor de nuestra librería onnx.

const inputs = [
new Tensor(float_array, "float32", [1, 3, 224, 224]),
];
const outputMap = await session.run(inputs);

Finalmente, tendrás un objeto enumerable con las probabilidades para cada clase ordenadas desde el valor mas probable al menos provable. Por si no lo has intuido, tendrás que exportar desde el modelo de python el mapa model.config.id2label para poder determinar que vale cada integer.