torch.nested

Введение

Предупреждение

API PyTorch для вложенных тензоров находится на стадии прототипирования и будет изменён в ближайшем будущем.

NestedTensor позволяет пользователю упаковать список тензоров в единую, эффективную структуру данных.

Единственное ограничение для входных тензоров заключается в том, что их размерность должна совпадать.

Это позволяет более эффективно представлять метаданные и получать доступ к специально разработанным ядрам.

Одно из применений NestedTensor — выражение последовательных данных в различных областях. Хотя традиционный подход заключается в дополнении последовательностей переменной длины, NestedTensor позволяет пользователям обойтись без дополнения. API для вызова операций над вложенным тензором ничем не отличается от API для обычного torch.Tensor, что должно позволить бесшовную интеграцию с существующими моделями, основное различие заключается в создании входных данных.

Поскольку эта функция находится на стадии прототипирования, поддерживаемые операции по-прежнему ограничены. Тем не менее, мы приветствуем сообщения об ошибках, запросы на новые функции и вклад. Более подробную информацию о сотрудничестве можно найти в этом файле Readme.

Создание

Создание — простая задача, которая включает передачу списка тензоров в конструктор torch.nested.nested_tensor.

>>> a, b = torch.arange(3), torch.arange(5) + 3
>>> a
tensor([0, 1, 2])
>>> b
tensor([3, 4, 5, 6, 7])
>>> nt = torch.nested.nested_tensor([a, b])
>>> nt
nested_tensor([
  tensor([0, 1, 2]),
    tensor([3, 4, 5, 6, 7])
    ])

Тип данных, устройство и требуется ли вычисление градиента можно выбрать с помощью обычных ключевых аргументов.

>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32, device="cuda", requires_grad=True)
>>> nt
nested_tensor([
  tensor([0., 1., 2.], device='cuda:0', requires_grad=True),
  tensor([3., 4., 5., 6., 7.], device='cuda:0', requires_grad=True)
], device='cuda:0', requires_grad=True)

В духе torch.as_tensor, torch.nested.as_nested_tensor можно использовать для сохранения истории автоградиента из тензоров, переданных в конструктор. Для получения дополнительной информации см. раздел Функции создания и преобразования вложенных тензоров.

Для формирования допустимого вложенного тензора все переданные тензоры должны совпадать по размерности, но ни один из других атрибутов не должен совпадать.

>>> a = torch.randn(3, 50, 70) # image 1
>>> b = torch.randn(3, 128, 64) # image 2
>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32)
>>> nt.dim()
4

Если какая-либо из размерностей не совпадает, конструктор выдаёт ошибку.

>>> a = torch.randn(50, 128) # text 1
>>> b = torch.randn(3, 128, 64) # image 2
>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: All Tensors given to nested_tensor must have the same dimension. Found dimension 3 for Tensor at index 1 and dimension 2 for Tensor at index 0.

Обратите внимание, что переданные тензоры копируются в непрерывный участок памяти. Результирующий NestedTensor выделяет новую память для их хранения и не сохраняет ссылку.

В настоящее время мы поддерживаем только один уровень вложенности, то есть простой плоский список тензоров. В будущем мы можем добавить поддержку нескольких уровней вложенности, например, список, состоящий полностью из списков тензоров. Обратите внимание, что для этого расширения важно поддерживать одинаковый уровень вложенности в записях, чтобы результирующий NestedTensor имел чётко определённую размерность. Если у вас есть потребность в этой функции, пожалуйста, не стесняйтесь открыть запрос на новую функцию, чтобы мы могли отследить её и спланировать соответствующим образом.

Размер

Несмотря на то, что NestedTensor не поддерживает .size() (или .shape), он поддерживает .size(i), если размерность i является регулярной.

>>> a = torch.randn(50, 128) # text 1
>>> b = torch.randn(32, 128) # text 2
>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32)
>>> nt.size(0)
2
>>> nt.size(1)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: Given dimension 1 is irregular and does not have a size.
>>> nt.size(2)
128

Если все размерности являются регулярными, вложенный тензор должен быть семантически неотличимым от обычного torch.Tensor.

>>> a = torch.randn(20, 128) # text 1
>>> nt = torch.nested.nested_tensor([a, a], dtype=torch.float32)
>>> nt.size(0)
2
>>> nt.size(1)
20
>>> nt.size(2)
128
>>> torch.stack(nt.unbind()).size()
torch.Size([2, 20, 128])
>>> torch.stack([a, a]).size()
torch.Size([2, 20, 128])
>>> torch.equal(torch.stack(nt.unbind()), torch.stack([a, a]))
True

В будущем мы можем упростить обнаружение этого условия и обеспечить бесшовное преобразование.

Пожалуйста, откройте запрос на новую функцию, если у вас есть потребность в этом (или любой другой связанной функции).

Разъединение

unbind позволяет получить представление составных элементов.

>>> import torch
>>> a = torch.randn(2, 3)
>>> b = torch.randn(3, 4)
>>> nt = torch.nested.nested_tensor([a, b], dtype=torch.float32)
>>> nt
nested_tensor([
  tensor([[ 1.2286, -1.2343, -1.4842],
          [-0.7827,  0.6745,  0.0658]]),
  tensor([[-1.1247, -0.4078, -1.0633,  0.8083],
          [-0.2871, -0.2980,  0.5559,  1.9885],
          [ 0.4074,  2.4855,  0.0733,  0.8285]])
])
>>> nt.unbind()
(tensor([[ 1.2286, -1.2343, -1.4842],
        [-0.7827,  0.6745,  0.0658]]), tensor([[-1.1247, -0.4078, -1.0633,  0.8083],
        [-0.2871, -0.2980,  0.5559,  1.9885],
        [ 0.4074,  2.4855,  0.0733,  0.8285]]))
>>> nt.unbind()[0] is not a
True
>>> nt.unbind()[0].mul_(3)
tensor([[ 3.6858, -3.7030, -4.4525],
        [-2.3481,  2.0236,  0.1975]])
>>> nt
nested_tensor([
  tensor([[ 3.6858, -3.7030, -4.4525],
          [-2.3481,  2.0236,  0.1975]]),
  tensor([[-1.1247, -0.4078, -1.0633,  0.8083],
          [-0.2871, -0.2980,  0.5559,  1.9885],
          [ 0.4074,  2.4855,  0.0733,  0.8285]])
])

Обратите внимание, что nt.unbind()[0] не является копией, а является фрагментом базовой памяти, представляющей первый элемент или составную часть вложенного тензора.

Функции создания и преобразования вложенных тензоров

Следующие функции относятся к вложенным тензорам:

torch.nested.nested_tensor(tensor_list, *, dtype=None, device=None, requires_grad=False, pin_memory=False) → Tensor

Создаёт вложенный тензор без истории автоградиента (также известный как «листовой тензор», см. Механизмы автоградиента) из tensor_list списка тензоров.

Параметры

tensor_list (List[array_like]) – список тензоров или то, что можно передать в torch.tensor,
размерность. (где каждый элемент списка имеет одинаковую) –

Ключевые аргументы

dtype (torch.dtype, необязательно) – желаемый тип возвращаемого вложенного тензора. По умолчанию: если None, то тот же torch.dtype, что и у левого тензора в списке.
device (torch.device, необязательно) – желаемое устройство возвращаемого вложенного тензора. По умолчанию: если None, то то же torch.device, что и у левого тензора в списке.
requires_grad (bool, необязательно) – Если autograd должен записывать операции над возвращённым вложенным тензором. По умолчанию: False.
pin_memory (bool, необязательно) – Если установлено, возвращаемый вложенный тензор будет выделен в закреплённой памяти. Работает только для тензоров CPU. По умолчанию: False.

Пример:

>>> a = torch.arange(3, dtype=torch.float, requires_grad=True)
>>> b = torch.arange(5, dtype=torch.float, requires_grad=True)
>>> nt = torch.nested.nested_tensor([a, b], requires_grad=True)
>>> nt.is_leaf
True

torch.nested.as_nested_tensor(tensor_list, dtype=None, device=None) [source]

Создаёт вложенный тензор, сохраняя историю автоградиента из tensor_list списка тензоров.

Примечание

Тензоры внутри списка всегда копируются этой функцией из-за текущей семантики вложенных тензоров.

Параметры

tensor_list (List[Tensor]) – список тензоров с одинаковым числом измерений

Ключевые аргументы

dtype (torch.dtype, необязательно) – желаемый тип возвращаемого вложенного тензора. По умолчанию: если None, то тот же torch.dtype, что и у левого тензора в списке.
device (torch.device, необязательно) – желаемое устройство возвращаемого вложенного тензора. По умолчанию: если None, то то же torch.device, что и у левого тензора в списке.

Тип возвращаемого значения

Tensor

Пример:

>>> a = torch.arange(3, dtype=torch.float, requires_grad=True)
>>> b = torch.arange(5, dtype=torch.float, requires_grad=True)
>>> nt = torch.nested.as_nested_tensor([a, b])
>>> nt.is_leaf
False
>>> fake_grad = torch.nested.nested_tensor([torch.ones_like(a), torch.zeros_like(b)])
>>> nt.backward(fake_grad)
>>> a.grad
tensor([1., 1., 1.])
>>> b.grad
tensor([0., 0., 0., 0., 0.])

torch.nested.to_padded_tensor(input, padding, output_size=None, out=None) → Tensor

Возвращает новый (невложенный) тензор, заполняя input вложенный тензор. Ведущие элементы будут заполнены данными, а хвостовые — дополнены.

Предупреждение

to_padded_tensor() всегда копирует базовые данные, так как вложенные и невложенные тензоры отличаются по структуре памяти.

Параметры

padding (float) – Значение заполнения для хвостовых элементов.

Ключевые аргументы

output_size (Tuple[int]) – Размер выходного тензора. Если задано, оно должно быть достаточно большим, чтобы содержать все вложенные данные; в противном случае вычисляется путём взятия максимального размера каждого вложенного подтензора по каждому измерению.
out (Tensor, необязательно) – выходной тензор.

Пример:

>>> nt = torch.nested.nested_tensor([torch.randn((2, 5)), torch.randn((3, 4))])
nested_tensor([
  tensor([[ 1.6862, -1.1282,  1.1031,  0.0464, -1.3276],
          [-1.9967, -1.0054,  1.8972,  0.9174, -1.4995]]),
  tensor([[-1.8546, -0.7194, -0.2918, -0.1846],
          [ 0.2773,  0.8793, -0.5183, -0.6447],
          [ 1.8009,  1.8468, -0.9832, -1.5272]])
])
>>> pt_infer = torch.nested.to_padded_tensor(nt, 0.0)
tensor([[[ 1.6862, -1.1282,  1.1031,  0.0464, -1.3276],
         [-1.9967, -1.0054,  1.8972,  0.9174, -1.4995],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000]],
        [[-1.8546, -0.7194, -0.2918, -0.1846,  0.0000],
         [ 0.2773,  0.8793, -0.5183, -0.6447,  0.0000],
         [ 1.8009,  1.8468, -0.9832, -1.5272,  0.0000]]])
>>> pt_large = torch.nested.to_padded_tensor(nt, 1.0, (2, 4, 6))
tensor([[[ 1.6862, -1.1282,  1.1031,  0.0464, -1.3276,  1.0000],
         [-1.9967, -1.0054,  1.8972,  0.9174, -1.4995,  1.0000],
         [ 1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000],
         [ 1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000]],
        [[-1.8546, -0.7194, -0.2918, -0.1846,  1.0000,  1.0000],
         [ 0.2773,  0.8793, -0.5183, -0.6447,  1.0000,  1.0000],
         [ 1.8009,  1.8468, -0.9832, -1.5272,  1.0000,  1.0000],
         [ 1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000]]])
>>> pt_small = torch.nested.to_padded_tensor(nt, 2.0, (2, 2, 2))
RuntimeError: Value in output_size is less than NestedTensor padded size. Truncation is not supported.

Поддерживаемые операции

В этом разделе мы подытоживаем операции, которые в настоящее время поддерживаются для вложенных тензоров, и любые ограничения, с которыми они связаны.

Операция PyTorch	Ограничения
`torch.matmul()`	Поддерживает матричное умножение между двумя (>= 3d) вложенными тензорами, где две последние размерности — это размерности матриц, а ведущие (пакетные) размерности имеют одинаковый размер (т.е. поддержка трансляции для пакетных размерностей пока отсутствует).
`torch.bmm()`	Поддерживает пакетное матричное умножение двух 3-мерных вложенных тензоров.
`torch.nn.Linear()`	Поддерживает 3-мерный вложенный вход и плотную 2-мерную матрицу весов.
`torch.nn.functional.softmax()`	Поддерживает softmax по всем размерностям, кроме dim=0.
`torch.nn.Dropout()`	Поведение такое же, как и для обычных тензоров.
`torch.Tensor.masked_fill()`	Поведение такое же, как и для обычных тензоров.
`torch.relu()`	Поведение такое же, как и для обычных тензоров.
`torch.gelu()`	Поведение такое же, как и для обычных тензоров.
`torch.silu()`	Поведение такое же, как и для обычных тензоров.
`torch.abs()`	Поведение такое же, как и для обычных тензоров.
`torch.sgn()`	Поведение такое же, как и для обычных тензоров.
`torch.logical_not()`	Поведение такое же, как и для обычных тензоров.
`torch.neg()`	Поведение такое же, как и для обычных тензоров.
`torch.sub()`	Поддерживает поэлементное вычитание двух вложенных тензоров.
`torch.add()`	Поддерживает поэлементное сложение двух вложенных тензоров. Поддерживает сложение скаляра с вложенным тензором.
`torch.mul()`	Поддерживает поэлементное умножение двух вложенных тензоров. Поддерживает умножение вложенного тензора на скаляр.
`torch.select()`	Поддерживает выбор по всем размерностям.
`torch.clone()`	Поведение такое же, как и для обычных тензоров.
`torch.detach()`	Поведение такое же, как и для обычных тензоров.
`torch.unbind()`	Поддерживает отсоединение только по `dim=0`.
`torch.reshape()`	Поддерживает изменение формы с сохранением размера `dim=0` (т.е. количество вложенных тензоров изменить нельзя). В отличие от обычных тензоров, размер `-1` здесь означает, что размер наследуется. В частности, единственный допустимый размер для нерегулярной размерности — `-1`. Вывод размера пока не реализован, поэтому для новых размерностей размер нельзя `-1`.
`torch.Tensor.reshape_as()`	Аналогичные ограничения, как и для `reshape`.
`torch.transpose()`	Поддерживает транспонирование всех размерностей, кроме `dim=0`.
`torch.Tensor.view()`	Правила для новой формы аналогичны правилам для `reshape`.
`torch.empty_like()`	Аналогичное поведение обычных тензоров; возвращает новый пустой вложенный тензор (т.е. с неинициализированными значениями), соответствующий вложенной структуре входных данных.
`torch.randn_like()`	Аналогичное поведение обычных тензоров; возвращает новый вложенный тензор со значениями, случайным образом инициализированными в соответствии со стандартным нормальным распределением, соответствующим вложенной структуре входных данных.
`torch.zeros_like()`	Аналогичное поведение обычных тензоров; возвращает новый вложенный тензор со всеми нулевыми значениями, соответствующий вложенной структуре входных данных.
`torch.nn.LayerNorm()`	Аргумент `normalized_shape` ограничен, чтобы не распространяться на нерегулярные размерности NestedTensor.

© 2024, PyTorch Contributors
PyTorch has a BSD-style license, as found in the LICENSE file.
https://pytorch.org/docs/2.1/nested.html