Векторизация скользящего окна в PyTorch: unfold вместо цикла

Как векторизовать скользящее окно для 1D тензоров в PyTorch: используем tensor.unfold и broadcasting вместо torch.roll и циклов. Пример кода и прирост скорости.

Скользящие окна по одномерным тензорам — типичный прием в анализе временных рядов, обработке сигналов и моделировании последовательностей. Простой цикл с torch.roll работает, но теряет в производительности и плохо масштабируется. Цель — посчитать результаты для всех окон за один векторизованный проход и собрать их в 2D‑тензор.

Базовая реализация через цикл

Следующий код обрабатывает 1D‑тензор, каждый раз беря четыре соседних значения, передавая их функции и сдвигая стартовую позицию на один элемент на каждой итерации. Промежуточные результаты складываются в тензор 5×4.

import torch


def apply_block(win: torch.tensor, scale: float) -> torch.tensor:
    return win * scale

win_width = 4
series = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
idx_pick = torch.tensor(list(range(win_width)))

chunks = []
for _ in range(5):
    slice_vals = torch.index_select(series, 0, idx_pick.view(-1))
    out_vec = apply_block(slice_vals, series[0])
    chunks.append(out_vec)
    series = torch.roll(series, -1)

stacked_loop = torch.stack(chunks)

В чем узкое место на самом деле

Сама задача — это классическое «скользящее окно»: на каждом шаге берутся четыре подряд идущих элемента. Если реализовать это на Python‑цикле с повторяющимися индексными операциями, работа уходит из быстрых векторизованных операций с тензорами. Лучше сразу построить двумерное представление всех окон и подать его в ту же функцию. Такое представление можно получить без копирования данных с помощью tensor.unfold — окна будут отображаться как строки матрицы.

Векторизованный подход с видом скользящего окна

Ниже та же операция выполняется за один вызов: строим двумерный «вид скользящих окон» и подаем его в функцию. Масштаб берется из первого элемента каждого окна и растягивается (broadcast) до нужных размеров.

import torch


def apply_block(win: torch.tensor, scale: float) -> torch.tensor:
    return win * scale

win_width = 4
base = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])

windows = base.unfold(0, win_width, 1)
vectorized_out = apply_block(windows, windows[:, 0:1])

# Дополнительная сверка с версией на цикле
check_series = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
idx_pick = torch.tensor(list(range(win_width)))
accum = []
for _ in range(5):
    sl = torch.index_select(check_series, 0, idx_pick.view(-1))
    accum.append(apply_block(sl, check_series[0]))
    check_series = torch.roll(check_series, -1)
stacked_loop = torch.stack(accum)

assert torch.allclose(stacked_loop, vectorized_out)

Почему это важно

Заменив циклы на стороне Python одной тензорной операцией, вы сохраняете вычисления внутри оптимизированного конвейера PyTorch. Код становится нагляднее: двумерная матрица окон явно выражает задумку, а распространение масштаба по второй размерности естественно соответствует структуре данных. Это особенно полезно при большом числе окон, когда накладные расходы циклов и повторное индексирование начинают заметно влиять на скорость.

Выводы

Когда нужны последовательные группы элементов тензора, создавайте вид скользящего окна через unfold и применяйте вычисление сразу ко всем окнам. Сохраняйте входы тензорными, чтобы при необходимости использовать torch.index_select для индексирования, а для согласования форм полагайтесь на broadcasting — без лишних копий. Такой подход сохраняет исходную логику и одновременно повышает читаемость и масштабируемость.