Функция PyTorch torch.compile
Справочное руководство PyTorch torch
torch.compile— это функция в PyTorch, используемая для компиляции и оптимизации моделей, которая может значительно повысить производительность модели во время выполнения.
Определение функции
torch.compile(model, *, mode="default", backend="inductor", fullgraph=False, dynamic=None, options=None, disable=False)
Пример использования
Пример
import torch
# Определим простую модель
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()
# Скомпилировать модель
compiled_model = torch.compile(model)
# Использовать скомпилированную модель для инференса
x = torch.randn(1, 10)
output = compiled_model(x)
print(output.shape)
# Определим простую модель
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()
# Скомпилировать модель
compiled_model = torch.compile(model)
# Использовать скомпилированную модель для инференса
x = torch.randn(1, 10)
output = compiled_model(x)
print(output.shape)
Другие расширения