Función PyTorch torch.compile
Manual de referencia de Pytorch torch
torch.compileEs una función en PyTorch utilizada para compilar y optimizar modelos, que puede mejorar significativamente el rendimiento de ejecución del modelo.
Definición de la función
torch.compile(model, *, mode="default", backend="inductor", fullgraph=False, dynamic=None, options=None, disable=False)
Ejemplo de uso
Ejemplo
import torch
# Definir un modelo simple
class SimpleModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(10, 10)
def forward(self, x):
return self.linear(x)
model = SimpleModel()
# Compilar el modelo
compiled_model = torch.compile(model)
# Usar el modelo compilado para inferencia
x = torch.randn(1, 10)
output = compiled_model(x)
print(output.shape)
# Definir un modelo simple
class SimpleModel(torch.nn.Module):
def __init__(self):
super().__init__()
self.linear = torch.nn.Linear(10, 10)
def forward(self, x):
return self.linear(x)
model = SimpleModel()
# Compilar el modelo
compiled_model = torch.compile(model)
# Usar el modelo compilado para inferencia
x = torch.randn(1, 10)
output = compiled_model(x)
print(output.shape)
Otras extensiones