Inferencia utilizando torch.compile(): el presente con PyTorch 2.X

Hace unas clases me dejó de funcionar DBpedia y leí en uno de los aportes que se podía hacer con AG\_NEWS. Así que lo hice con AG\_NEWS. Estos son los AG\_NEWS\_labels: `AG_NEWS_labels = {    0:"World",    1:"Sports",    2:"Business",    3:"Sci/Tech"}`

Les dejo un ejemplo del dataset AG\_NEWS: 'The Race is On: Second Private Team Sets Launch Date for Human Spaceflight (SPACE.com) SPACE.com - TORONTO, Canada -- A second\\\team of rocketeers competing for the #36;10 million Ansari X Prize, a contest for\\\privately funded suborbital space flight, has officially announced the first\\\launch date for its manned rocket.' Pd: Su categoria es 'Sci/Tech'

En PyTorch 2.X, se introduce `torch.compile()`, que optimiza el modelo utilizando compiladores JIT (just-in-time) como parte de su proceso de inferencia y entrenamiento. Esto permite una ejecución más rápida al aplicar optimizaciones como fusión de operadores, eliminación de código redundante, y otras mejoras a nivel de rendimiento. La idea es que puedas aprovechar este tipo de optimización con solo una línea adicional en tu código. Aquí te muestro cómo puedes usar `torch.compile()` en la \*\*inferencia\*\* de un modelo de clasificación de texto con PyTorch 2.X. \### Paso 1: Configuración del modelo Supongamos que ya tienes un modelo entrenado de clasificación de texto como en el ejemplo anterior: ```python \# Modelo de clasificación de texto class TextClassificationModel(nn.Module): def \_\_init\_\_(self, vocab\_size, embed\_dim, num\_class): super(TextClassificationModel, self).\_\_init\_\_() self.embedding = nn.EmbeddingBag(vocab\_size, embed\_dim, sparse=True) self.fc = nn.Linear(embed\_dim, num\_class) self.init\_weights() def init\_weights(self): initrange = 0.5 self.embedding.weight.data.uniform\_(-initrange, initrange) self.fc.weight.data.uniform\_(-initrange, initrange) self.fc.bias.data.zero\_() def forward(self, text, offsets): embedded = self.embedding(text, offsets) return self.fc(embedded) \# Definimos el vocabulario y las dimensiones del modelo VOCAB\_SIZE = len(vocab) EMBED\_DIM = 64 NUM\_CLASS = len(set(\[label for (label, text) in train\_iter])) model = TextClassificationModel(VOCAB\_SIZE, EMBED\_DIM, NUM\_CLASS) ``` \### Paso 2: Optimización del modelo para inferencia con `torch.compile()` Con PyTorch 2.X, puedes mejorar la inferencia compilando el modelo de manera directa con `torch.compile()`: ```python import torch \# Compilación del modelo para optimización compiled\_model = torch.compile(model) \# Asumiendo que tienes un tensor de entrada para la inferencia \# Ejemplo: texto de entrada tokenizado y convertidos a tensores def predict(text, offsets): \# Colocamos el modelo en modo evaluación compiled\_model.eval() with torch.no\_grad(): # Desactivamos la actualización de gradientes output = compiled\_model(text, offsets) return output.argmax(1).item() \# Simulación de un tensor de texto de prueba example\_text = torch.tensor(\[1, 2, 3, 4, 5], dtype=torch.int64) # Un ejemplo tokenizado example\_offsets = torch.tensor(\[0], dtype=torch.int64) \# Inferencia optimizada prediccion = predict(example\_text, example\_offsets) print(f"Predicción: {prediccion}") ``` \### Paso 3: Entrenamiento también con `torch.compile()` Si también quisieras optimizar el \*\*entrenamiento\*\* del modelo, puedes envolver el modelo con `torch.compile()` de la misma forma: ```python \# Compilamos el modelo para optimizar el entrenamiento compiled\_model = torch.compile(model) \# Entrenamiento con el modelo compilado for epoch in range(epochs): train(train\_dataloader, compiled\_model, criterion, optimizer) acc = evaluate(test\_dataloader, compiled\_model, criterion) print(f'Epoch {epoch+1}: Test Accuracy: {acc:.4f}') ``` \### Paso 4: Consideraciones \- `torch.compile()` es especialmente útil para modelos complejos o para ejecutar el modelo en hardware acelerado como GPUs. \- No afecta la precisión del modelo, pero puede mejorar la velocidad de ejecución. \- El compilador JIT subyacente puede aplicar optimizaciones en tiempo de ejecución, lo que hace que la primera ejecución sea más lenta, pero las posteriores sean más rápidas. \### Conclusión Con PyTorch 2.X y `torch.compile()`, puedes fácilmente optimizar tu modelo tanto para \*\*inferencia\*\* como para \*\*entrenamiento\*\* con cambios mínimos en el código. Esta nueva característica facilita la integración de optimizaciones a nivel de compilador sin necesidad de reescribir todo el modelo o flujo de datos.

Está genial. Dado que, como a varios no me funcionó el DBPedia, y estoy usando el AG\_NEWS; le pedí a ChatGPT que me escribiera tres titulares de noticias, etiquetándolas con las categorías del dataset. Le pasé esos titulares al modelo y: 3/3.

Les dejo el diccionario de etiquetas:

DBpedia_label = {
                  1: 'Company',
                  2: 'EducationalInstitution',
                  3: 'Artist',
                  4: 'Athlete',
                  5: 'OfficeHolder',
                  6: 'MeanOfTransportation',
                  7: 'Building',
                  8: 'NaturalPlace',
                  9: 'Village',
                  10: 'Animal',
                  11: 'Plant',
                  12: 'Album',
                  13: 'Film',
                  14: 'WrittenWork'
                }

y el texto ejemplo:

ejemplo_1 = “Nithari is a village in the western part of the state of Uttar Pradesh India bordering on New Delli”

¿torch.compile tiene algún tipo de cache? percibo que la primera vez que se ejecuta la función predict, toma bastante mas tiempo que cuando se ejecuta una segunda o X vez.