torch.sparse

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

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

Обратите внимание, что мы поддерживаем только тензоры CUDA, поскольку аппаратная поддержка полуструктурированной разреженности ограничена GPU 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, когда в вашем тензоре уже есть нули в полуструктурированном разреженном формате, например:

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

torch.sparse.semi_structured.to_sparse_semi_structured(original_tensor, transposed=False, alg_id=0) [исходный код]

Эта функция преобразует плотный тензор в полуструктурированный разреженный тензор. Она возвращает SparseSemiStructuredTensor — подкласс torch.Tensor.

Эта функция проверяет, что у плотного тензора правильные тип данных, размер, размерность и устройство. В настоящее время мы поддерживаем полуструктурированные разреженные тензоры только для двумерных тензоров CUDA. Кроме того, размер вашего тензора должен быть положительным числом, кратным минимальному размеру разреженного блока, указанному в _DTYPE_TO_SHAPE_CONSTRAINTS для каждого типа данных (float32, float16, bfloat16, int8).

Параметры:

• original_tensor (Tensor) – плотный тензор для преобразования
• transposed (bool, optional) – устаревший аргумент, который будет удалён в одном из следующих выпусков. Не используйте его.
• alg_id (int, optional) – идентификатор алгоритма для умножения матриц cuSPARSELt. По умолчанию равен 0. Его можно получить с помощью torch._cslt_sparse_mm_search.

Возвращает:

Разреженный полуструктурированный тензор, созданный из заданного original_tensor

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

SparseSemiStructuredTensor

Исключения:

None –

Пример

>>> 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]))
>>> A_sparse.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),
>>> A_sparse.indices()
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))

Создание

Разреженный тензор 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 — значения хранятся в виде тензоров со схемой хранения strided.

Несобранные разреженные тензоры COO

Формат разреженных тензоров COO в PyTorch допускает разреженные несобранные тензоры, в индексах которых могут повторяться координаты. В этом случае значением по такому индексу считается сумма всех значений с повторяющимися индексами. Например, можно указать несколько значений — 3 и 4 — для одного индекса 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, предполагая, что значение заполнения равно минус бесконечности.

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

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

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

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

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

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

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

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

>>> 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 для хранения двумерных тензоров и дополнен поддержкой пакетов разреженных тензоров 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).

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

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

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

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

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

>>> 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] обозначает матрицу (двумерный тензор PyTorch), а V[layout] — вектор (одномерный тензор PyTorch). Кроме того, f обозначает скаляр (число с плавающей точкой или нульмерный тензор PyTorch), * — поэлементное умножение, а @ — матричное умножение.

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

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

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

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

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

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

Следующие методы Tensor поддерживают разреженные тензоры 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() см.:

Унарные функции, сохраняющие ноль

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

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

В настоящее время следующие операторы поддерживают входные тензоры sparse 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()