Экспорт моделей с помощью TorchScript

Масштабируемые модели ИИ с PyTorch Lightning

Sergiy Tkachuk

Director, GenAI Productivity

Что такое TorchScript?

Независим от Python
Эффективен в производственной среде
Примеры:
- Развёртывание на мобильных устройствах

Мобильный телефон

Что такое TorchScript?

Независим от Python
Эффективен в производственной среде
Примеры:
- Развёртывание на мобильных устройствах
- Высокопроизводительный инференс в производственных системах

Мобильный телефон и ноутбук с шестерёнками, символизирующими производственную среду

Преобразование моделей в TorchScript

Два способа преобразования:

torch.jit.trace: трассирует выполнение на основе примеров входных данных
torch.jit.script: компилирует модель, анализируя код Python

Когда использовать:

trace — для простых моделей
script — для моделей с управляющими конструкциями (например, циклами)

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def forward(self, x):
        return x * 2

model = SimpleModel()
scripted_model = torch.jit.script(model)

Сохранение и загрузка моделей TorchScript

Сохранение модели:
- torch.jit.save: сохраняет скриптованную модель в файл
Загрузка модели:
- torch.jit.load: загружает модель для инференса

# Save the model
torch.jit.save(scripted_mod,"model.pt")

# Load the model
loaded_model=torch.jit.load("model.pt")

Инференс с помощью TorchScript

Шаги:
- Загрузить модель TorchScript
- Передать входные данные для получения предсказаний
- Результаты идентичны предсказаниям PyTorch

Пример входных данных:

Входной тензор: [1.0, 2.0, 3.0]

Пример выходных данных:

Выходной тензор: [2.0, 4.0, 6.0]

# Perform inference
input_arr = [1.0, 2.0, 3.0]
input_tensor = torch.tensor(input_arr)

output = loaded_model(input_tensor)
print(output)

TorchScript: коротко о главном

torch.jit.trace: подходит для статических моделей
torch.jit.script: поддерживает динамические управляющие конструкции
torch.jit.save: сохраняет скриптованную модель
torch.jit.load: загружает модель для инференса

Давайте потренируемся!

Масштабируемые модели ИИ с PyTorch Lightning