Библиотека дата-сайентиста | Data Science, Machine learning, анализ данных, машинное обучение

🛠 Команда дня: ускоряем динамический инференс

▶️ Проблема: у модели меняется размер входов — обычно это мешает трассировке и оптимизации.

✅ Решение: torch.jit.trace умеет кэшировать разные формы входов. TorchScript сохраняет шаблоны и переиспользует их, ускоряя инференс без потерь гибкости.

Пример:

import torch
from torch import nn

# Простая нейросеть
model = nn.Sequential(
    nn.Linear(128, 64),
    nn.ReLU(),
    nn.Linear(64, 10)
)

# Пример входа для трассировки
example_input = torch.randn(1, 128)

# Трассировка модели — включаем TorchScript
traced_model = torch.jit.trace(model, example_input)

# Первый вызов с новым размером — форма кэшируется
output1 = traced_model(torch.randn(32, 128))

# Второй вызов с такой же формой — используется кэш
output2 = traced_model(torch.randn(32, 128))

# Новый батч — снова кэшируется
output3 = traced_model(torch.randn(64, 128))

➡️ Что это даёт:
• Ускорение инференса за счёт shape specialization
• Гибкость — не нужно фиксировать размер входа
• TorchScript автоматически переиспользует кэш при повторении формы

Библиотека дата-сайентиста #буст

Please open Telegram to view this post

VIEW IN TELEGRAM

👍4❤2⚡1👏1

www.tgoop.com/dsproglib/6390

1.8K viewsedited  Apr 25 at 07:01

🛠 Команда дня: ускоряем динамический инференс

▶️ Проблема: у модели меняется размер входов — обычно это мешает трассировке и оптимизации.

✅ Решение: torch.jit.trace умеет кэшировать разные формы входов. TorchScript сохраняет шаблоны и переиспользует их, ускоряя инференс без потерь гибкости.

Пример:

import torch
from torch import nn

# Простая нейросеть
model = nn.Sequential(
    nn.Linear(128, 64),
    nn.ReLU(),
    nn.Linear(64, 10)
)

# Пример входа для трассировки
example_input = torch.randn(1, 128)

# Трассировка модели — включаем TorchScript
traced_model = torch.jit.trace(model, example_input)

# Первый вызов с новым размером — форма кэшируется
output1 = traced_model(torch.randn(32, 128))

# Второй вызов с такой же формой — используется кэш
output2 = traced_model(torch.randn(32, 128))

# Новый батч — снова кэшируется
output3 = traced_model(torch.randn(64, 128))