torch.sparse

Создание

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

Разреженный тензор 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.

>>> 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 для хранения 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).

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

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

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

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

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

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

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

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

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

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