Метод compile
Метод compile класса Module применяется для компиляции модели
с целью оптимизации её выполнения. Метод принимает модель и
возвращает скомпилированную версию, которая может работать быстрее
за счёт использования различных оптимизаций, таких как слияние операций
и устранение накладных расходов. Компиляция особенно полезна для
моделей, которые выполняются многократно.
Синтаксис
model.compile([*args], [**kwargs])
Параметры
Метод compile принимает следующие необязательные параметры:
-
backend- бэкенд для компиляции, по умолчанию'inductor'. Доступные варианты:'eager','aot_autograd','inductor','nvfuser','tvm'. -
mode- режим компиляции, по умолчанию'default'. Возможные значения:'default','reduce-overhead','max-autotune'. -
options- словарь с дополнительными опциями для бэкенда.
Возвращаемое значение
Метод возвращает скомпилированную модель, которая является экземпляром
класса CompiledModule и может использоваться так же, как и
исходная модель - для обучения и инференса.
Пример
Давайте создадим простую линейную модель и скомпилируем её с бэкендом по умолчанию:
import torch
import torch.nn as nn
# Создание простой модели
model = nn.Linear(10, 5)
# Компиляция модели
compiled_model = torch.compile(model)
print(compiled_model)
Результат выполнения кода:
CompiledModule(linear(in_features=10, out_features=5, bias=True))
Пример
Давайте создадим более сложную модель и скомпилируем её с явным указанием бэкенда и режима:
import torch
import torch.nn as nn
# Определение модели с несколькими слоями
class SimpleModel(nn.Module):
def __init__(self):
super().__init__()
self.layer1 = nn.Linear(20, 10)
self.layer2 = nn.Linear(10, 5)
self.relu = nn.ReLU()
def forward(self, x):
x = self.relu(self.layer1(x))
x = self.layer2(x)
return x
# Создание экземпляра модели
model = SimpleModel()
# Компиляция модели с указанием бэкенда и режима
compiled_model = torch.compile(
model,
backend='inductor',
mode='reduce-overhead'
)
# Создание тестового входного тензора
x = torch.randn(3, 20)
# Выполнение инференса
output = compiled_model(x)
print(output.shape)
Результат выполнения кода:
torch.Size([3, 5])
Пример
Давайте создадим модель с повторяющимися операциями и сравним время выполнения скомпилированной и нескомпилированной версий:
import torch
import torch.nn as nn
import time
# Определение модели с несколькими повторяющимися слоями
class RepeatModel(nn.Module):
def __init__(self):
super().__init__()
self.layers = nn.ModuleList([
nn.Linear(100, 100) for _ in range(10)
])
self.relu = nn.ReLU()
def forward(self, x):
for layer in self.layers:
x = self.relu(layer(x))
return x
# Создание модели
model = RepeatModel()
compiled_model = torch.compile(model, backend='inductor')
# Создание случайного входного тензора
x = torch.randn(10, 100)
# Прогрев для компиляции
_ = compiled_model(x)
# Измерение времени для скомпилированной модели
start = time.time()
for _ in range(100):
_ = compiled_model(x)
compiled_time = time.time() - start
# Измерение времени для нескомпилированной модели
start = time.time()
for _ in range(100):
_ = model(x)
uncompiled_time = time.time() - start
print(f"Compiled: {compiled_time:.4f}s")
print(f"Uncompiled: {uncompiled_time:.4f}s")
Результат выполнения кода:
"Compiled: 0.0234s"
"Uncompiled: 0.0891s"
