torch.nested

Введение

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

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

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

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

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

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

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

Создание

Создание простое и заключается в передаче списка тензоров в конструктор 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 можно использовать для сохранения истории автоградиента из тензоров, переданных в конструктор. Для получения дополнительной информации см. раздел Конструкторы и функции преобразования вложенных тензоров.

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

>>> 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

Если все размерности являются обычными, NestedTensor предполагается семантически неотличимым от обычного 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

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] не является копией, а представляет собой фрагмент базовой памяти, который представляет собой первый элемент или составную часть NestedTensor.

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

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

torch.nested.nested_tensor(tensor_list, *, dtype=None, device=None, requires_grad=False, pin_memory=False) [source]

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

Параметры:

tensor_list (Список[Тензор]) – список тензоров с одинаковой размерностью

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

dtype (torch.dtype, необязательно) – желаемый тип возвращаемого вложенного тензора. По умолчанию: если None, совпадает с torch.dtype левого тензора в списке.
device (torch.device, необязательно) – желаемое устройство возвращаемого вложенного тензора. По умолчанию: если None, совпадает с torch.device левого тензора в списке.
requires_grad (bool, необязательно) – Если автоградиент должен записывать операции над возвращаемым вложенным тензором. По умолчанию: 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 (Список[Тензор]) – список тензоров с одинаковой размерностью

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

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

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

Тензор

Пример:

>>> 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_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 (Кортеж[int]) – размер выходного тензора. Если задано, оно должно быть достаточно большим, чтобы содержать все данные; в противном случае, будет определяться путём взятия максимального размера каждого вложенного подтензора по каждой размерности.
out (Тензор, необязательно) – выходной тензор.

Пример:

>>> 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.

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

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

Операция PyTorch	Ограничения
`torch.matmul()`	Поддерживает матричное умножение двух (>= 3d) вложенных тензоров, где две последние размерности являются матричными размерностями, а ведущие (пакетные) размерности имеют одинаковый размер (т.е. поддержка широковещательной рассылки для пакетных размерностей пока отсутствует).
`torch.bmm()`	Поддерживает пакетное матричное умножение двух 3-мерных вложенных тензоров.
`torch.nn.Linear()`	Поддерживает 3-мерный вложенный вход и плотный 2-мерный весовой массив.
`torch.nn.functional.softmax()`	Поддерживает softmax вдоль всех измерений, кроме dim=0.
`torch.nn.Dropout()`	Поведение такое же, как и для обычных тензоров.
`torch.relu()`	Поведение такое же, как и для обычных тензоров.
`torch.gelu()`	Поведение такое же, как и для обычных тензоров.
`torch.add()`	Поддерживает поэлементное сложение двух вложенных тензоров. Поддерживает добавление скаляра к вложенному тензору.
`torch.mul()`	Поддерживает поэлементное умножение двух вложенных тензоров. Поддерживает умножение вложенного тензора на скаляр.
`torch.select()`	Поддерживает выбор по `dim=0` только (аналогично `nt[i]`).
`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`.

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