Функция unfold

Функция F.unfold из модуля torch.nn.functional извлекает скользящие локальные блоки из входного тензора. Она полезна для операций, требующих доступа к локальным окрестностям, например, при реализации пользовательских свёрточных слоёв, обработке изображений или анализе паттернов. Функция принимает входной тензор размерностью (N, C, H, W), размер ядра kernel_size, шаг stride и отступ padding. Результатом является тензор формы (N, C * kernel_size[0] * kernel_size[1], L), где L - количество извлечённых блоков.

Синтаксис

torch.nn.functional.unfold(
	input,
	kernel_size,
	stride=1,
	padding=0,
	dilation=1
)

Параметры

Функция unfold принимает следующие параметры:

input - входной тензор размерностью (N, C, H, W), где N - размер батча, C - число каналов, H и W - высота и ширина.
kernel_size - размер ядра (окна), может быть числом или кортежем (kernel_height, kernel_width).
stride - шаг скольжения, может быть числом или кортежем (stride_height, stride_width). По умолчанию 1.
padding - отступ, добавляемый к краям тензора, может быть числом или кортежем. По умолчанию 0.
dilation - коэффициент расширения ядра. По умолчанию 1.

Пример работы с изображением

Давайте извлечём блоки размером 2x2 из одноканального изображения размером 4x4 с шагом 1:

import torch
import torch.nn.functional as F

t = torch.tensor([
	[1, 2, 3, 4],
	[5, 6, 7, 8],
	[9, 10, 11, 12],
	[13, 14, 15, 16]
])
t = t.unsqueeze(0).unsqueeze(0)  # (1, 1, 4, 4)

res = F.unfold(t, kernel_size=2, stride=1)
print(res.shape)
print(res)

Результат выполнения кода:

torch.Size([1, 4, 9])
tensor([[
	[ 1.,  2.,  3.,  5.,  6.,  7.,  9., 10., 11.],
	[ 2.,  3.,  4.,  6.,  7.,  8., 10., 11., 12.],
	[ 5.,  6.,  7.,  9., 10., 11., 13., 14., 15.],
	[ 6.,  7.,  8., 10., 11., 12., 14., 15., 16.]
]])

Пример с многоканальным тензором

Рассмотрим извлечение блоков из трёхканального изображения размером 5x5 с шагом 2:

import torch
import torch.nn.functional as F

t = torch.randn(1, 3, 5, 5)

res = F.unfold(t, kernel_size=3, stride=2)
print(res.shape)

Результат выполнения кода:

torch.Size([1, 27, 4])

Пример с паддингом

Использование паддинга позволяет извлекать блоки на краях тензора. Добавим отступ 1 к изображению 3x3 и извлечём блоки 3x3 с шагом 1:

import torch
import torch.nn.functional as F

t = torch.tensor([
	[1, 2, 3],
	[4, 5, 6],
	[7, 8, 9]
])
t = t.unsqueeze(0).unsqueeze(0)  # (1, 1, 3, 3)

res = F.unfold(t, kernel_size=3, stride=1, padding=1)
print(res.shape)

Результат выполнения кода:

torch.Size([1, 9, 9])

Смотрите также

функцию F.fold,
которая выполняет обратное преобразование блоков в тензор
функцию F.conv2d,
которая применяет свёртку к входному тензору
функцию F.max_pool2d,
которая выполняет операцию максимального пулинга
функцию F.avg_pool2d,
которая выполняет операцию среднего пулинга