Exportar modelos PyTorch del pipeline KYC a formato ONNX para portabilidad y optimización
Convierte los modelos PyTorch del sistema de verificación de identidad (ArcFace para reconocimiento facial, MiniFASNet para liveness detection, YOLOv8 para detección de documentos) al formato ONNX. Esto permite desacoplar la inferencia del framework de entrenamiento, habilitando optimizaciones posteriores como cuantización y ejecución en ONNX Runtime o TensorRT.
Usa esta skill cuando necesites convertir modelos PyTorch a ONNX dentro del model_server_agent. Aplica cuando se tenga un modelo entrenado o fine-tuned y se requiera exportarlo para su despliegue en producción con runtimes optimizados, o cuando se necesite portabilidad entre diferentes plataformas de inferencia.
Cargar el modelo PyTorch con sus pesos entrenados:
import torch
from models.arcface import ArcFaceModel
model = ArcFaceModel()
model.load_state_dict(torch.load("weights/arcface_r100.pth", map_location="cpu"))
model.eval()
Crear un tensor dummy con las dimensiones de entrada correctas para cada modelo:
# ArcFace: batch x 3 x 112 x 112
dummy_input = torch.randn(1, 3, 112, 112)
# MiniFASNet (liveness): batch x 3 x 80 x 80
# YOLOv8 (documentos): batch x 3 x 640 x 640
Exportar el modelo a ONNX con nombres de entrada/salida descriptivos y ejes dinámicos para batching:
torch.onnx.export(
model,
dummy_input,
"models/arcface.onnx",
input_names=["face_image"],
output_names=["embedding"],
dynamic_axes={
"face_image": {0: "batch_size"},
"embedding": {0: "batch_size"}
},
opset_version=17,
do_constant_folding=True
)
Verificar la validez del modelo ONNX exportado:
import onnx
model_onnx = onnx.load("models/arcface.onnx")
onnx.checker.check_model(model_onnx)
print("Modelo ONNX válido")
Simplificar el grafo ONNX para eliminar nodos redundantes:
import onnxsim
model_simplified, check = onnxsim.simplify(model_onnx)
assert check, "La simplificación falló"
onnx.save(model_simplified, "models/arcface_simplified.onnx")
Validar numéricamente que las salidas del modelo ONNX coincidan con las del modelo PyTorch original:
import onnxruntime as ort
import numpy as np
session = ort.InferenceSession("models/arcface.onnx")
onnx_output = session.run(None, {"face_image": dummy_input.numpy()})[0]
with torch.no_grad():
torch_output = model(dummy_input).numpy()
np.testing.assert_allclose(torch_output, onnx_output, rtol=1e-3, atol=1e-5)
print("Validación numérica exitosa")
Documentar los metadatos del modelo exportado (versión, dimensiones, opset, fecha) en el registro de modelos del sistema.
opset_version=17 o superior para asegurar compatibilidad con las operaciones utilizadas por ArcFace y YOLOv8; versiones anteriores pueden no soportar ciertos operadores.