torch.sparse

Зачем и когда использовать разреженность

По умолчанию PyTorch хранит torch.Tensor хранит элементы непрерывно в физической памяти. Это приводит к эффективной реализации различных алгоритмов обработки массивов, которые требуют быстрого доступа к элементам.

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

За прошедшие годы были разработаны различные разреженные форматы хранения, такие как COO, CSR/CSC, полуструктурированные, LIL и т. д. Хотя они различаются по точным схемам расположения, все они сжимают данные с помощью эффективного представления элементов с нулевым значением. Мы называем несжатые значения указанными в отличие от неуказанных, сжатых элементов.

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

Как и многие другие оптимизации производительности, разреженные форматы хранения не всегда являются предпочтительными. При попытке использовать разреженные форматы для вашего случая вы можете обнаружить, что время выполнения увеличивается, а не уменьшается.

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

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

Обзор функциональности

Мы стремимся к простоте создания разреженного тензора из заданного плотного тензора, предоставив процедуры преобразования для каждого макета.

В следующем примере мы преобразуем 2D тензор с плотным (уступающим) макетом по умолчанию в 2D тензор, поддерживаемый макетом памяти COO. В этом случае хранятся только значения и индексы элементов, отличных от нуля.

>>> a = torch.tensor([[0, 2.], [3, 0]])
>>> a.to_sparse()
tensor(indices=tensor([[0, 1],
                       [1, 0]]),
       values=tensor([2., 3.]),
       size=(2, 2), nnz=2, layout=torch.sparse_coo)

В настоящее время PyTorch поддерживает COO, CSR, CSC, BSR и BSC.

У нас также есть прототип реализации для поддержки :ref: semi-structured sparsity<sparse-semi-structured-docs>. Дополнительную информацию см. в ссылках.

Обратите внимание, что мы предоставляем незначительные обобщения этих форматов.

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

В настоящее время мы предлагаем очень простую версию пакетирования, где каждый компонент разреженного формата сам пакетируется. Это также требует одинакового количества указанных элементов на запись пакета. В этом примере мы создаём 3D (пакетированный) CSR-тензор из 3D плотного тензора.

>>> t = torch.tensor([[[1., 0], [2., 3.]], [[4., 0], [5., 6.]]])
>>> t.dim()
3
>>> t.to_sparse_csr()
tensor(crow_indices=tensor([[0, 1, 3],
                            [0, 1, 3]]),
       col_indices=tensor([[0, 0, 1],
                           [0, 0, 1]]),
       values=tensor([[1., 2., 3.],
                      [4., 5., 6.]]), size=(2, 2, 2), nnz=3,
       layout=torch.sparse_csr)

Размеры плотных тензоров: С другой стороны, некоторые данные, такие как вложения графов, могут быть лучше представлены как разреженные коллекции векторов, а не скаляров.

В этом примере мы создаём 3D гибридный COO-тензор с 2 разреженными и 1 плотной размерностью из 3D тензора с уступками. Если вся строка в 3D тензоре с уступками равна нулю, она не сохраняется. Однако если любые значения в строке отличны от нуля, они сохраняются полностью. Это уменьшает количество индексов, так как нам нужен один индекс на строку, а не на элемент. Но это также увеличивает объём хранения для значений. Поскольку могут быть выданы только строки, полностью равные нулю, а наличие любых значений, отличных от нуля, приводит к сохранению всей строки.

>>> t = torch.tensor([[[0., 0], [1., 2.]], [[0., 0], [3., 4.]]])
>>> t.to_sparse(sparse_dim=2)
tensor(indices=tensor([[0, 1],
                       [1, 1]]),
       values=tensor([[1., 2.],
                      [3., 4.]]),
       size=(2, 2, 2), nnz=2, layout=torch.sparse_coo)

Обзор операторов

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

Мы активно увеличиваем покрытие операторов для разреженных тензоров. Пользователи пока не должны ожидать поддержки такого же уровня, как для плотных тензоров. Список см. в нашей документации по операторам.

>>> b = torch.tensor([[0, 0, 1, 2, 3, 0], [4, 5, 0, 6, 0, 0]])
>>> b_s = b.to_sparse_csr()
>>> b_s.cos()
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: unsupported tensor layout: SparseCsr
>>> b_s.sin()
tensor(crow_indices=tensor([0, 3, 6]),
       col_indices=tensor([2, 3, 4, 0, 1, 3]),
       values=tensor([ 0.8415,  0.9093,  0.1411, -0.7568, -0.9589, -0.2794]),
       size=(2, 6), nnz=6, layout=torch.sparse_csr)

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

>>> b_s.to_dense().cos()
tensor([[ 1.0000, -0.4161],
        [-0.9900,  1.0000]])

Мы знаем, что некоторые пользователи хотят игнорировать сжатые нули при операциях, таких как cos, а не сохранять точную семантику операции. Для этого мы можем указать на torch.masked и его MaskedTensor, который, в свою очередь, также поддерживается и работает с разреженными форматами хранения и ядрами.

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

>>> a + b.to_sparse()
tensor([[0., 3.],
        [3., 0.]])

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

Разреженные полуструктурированные тензоры

Полуструктурированная разреженность — это макет разреженных данных, который был впервые представлен в архитектуре NVIDIA Ampere. Он также известен как тонкозернистая структурированная разреженность или структурированная разреженность 2:4.

Этот разреженный макет хранит n элементов из каждых 2n элементов, причём n определяется шириной типа данных (dtype) тензора. Наиболее часто используемый тип данных — float16, где n=2, поэтому и термин «структурированная разреженность 2:4».

Полуструктурированная разреженность более подробно описана в этой статье блога NVIDIA.

В PyTorch полуструктурированная разреженность реализована через подкласс Tensor. Наследуя подкласс, мы можем переопределить __torch_dispatch__ , что позволит нам использовать более быстрые разреженные ядра при выполнении умножения матриц. Мы также можем хранить тензор в его сжатом виде внутри подкласса, чтобы уменьшить объём памяти.

В этом сжатом виде разреженный тензор сохраняется путём сохранения только указанных элементов и некоторых метаданных, которые кодируют маску.

Для разреженных тензоров 2:4 метаданные — незначительные — всего 2 бита на указанный элемент.

Здесь мы разбираем, как вычислить коэффициент сжатия (размер плотного / размер разреженного) для разреженного тензора 2:4.

Пусть (r, c) = tensor.shape и e = bitwidth(tensor.dtype), так что e = 16 для torch.float16 и torch.bfloat16 и e = 8 для torch.int8.

Используя эти вычисления, мы можем определить общий объём памяти как для исходного плотного, так и для нового разреженного представления.

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

Используя эту формулу, мы получаем, что коэффициент сжатия составляет 56,25% для torch.float16 или torch.bfloat16, и 62,5% для torch.int8.

Создание разреженных полуструктурированных тензоров

Вы можете преобразовать плотный тензор в разреженный полуструктурированный тензор, просто используя функцию torch.to_sparse_semi_structured.

Обратите также внимание, что мы поддерживаем только тензоры CUDA, так как совместимость оборудования для полуструктурированной разреженности ограничена графическими процессорами NVIDIA.

Ниже приведены поддерживаемые типы данных для полуструктурированной разреженности. Обратите внимание, что каждый тип данных имеет свои ограничения формы и коэффициент сжатия.

Чтобы создать разреженный полуструктурированный тензор, начните с создания обычного плотного тензора, который соответствует разреженному формату 2:4 (или полуструктурированному). Для этого мы разделяем небольшой фрагмент 1x4, чтобы создать тензор float16 размером 16x16. После этого мы можем вызвать функцию to_sparse_semi_structured для его сжатия для ускорения вывода.

>>> from torch.sparse import to_sparse_semi_structured
>>> A = torch.Tensor([0, 0, 1, 1]).tile((128, 32)).half().cuda()
tensor([[0., 0., 1.,  ..., 0., 1., 1.],
        [0., 0., 1.,  ..., 0., 1., 1.],
        [0., 0., 1.,  ..., 0., 1., 1.],
        ...,
        [0., 0., 1.,  ..., 0., 1., 1.],
        [0., 0., 1.,  ..., 0., 1., 1.],
        [0., 0., 1.,  ..., 0., 1., 1.]], device='cuda:0', dtype=torch.float16)
>>> A_sparse = to_sparse_semi_structured(A)
SparseSemiStructuredTensor(shape=torch.Size([128, 128]), transposed=False, values=tensor([[1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        ...,
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.],
        [1., 1., 1.,  ..., 1., 1., 1.]], device='cuda:0', dtype=torch.float16), metadata=tensor([[-4370, -4370, -4370,  ..., -4370, -4370, -4370],
        [-4370, -4370, -4370,  ..., -4370, -4370, -4370],
        [-4370, -4370, -4370,  ..., -4370, -4370, -4370],
        ...,
        [-4370, -4370, -4370,  ..., -4370, -4370, -4370],
        [-4370, -4370, -4370,  ..., -4370, -4370, -4370],
        [-4370, -4370, -4370,  ..., -4370, -4370, -4370]], device='cuda:0',
dtype=torch.int16))

Операции с разреженными полуструктурированными тензорами

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

• torch.addmm(bias, dense, sparse.t())
• torch.mm(dense, sparse)
• torch.mm(sparse, dense)
• aten.linear.default(dense, sparse, bias)
• aten.t.default(sparse)
• aten.t.detach(sparse)

Чтобы использовать эти операции, просто передайте результат to_sparse_semi_structured(tensor) вместо tensor после того, как в вашем тензоре будут 0 в полуструктурированном разреженном формате, например:

>>> a = torch.Tensor([0, 0, 1, 1]).tile((64, 16)).half().cuda()
>>> b = torch.rand(64, 64).half().cuda()
>>> c = torch.mm(a, b)
>>> a_sparse = to_sparse_semi_structured(a)
>>> torch.allclose(c, torch.mm(a_sparse, b))
True

Ускорение nn.Linear с полуструктурированной разреженностью

Вы можете ускорить слои линейного преобразования в вашей модели, если веса уже являются разреженными полуструктурированными, всего несколькими строками кода:

>>> input = torch.rand(64, 64).half().cuda()
>>> mask = torch.Tensor([0, 0, 1, 1]).tile((64, 16)).cuda().bool()
>>> linear = nn.Linear(64, 64).half().cuda()
>>> linear.weight = nn.Parameter(to_sparse_semi_structured(linear.weight.masked_fill(~mask, 0)))

Разреженные COO-тензоры

PyTorch реализует так называемый формат координат, или формат COO, как один из форматов хранения для реализации разреженных тензоров. В формате COO указанные элементы хранятся в виде кортежей индексов элементов и соответствующих значений. В частности,

• индексы указанных элементов собираются в тензоре indices размера (ndim, nse) и с типом элементов torch.int64,
• соответствующие значения собираются в тензоре values размера (nse,) и с произвольным целочисленным или плавающей точкой типом элементов,

где ndim — размерность тензора, а nse — количество указанных элементов.

Создание

Разреженный COO-тензор можно создать, предоставив два тензора индексов и значений, а также размер разреженного тензора (если он не может быть выведен из тензоров индексов и значений) в функцию torch.sparse_coo_tensor().

Предположим, что мы хотим определить разреженный тензор с элементом 3 в позиции (0, 2), элементом 4 в позиции (1, 0) и элементом 5 в позиции (1, 2). Неуказанные элементы предполагаются имеющими одинаковое значение, значение заполнения, которое по умолчанию равно нулю. Тогда мы напишем:

>>> i = [[0, 1, 1],
         [2, 0, 2]]
>>> v =  [3, 4, 5]
>>> s = torch.sparse_coo_tensor(i, v, (2, 3))
>>> s
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([3, 4, 5]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo)
>>> s.to_dense()
tensor([[0, 0, 3],
        [4, 0, 5]])

Обратите внимание, что вход i НЕ является списком кортежей индексов. Если вы хотите записать свои индексы таким образом, вы должны транспонировать их перед передачей в конструктор разреженного тензора:

>>> i = [[0, 2], [1, 0], [1, 2]]
>>> v =  [3,      4,      5    ]
>>> s = torch.sparse_coo_tensor(list(zip(*i)), v, (2, 3))
>>> # Or another equivalent formulation to get s
>>> s = torch.sparse_coo_tensor(torch.tensor(i).t(), v, (2, 3))
>>> torch.sparse_coo_tensor(i.t(), v, torch.Size([2,3])).to_dense()
tensor([[0, 0, 3],
        [4, 0, 5]])

Пустой разреженный COO-тензор можно создать, указав только его размер:

>>> torch.sparse_coo_tensor(size=(2, 3))
tensor(indices=tensor([], size=(2, 0)),
       values=tensor([], size=(0,)),
       size=(2, 3), nnz=0, layout=torch.sparse_coo)

Разреженные гибридные COO-тензоры

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

Гибридный COO-тензор PyTorch расширяет разреженный COO-тензор, позволяя тензору values быть многомерным тензором, так что у нас есть:

• индексы указанных элементов собираются в тензоре indices размера (sparse_dims, nse) и с типом элементов torch.int64,
• соответствующие (тензорные) значения собираются в тензоре values размера (nse, dense_dims) и с произвольным целочисленным или плавающей точкой типом элементов.

Предположим, что мы хотим создать (2 + 1)-мерный тензор с элементом [3, 4] в позиции (0, 2), элементом [5, 6] в позиции (1, 0) и элементом [7, 8] в позиции (1, 2). Мы напишем

>>> i = [[0, 1, 1],
         [2, 0, 2]]
>>> v =  [[3, 4], [5, 6], [7, 8]]
>>> s = torch.sparse_coo_tensor(i, v, (2, 3, 2))
>>> s
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([[3, 4],
                      [5, 6],
                      [7, 8]]),
       size=(2, 3, 2), nnz=3, layout=torch.sparse_coo)

>>> s.to_dense()
tensor([[[0, 0],
         [0, 0],
         [3, 4]],
        [[5, 6],
         [0, 0],
         [7, 8]]])

В общем случае, если s является разреженным COO-тензором и M = s.sparse_dim(), K = s.dense_dim(), то у нас есть следующие инварианты:

• M + K == len(s.shape) == s.ndim — размерность тензора равна сумме числа разреженных и плотных измерений,
• s.indices().shape == (M, nse) — разреженные индексы хранятся явно,
• s.values().shape == (nse,) + s.shape[M : M + K] — значения гибридного тензора являются K-мерными тензорами,
• s.values().layout == torch.strided — значения хранятся как строенные тензоры.

Разреженные неслитые COO-тензоры

Формат PyTorch разреженных COO-тензоров допускает разреженные *неслитые* тензоры, где в индексах могут быть дублирующие координаты; в этом случае интерпретация заключается в том, что значение по этому индексу является суммой всех дублирующих значений. Например, можно указать несколько значений, 3 и 4, для одного и того же индекса 1, что приводит к 1-мерному неслитому тензору:

>>> i = [[1, 1]]
>>> v =  [3, 4]
>>> s=torch.sparse_coo_tensor(i, v, (3,))
>>> s
tensor(indices=tensor([[1, 1]]),
       values=tensor(  [3, 4]),
       size=(3,), nnz=2, layout=torch.sparse_coo)

в то время как процесс слияния накопит многозначные элементы в одно значение с использованием суммирования:

>>> s.coalesce()
tensor(indices=tensor([[1]]),
       values=tensor([7]),
       size=(3,), nnz=1, layout=torch.sparse_coo)

В общем случае результат метода torch.Tensor.coalesce() — это разреженный тензор со следующими свойствами:

• индексы указанных элементов тензора уникальны,
• индексы упорядочены лексикографически,
• torch.Tensor.is_coalesced() возвращает True.

>>> a = torch.sparse_coo_tensor([[1, 1]], [5, 6], (2,))
>>> b = torch.sparse_coo_tensor([[0, 0]], [7, 8], (2,))
>>> a + b
tensor(indices=tensor([[0, 0, 1, 1]]),
       values=tensor([7, 8, 5, 6]),
       size=(2,), nnz=4, layout=torch.sparse_coo)

Работа с разреженными COO-тензорами

Рассмотрим следующий пример:

>>> i = [[0, 1, 1],
         [2, 0, 2]]
>>> v =  [[3, 4], [5, 6], [7, 8]]
>>> s = torch.sparse_coo_tensor(i, v, (2, 3, 2))

Как упоминалось выше, разреженный COO-тензор — это экземпляр torch.Tensor, и чтобы отличить его от экземпляров Tensor, использующих другой макет, можно использовать свойства torch.Tensor.is_sparse или torch.Tensor.layout:

>>> isinstance(s, torch.Tensor)
True
>>> s.is_sparse
True
>>> s.layout == torch.sparse_coo
True

Количество разреженных и плотных измерений можно получить, используя методы torch.Tensor.sparse_dim() и torch.Tensor.dense_dim() соответственно. Например:

>>> s.sparse_dim(), s.dense_dim()
(2, 1)

Если s — это разреженный COO-тензор, то данные его формата COO можно получить, используя методы torch.Tensor.indices() и torch.Tensor.values().

>>> s.indices()
RuntimeError: Cannot get indices on an uncoalesced tensor, please call .coalesce() first

>>> s._indices()
tensor([[0, 1, 1],
        [2, 0, 2]])

Создание нового разреженного COO-тензора приводит к тензору, который не объединён:

>>> s.is_coalesced()
False

но можно создать объединённую копию разреженного COO-тензора, используя метод torch.Tensor.coalesce():

>>> s2 = s.coalesce()
>>> s2.indices()
tensor([[0, 1, 1],
       [2, 0, 2]])

При работе с не объединёнными разреженными COO-тензорами, необходимо учитывать аддитивный характер не объединённых данных: значения одних и тех же индексов являются членами суммы, вычисление которой даёт значение соответствующего элемента тензора. Например, скалярное умножение на разреженный не объединённый тензор можно реализовать, умножив все значения на скаляр, так как c * (a + b) == c * a + c * b верно. Однако любая нелинейная операция, например, квадратный корень, не может быть реализована путём применения операции к не объединённым данным, так как sqrt(a + b) == sqrt(a) + sqrt(b) в общем случае не верно.

Срезы (с положительным шагом) разреженного COO-тензора поддерживаются только для плотных измерений. Индексирование поддерживается для разреженных и плотных измерений:

>>> s[1]
tensor(indices=tensor([[0, 2]]),
       values=tensor([[5, 6],
                      [7, 8]]),
       size=(3, 2), nnz=2, layout=torch.sparse_coo)
>>> s[1, 0, 1]
tensor(6)
>>> s[1, 0, 1:]
tensor([6])

В PyTorch значение заполнения разреженного тензора не может быть задано явно и, как правило, предполагается равным нулю. Однако существуют операции, которые могут интерпретировать значение заполнения по-разному. Например, torch.sparse.softmax() вычисляет softmax с предположением, что значение заполнения равно отрицательной бесконечности.

Разреженные сжатые тензоры

Разреженные сжатые тензоры представляют собой класс разреженных тензоров, у которых общая особенность — сжатие индексов определённого измерения с использованием кодирования, которое позволяет выполнять определённые оптимизации на линейных алгебраических ядрах разреженных сжатых тензоров. Это кодирование основано на формате Compressed Sparse Row (CSR), который PyTorch разреженные сжатые тензоры расширяют, добавив поддержку партий разреженных тензоров, позволяя использовать многомерные значения тензора и хранить значения разреженного тензора в плотных блоках.

Разреженный CSR-тензор

Основное преимущество формата CSR по сравнению с форматом COO — более эффективное использование памяти и намного более быстрые вычисления, такие как умножение разреженной матрицы на вектор, используя бэкэнды MKL и MAGMA.

В простейшем случае (0 + 2 + 0)-мерный разреженный CSR-тензор состоит из трёх 1-мерных тензоров: crow_indices, col_indices и values:

• Тензор crow_indices состоит из сжатых индексов строк. Это 1-мерный тензор размера nrows + 1 (количество строк плюс 1). Последний элемент crow_indices — это количество указанных элементов, nse. Этот тензор кодирует индекс в values и col_indices в зависимости от того, где начинается заданная строка. Каждое последующее число в тензоре, вычтенное из предыдущего, обозначает количество элементов в заданной строке.
• Тензор col_indices содержит индексы столбцов каждого элемента. Это 1-мерный тензор размера nse.
• Тензор values содержит значения элементов CSR-тензора. Это 1-мерный тензор размера nse.

В общем случае (B + 2 + K)-мерный разреженный CSR-тензор состоит из двух (B + 1)-мерных тензоров индексов crow_indices и col_indices, и (1 + K)-мерного тензора values, такого что

• crow_indices.shape == (*batchsize, nrows + 1)
• col_indices.shape == (*batchsize, nse)
• values.shape == (nse, *densesize)

в то время как форма разреженного CSR-тензора — (*batchsize, nrows, ncols, *densesize), где len(batchsize) == B и len(densesize) == K.

Создание CSR-тензоров

Разреженные CSR-тензоры можно напрямую создать, используя функцию torch.sparse_csr_tensor(). Пользователь должен предоставить тензоры индексов строк и столбцов, а также тензор значений отдельно, при этом индексы строк должны быть указаны с использованием кодирования CSR. Аргумент size является необязательным и будет выведен из crow_indices и col_indices, если он не указан.

>>> crow_indices = torch.tensor([0, 2, 4])
>>> col_indices = torch.tensor([0, 1, 0, 1])
>>> values = torch.tensor([1, 2, 3, 4])
>>> csr = torch.sparse_csr_tensor(crow_indices, col_indices, values, dtype=torch.float64)
>>> csr
tensor(crow_indices=tensor([0, 2, 4]),
       col_indices=tensor([0, 1, 0, 1]),
       values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4,
       dtype=torch.float64)
>>> csr.to_dense()
tensor([[1., 2.],
        [3., 4.]], dtype=torch.float64)

Самый простой способ построения 2-мерного разреженного тензора CSR из строкового или разреженного COO-тензора — это использование метода torch.Tensor.to_sparse_csr(). Любые нули в (строковом) тензоре будут интерпретироваться как отсутствующие значения в разреженном тензоре:

>>> a = torch.tensor([[0, 0, 1, 0], [1, 2, 0, 0], [0, 0, 0, 0]], dtype=torch.float64)
>>> sp = a.to_sparse_csr()
>>> sp
tensor(crow_indices=tensor([0, 1, 3, 3]),
      col_indices=tensor([2, 0, 1]),
      values=tensor([1., 1., 2.]), size=(3, 4), nnz=3, dtype=torch.float64)

Операции с тензором CSR

Умножение разреженной матрицы на вектор можно выполнить с помощью метода tensor.matmul(). В настоящее время это единственная поддерживаемая математическая операция для тензоров CSR.

>>> vec = torch.randn(4, 1, dtype=torch.float64)
>>> sp.matmul(vec)
tensor([[0.9078],
        [1.3180],
        [0.0000]], dtype=torch.float64)

Разреженный тензор CSC

Формат разреженного CSC (Compressed Sparse Column) тензора реализует формат CSC для хранения 2-мерных тензоров с расширением для поддержки пакетов разреженных тензоров CSC и значений, являющихся многомерными тензорами.

Аналогично разреженным CSR-тензорам, разреженный CSC-тензор состоит из трех тензоров: ccol_indices, row_indices и values:

• Тензор ccol_indices состоит из сжатых индексов столбцов. Это (B + 1)-мерный тензор формы (*batchsize, ncols + 1). Последний элемент — количество указанных элементов, nse. Этот тензор кодирует индекс в values и row_indices в зависимости от того, где начинается данный столбец. Каждое последующее число в тензоре, вычтенное из предыдущего, обозначает количество элементов в данном столбце.
• Тензор row_indices содержит индексы строк каждого элемента. Это (B + 1)-мерный тензор формы (*batchsize, nse).
• Тензор values содержит значения элементов тензора CSC. Это (1 + K)-мерный тензор формы (nse, *densesize).

Создание тензоров CSC

Разреженные CSC-тензоры можно непосредственно создать, используя функцию torch.sparse_csc_tensor(). Пользователь должен предоставить тензоры индексов строк и столбцов и значений отдельно, где индексы столбцов должны быть указаны с использованием кодирования сжатия CSR. Аргумент size является необязательным и будет выведен из тензоров row_indices и ccol_indices при его отсутствии.

>>> ccol_indices = torch.tensor([0, 2, 4])
>>> row_indices = torch.tensor([0, 1, 0, 1])
>>> values = torch.tensor([1, 2, 3, 4])
>>> csc = torch.sparse_csc_tensor(ccol_indices, row_indices, values, dtype=torch.float64)
>>> csc
tensor(ccol_indices=tensor([0, 2, 4]),
       row_indices=tensor([0, 1, 0, 1]),
       values=tensor([1., 2., 3., 4.]), size=(2, 2), nnz=4,
       dtype=torch.float64, layout=torch.sparse_csc)
>>> csc.to_dense()
tensor([[1., 3.],
        [2., 4.]], dtype=torch.float64)

Разреженные CSC-тензоры размерности (0 + 2 + 0) можно создать из любого двумерного тензора, используя метод torch.Tensor.to_sparse_csc(). Любые нули в (строковом) тензоре будут интерпретироваться как отсутствующие значения в разреженном тензоре:

>>> a = torch.tensor([[0, 0, 1, 0], [1, 2, 0, 0], [0, 0, 0, 0]], dtype=torch.float64)
>>> sp = a.to_sparse_csc()
>>> sp
tensor(ccol_indices=tensor([0, 1, 2, 3, 3]),
       row_indices=tensor([1, 1, 0]),
       values=tensor([1., 2., 1.]), size=(3, 4), nnz=3, dtype=torch.float64,
       layout=torch.sparse_csc)

Разреженный тензор BSR

Формат разреженного BSR (Block compressed Sparse Row) тензора реализует формат BSR для хранения двумерных тензоров с расширением для поддержки пакетов разреженных тензоров BSR и блоков многомерных тензоров в качестве значений.

Разреженный BSR-тензор состоит из трех тензоров: crow_indices, col_indices и values:

• Тензор crow_indices состоит из сжатых индексов строк. Это (B + 1)-мерный тензор формы (*batchsize, nrowblocks + 1). Последний элемент — количество указанных блоков, nse. Этот тензор кодирует индекс в values и col_indices в зависимости от того, где начинается данный столбец-блок. Каждое последующее число в тензоре, вычтенное из предыдущего, обозначает количество блоков в данной строке.
• Тензор col_indices содержит индексы столбцов-блоков каждого элемента. Это (B + 1)-мерный тензор формы (*batchsize, nse).
• Тензор values содержит значения элементов разреженного BSR-тензора, собранные в двумерные блоки. Это (1 + 2 + K)-мерный тензор формы (nse, nrowblocks, ncolblocks, *densesize).

Создание тензоров BSR

Разреженные BSR-тензоры можно непосредственно создать, используя функцию torch.sparse_bsr_tensor(). Пользователь должен предоставить тензоры индексов строк-блоков и столбцов-блоков и значений отдельно, где индексы строк-блоков должны быть указаны с использованием кодирования сжатия CSR. Аргумент size является необязательным и будет выведен из тензоров crow_indices и col_indices при его отсутствии.

>>> crow_indices = torch.tensor([0, 2, 4])
>>> col_indices = torch.tensor([0, 1, 0, 1])
>>> values = torch.tensor([[[0, 1, 2], [6, 7, 8]],
...                        [[3, 4, 5], [9, 10, 11]],
...                        [[12, 13, 14], [18, 19, 20]],
...                        [[15, 16, 17], [21, 22, 23]]])
>>> bsr = torch.sparse_bsr_tensor(crow_indices, col_indices, values, dtype=torch.float64)
>>> bsr
tensor(crow_indices=tensor([0, 2, 4]),
       col_indices=tensor([0, 1, 0, 1]),
       values=tensor([[[ 0.,  1.,  2.],
                       [ 6.,  7.,  8.]],
                      [[ 3.,  4.,  5.],
                       [ 9., 10., 11.]],
                      [[12., 13., 14.],
                       [18., 19., 20.]],
                      [[15., 16., 17.],
                       [21., 22., 23.]]]),
       size=(4, 6), nnz=4, dtype=torch.float64, layout=torch.sparse_bsr)
>>> bsr.to_dense()
tensor([[ 0.,  1.,  2.,  3.,  4.,  5.],
        [ 6.,  7.,  8.,  9., 10., 11.],
        [12., 13., 14., 15., 16., 17.],
        [18., 19., 20., 21., 22., 23.]], dtype=torch.float64)

Разреженные BSR-тензоры размерности (0 + 2 + 0) можно создать из любого двумерного тензора, используя метод torch.Tensor.to_sparse_bsr(), который также требует указания размера блока значений:

>>> dense = torch.tensor([[0, 1, 2, 3, 4, 5],
...                       [6, 7, 8, 9, 10, 11],
...                       [12, 13, 14, 15, 16, 17],
...                       [18, 19, 20, 21, 22, 23]])
>>> bsr = dense.to_sparse_bsr(blocksize=(2, 3))
>>> bsr
tensor(crow_indices=tensor([0, 2, 4]),
       col_indices=tensor([0, 1, 0, 1]),
       values=tensor([[[ 0,  1,  2],
                       [ 6,  7,  8]],
                      [[ 3,  4,  5],
                       [ 9, 10, 11]],
                      [[12, 13, 14],
                       [18, 19, 20]],
                      [[15, 16, 17],
                       [21, 22, 23]]]), size=(4, 6), nnz=4,
       layout=torch.sparse_bsr)

Разреженный тензор BSC

Формат разреженного BSC (Block compressed Sparse Column) тензора реализует формат BSC для хранения двумерных тензоров с расширением для поддержки пакетов разреженных тензоров BSC и блоков многомерных тензоров в качестве значений.

Разреженный BSC-тензор состоит из трех тензоров: ccol_indices, row_indices и values:

• Тензор ccol_indices состоит из сжатых индексов столбцов. Это (B + 1)-мерный тензор формы (*batchsize, ncolblocks + 1). Последний элемент — количество указанных блоков, nse. Этот тензор кодирует индекс в values и row_indices в зависимости от того, где начинается данный блок строки. Каждое последующее число в тензоре, вычтенное из предыдущего, обозначает количество блоков в данном столбце.
• Тензор row_indices содержит индексы блоков строк каждого элемента. Это (B + 1)-мерный тензор формы (*batchsize, nse).
• Тензор values содержит значения элементов разреженного BSC-тензора, собранные в двумерные блоки. Это (1 + 2 + K)-мерный тензор формы (nse, nrowblocks, ncolblocks, *densesize).

Создание тензоров BSC

Разреженные BSC-тензоры можно непосредственно создать, используя функцию torch.sparse_bsc_tensor(). Пользователь должен предоставить тензоры индексов блоков строк и столбцов и значений отдельно, где индексы блоков столбцов должны быть указаны с использованием кодирования сжатия CSR. Аргумент size является необязательным и будет выведен из тензоров ccol_indices и row_indices при его отсутствии.

>>> ccol_indices = torch.tensor([0, 2, 4])
>>> row_indices = torch.tensor([0, 1, 0, 1])
>>> values = torch.tensor([[[0, 1, 2], [6, 7, 8]],
...                        [[3, 4, 5], [9, 10, 11]],
...                        [[12, 13, 14], [18, 19, 20]],
...                        [[15, 16, 17], [21, 22, 23]]])
>>> bsc = torch.sparse_bsc_tensor(ccol_indices, row_indices, values, dtype=torch.float64)
>>> bsc
tensor(ccol_indices=tensor([0, 2, 4]),
       row_indices=tensor([0, 1, 0, 1]),
       values=tensor([[[ 0.,  1.,  2.],
                       [ 6.,  7.,  8.]],
                      [[ 3.,  4.,  5.],
                       [ 9., 10., 11.]],
                      [[12., 13., 14.],
                       [18., 19., 20.]],
                      [[15., 16., 17.],
                       [21., 22., 23.]]]), size=(4, 6), nnz=4,
       dtype=torch.float64, layout=torch.sparse_bsc)

Инструменты для работы с разреженными сжатыми тензорами

Все разреженные сжатые тензоры — CSR, CSC, BSR и BSC — концептуально очень похожи тем, что данные их индексов разделены на две части: так называемые сжатые индексы, использующие кодирование CSR, и так называемые обычные индексы, ортогональные сжатым индексам. Это позволяет различным инструментам для этих тензоров использовать одни и те же реализации, которые параметризуются структурой тензора.

Создание разреженных сжатых тензоров

Разреженные CSR, CSC, BSR и CSC тензоры можно создать, используя функцию torch.sparse_compressed_tensor(), которая имеет тот же интерфейс, что и вышеупомянутые функции-конструкторы torch.sparse_csr_tensor(), torch.sparse_csc_tensor(), torch.sparse_bsr_tensor() и torch.sparse_bsc_tensor() соответственно, но с дополнительным требуемым аргументом layout. Следующий пример иллюстрирует способ создания CSR и CSC тензоров с использованием одних и тех же входных данных, указав соответствующий параметр структуры в функции torch.sparse_compressed_tensor():

>>> compressed_indices = torch.tensor([0, 2, 4])
>>> plain_indices = torch.tensor([0, 1, 0, 1])
>>> values = torch.tensor([1, 2, 3, 4])
>>> csr = torch.sparse_compressed_tensor(compressed_indices, plain_indices, values, layout=torch.sparse_csr)
>>> csr
tensor(crow_indices=tensor([0, 2, 4]),
       col_indices=tensor([0, 1, 0, 1]),
       values=tensor([1, 2, 3, 4]), size=(2, 2), nnz=4,
       layout=torch.sparse_csr)
>>> csc = torch.sparse_compressed_tensor(compressed_indices, plain_indices, values, layout=torch.sparse_csc)
>>> csc
tensor(ccol_indices=tensor([0, 2, 4]),
       row_indices=tensor([0, 1, 0, 1]),
       values=tensor([1, 2, 3, 4]), size=(2, 2), nnz=4,
       layout=torch.sparse_csc)
>>> (csr.transpose(0, 1).to_dense() == csc.to_dense()).all()
tensor(True)

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

Операции линейной алгебры

В следующей таблице обобщены поддерживаемые операции линейной алгебры для разреженных матриц, где структуры операндов могут различаться. Здесь T[layout] обозначает тензор с заданной структурой. Аналогично, M[layout] обозначает матрицу (2-мерный тензор PyTorch), а V[layout] обозначает вектор (1-мерный тензор PyTorch). Кроме того, f обозначает скаляр (вещественное число или 0-мерный тензор PyTorch), * — поэлементное умножение, а @ — умножение матриц.

В столбце «Разреженный градиент?» указано, поддерживает ли операция PyTorch обратное распространение относительно аргумента разреженной матрицы. Все операции PyTorch, за исключением torch.smm(), поддерживают обратное распространение относительно аргументов шаговой матрицы.

Методы тензоров и разреженные

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

Следующие методы специфичны для разреженных CSR тензоров и разреженных BSR тензоров:

Следующие методы специфичны для разреженных CSC тензоров и разреженных BSC тензоров:

Следующие методы тензора поддерживают разреженные COO тензоры:

add() add_() addmm() addmm_() any() asin() asin_() arcsin() arcsin_() bmm() clone() deg2rad() deg2rad_() detach() detach_() dim() div() div_() floor_divide() floor_divide_() get_device() index_select() isnan() log1p() log1p_() mm() mul() mul_() mv() narrow_copy() neg() neg_() negative() negative_() numel() rad2deg() rad2deg_() resize_as_() size() pow() sqrt() square() smm() sspaddmm() sub() sub_() t() t_() transpose() transpose_() zero_()

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

Другие функции

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

cat() dstack() empty() empty_like() hstack() index_select() is_complex() is_floating_point() is_nonzero() is_same_size() is_signed() is_tensor() lobpcg() mm() native_norm() pca_lowrank() select() stack() svd_lowrank() unsqueeze() vstack() zeros() zeros_like()

Для управления проверкой инвариантов разрежённых тензоров см.:

Для использования разреженных тензоров с функцией gradcheck() см.:

Унарные функции

Мы стремимся поддерживать все унарные функции, сохраняющие нули.

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

Следующие операторы в настоящее время поддерживают в качестве входных данных разреженные тензоры COO/CSR/CSC/BSR/CSR.

abs() asin() asinh() atan() atanh() ceil() conj_physical() floor() log1p() neg() round() sin() sinh() sign() sgn() signbit() tan() tanh() trunc() expm1() sqrt() angle() isinf() isposinf() isneginf() isnan() erf() erfinv()