Fonction torch.compile de PyTorch
Manuel de référence PyTorch torch
torch.compileest une fonction dans PyTorch utilisée pour compiler et optimiser des modèles, ce qui peut considérablement améliorer les performances d'exécution du modèle.
Définition de fonction
torch.compile(model, *, mode="default", backend="inductor", fullgraph=False, dynamic=None, options=None, disable=False)
Exemple d'utilisation
Exemple
import torch
# Définir un modèle 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()
# Compiler le modèle
compiled_model = torch.compile(model)
# Utiliser le modèle compilé pour l'inférence
x = torch.randn(1, 10)
output = compiled_model(x)
print(output.shape)
# Définir un modèle 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()
# Compiler le modèle
compiled_model = torch.compile(model)
# Utiliser le modèle compilé pour l'inférence
x = torch.randn(1, 10)
output = compiled_model(x)
print(output.shape)
Autres extensions