Flujo de despliegue: PyTorch → Core ML → ANE

Parte de la ANE Knowledge Base. Fuentes: ambos artículos de Apple y el código de exportación en references/ml-ane-transformers (tutorial del README) y references/ml-vision-transformers-ane/export.py.

El pipeline es idéntico para modelos de texto y de visión:

PyTorch model (ANE-optimized, eval mode)
   → torch.jit.trace                      (TorchScript)
   → coremltools ct.convert(mlprogram)    (.mlpackage)
   → Xcode Performance tab                (verify ANE dispatch + measure)
   → ship in app, load asynchronously

1. Preparar el modelo#

• Construir/instanciar la variante optimizada para la ANE y cargar los pesos preentrenados (los hooks del state-dict gestionan Linear→Conv2d y el orden de LayerNorm — véase doc 03).
• model.eval() siempre; estos ports son solo de inferencia.
• Para modelos de HF: cargar el baseline con return_dict=False y torchscript=True (coremltools no soporta salidas de tipo dict).
• Envolver el tracing en torch.no_grad() (como hace export.py) para evitar el lastre de autograd.

2. Trazar a TorchScript#

# Text (fixed sequence length — pad to max_length)
tokenized = tokenizer(["sample"], return_tensors="pt",
                      max_length=128, padding="max_length")
traced = torch.jit.trace(model, (tokenized["input_ids"], tokenized["attention_mask"]))

# Vision (fixed input shape)
x = torch.rand((1, 3, 256, 256))
traced = torch.jit.trace(model, (x,))

El tracing fija shapes estáticos — elige la longitud de secuencia / resolución / tamaño de batch que vas a servir (los artefactos de Apple los codifican en el nombre de archivo: ..._seqLen128_batchSize1, ..._batch1_256x256_...). Exporta un package por cada shape que necesites.

3. Convertir con coremltools#

import coremltools as ct, numpy as np

mlpackage = ct.convert(
    traced,
    convert_to="mlprogram",                      # ML Program, not neuralnetwork
    inputs=[
        # text: token ids are int32
        ct.TensorType("input_ids", shape=(1, 128), dtype=np.int32),
        ct.TensorType("attention_mask", shape=(1, 128), dtype=np.int32),
        # vision: ct.TensorType("x", shape=x.shape)
    ],
    compute_units=ct.ComputeUnit.ALL,            # allow CPU+GPU+ANE (default)
)
mlpackage.save("model.mlpackage")

Puntos clave:

• convert_to="mlprogram" apunta al backend moderno ML Program (FP16 por defecto en la ANE).
• compute_units=ct.ComputeUnit.ALL deja que Core ML construya el plan híbrido; usa CPU_AND_NE durante la depuración si quieres detectar caídas a la GPU (fallbacks).
• Las entradas enteras (token ids, máscaras) se declaran como np.int32.

4. Verificar en el dispositivo: informes de rendimiento de Xcode#

1. Añade el .mlpackage como recurso en cualquier proyecto de Xcode.
2. Ábrelo → pestaña Performance → genera un informe en un dispositivo disponible localmente (el propio Mac o un iPhone/iPad conectado).
3. El informe muestra el despacho por op a cada compute-unit (CPU / GPU / ANE) y la latencia medida.

Cómo leerlo:

• Espera el tronco del transformer en la ANE. Algunas ops en CPU son normales — DistilBERT mantiene 4/606 ops (búsquedas de embeddings) en CPU porque ahí es genuinamente más rápido.
• El número de ops crece con el chunking → mayor tiempo de carga/compilación puntual; lo que mejora es la latencia en régimen estacionario. Carga el modelo de forma asíncrona al arrancar la app.
• Vuelve a ejecutar los informes en distintos dispositivos/versiones de OS objetivo — las decisiones de despacho pueden diferir según el chip y el OS (Apple publicó curvas para iPhone 12/13, iOS 15/16, M1 macOS 13).
• Si la latencia se mantiene plana mientras reduces la carga de trabajo, estás limitado por el ancho de banda → considera batches más grandes o cuantización/poda (Principio 4).

5. Checklist de errores comunes#

Síntoma	Causa probable	Solución
Ops que caen a GPU/CPU en mitad del grafo	Op no soportada o layout hostil (rango > 5, último eje pequeño)	reshapes de relevo 5D; layouts BC1S/NHWC; revisa cada op nueva en el informe
Salidas radicalmente incorrectas frente a PyTorch	FP16 (eps demasiado pequeño, máscara −inf/1e9, overflow)	eps ≥ 1e-7, máscaras −1e4, clamping opcional
Error de conversión sobre salidas dict	return_dict=True de HF	return_dict=False (+ torchscript=True)
El checkpoint no carga en el modelo optimizado	Shapes de pesos Linear vs Conv2d; orden de scale/bias en LayerNorm	pre-hooks de state-dict (linear_to_conv2d_map, corrección de bias/weight)
Gran latencia, primera carga lenta	cientos de ops fragmentadas compilándose	Acéptalo como coste puntual; carga asíncrona; cachea el modelo compilado
La latencia no baja con secuencias más cortas	Régimen limitado por ancho de banda	Aumenta el batch; cuantiza/poda los pesos
Explosión de padding / ralentización de 32×	Último eje pequeño o unitario	Reordena las dimensiones para que un eje grande (alineado a 64 bytes) quede al final; canales múltiplo de 32

6. Reproducir las exportaciones de Apple desde los repositorios clonados#

# Vision: exports tiny-moat-0 at 512x512 and 256x256, global and local attention
cd references/ml-vision-transformers-ane
pip install torch coremltools pytest timm
python export.py          # writes ./exported_model/*.mlpackage
pytest tests.py           # unit tests / usage examples

# Text: follow the README tutorial
cd ../ml-ane-transformers
pip install ane_transformers   # or: pip install -e .
# then run the DistilBERT tutorial code from README.md

Eso cierra el círculo: 01 restricciones de hardware → 02 principios → 03/04 arquitecturas aplicadas → este flujo de despliegue/verificación.