01.09.2025 12:13
badcasedaily1 0 338 Источник

Разработчики PyTorch предоставили модуль torch.sparse для работы с разреженными тензорами, где большинство элементов – нули. Зачем это нужно? Представьте матрицу смежности графа, сильно обрезанную сеть или облако точек – хранить такие данные плотным массивом без надобности расточительно. Разрежённая структура сохраняет только ненулевые элементы и их индексы, что сильно экономит память и ускоряет вычисления. Например, матрица размером 10,000 на 10,000 с 100 000 ненулевых float-значений в разрежённом COO-формате займёт не 400 МБ, а около 2 МБ.

Несмотря на перспективы, API разрежённых тензоров в PyTorch пока в бете и может менять крошечные детали. Будьте к этому готовы: часть операций поддерживается, часть – нет, и некоторые автоград-ячейки пока работают только для COO, а для CSR, например, градиент не считается. Но обо всём по порядку.

Основные форматы разрежённых тензоров

PyTorch поддерживает несколько форматов хранения sparse-данных: COO, CSR, CSC, а также блочно-разреженные BSR и BSC.

COO (coordinate) – самый базовый. Два 1D-тензора indices и values хранят индексы ненулей и их значения. indices имеет форму (ndim, nnz), где nnz – число ненулевых элементов, а values – (nnz,) (для скалярных значений) или (nnz, *dense_dims) для гибридных (tensor-valued) форм. Типы: индексы – int64, значения – любой числовой. Пустой COO-тензор можно создать, указав только размер. Пример создания и использования COO:

import torch
# создаём разреженный COO-тензор: ненулевые 3,4,5 в позициях (0,2),(1,0),(1,2) матрицы 2x3
i = [[0, 1, 1],   # строки ненулей
     [2, 0, 2]]   # столбцы ненулей
v = [3, 4, 5]     # значения ненулей
s = torch.sparse_coo_tensor(i, v, size=(2, 3))
print(s)
# tensor(indices=tensor([[0, 1, 1],
#                       [2, 0, 2]]),
#        values=tensor([3, 4, 5]),
#        size=(2, 3), nnz=3, layout=torch.sparse_coo)
print("Плотная матрица:\n", s.to_dense())
# Плотная матрица:
#  tensor([[0, 0, 3],
#          [4, 0, 5]])

Явно передаём индексы и значения. Метод to_dense() восстанавливает обычный (плотный) тензор. Если индексы неотсортированы или дублируются, у COO-тензора есть понятие uncoalesced и coalesced. В некой sense, некоалесцированный тензор может иметь несколько записей для одной и той же позиции (их интерпретируют как суммы). Метод s.coalesce() объединит дубликаты, просуммировав значения. Например:

# Пример: создаём COO с дублированием индексов (ненулевых элемент 1,1 два раза):
i = [[1, 1]]      # оба индекса (1,1)
v = [3, 4]        # два значения для одной позиции
s = torch.sparse_coo_tensor(i, v, (3,))  # одномерный 3-элементный тензор
print("Неcoalesced:", s)
# tensor(indices=tensor([[1, 1]]),
#        values=tensor([3, 4]),
#        size=(3,), nnz=2, layout=torch.sparse_coo)
s2 = s.coalesce()
print("Coalesced:", s2)
# tensor(indices=tensor([[1]]),
#        values=tensor([7]),    # 3+4=7
#        size=(3,), nnz=1, layout=torch.sparse_coo)

После coalesce() индексы становятся уникальными и отсортированными. Большинство операций с разрежёнными данными автоматически коалесцируются по необходимости (для линейных операций это безопасно). Следует лишь учесть: не коалесцированные тензоры могут не поддерживать все методы (например, нельзя просто взять .indices() без coalesce()). Учтите, что to_dense() всё равно вернёт правильный результат, не важно, coalesced тензор или нет.

CSR – хранит данные в CSR-формате: два вектора индексов (строк и столбцов) и значения. Конструкция: crow_indices размера (nrows+1,), col_indices длины nnz и values длины nnz. crow_indices[i] – это индекс в values/col_indices, с которого начинается строка i. Основное преимущество CSR – умножение разреженной матрицы на вектор/матрицу. В PyTorch множества sparse-операций в CSR идут быстрее, чем для COO. Пример создания CSR-тензора:

# Создаём разреженный CSR-тензор вручную для матрицы 2x2:
crow = [0, 2, 4]          # crow_indices: 0-я строка с 0 по 1, 1-я строка с 2 по 3 (nnz=4)
cols = [0, 1, 0, 1]      # индексы столбцов для каждого ненулевого значения
vals = [1, 2, 3, 4]      # сами значения
csr = torch.sparse_csr_tensor(
    torch.tensor(crow, dtype=torch.int64),
    torch.tensor(cols, dtype=torch.int64),
    torch.tensor(vals, dtype=torch.float32)
)
print(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.float32, layout=torch.sparse_csr)

Или проще: любая плотная матрица или COO-тензор A может быть сконвертирован в CSR одним методом A.to_sparse_csr(). Пример:

dense = torch.tensor([[0,1,0],[2,0,3]], dtype=torch.float32)
csr2 = dense.to_sparse_csr()  # плотный -> CSR
print(csr2)

В CSR автоматом воспринимаются все нули как отсутствующие элементы. Стандартно crow_indices и col_indices имеют тип int64, хотя для совместимости с MKL можно использовать int32. PyTorch требует, чтобы в CSR индексы были целыми (по умолчанию torch.int64).

CSR хорош для умножения разреженной матрицы на вектор или другую матрицу. В PyTorch для CSR доступны операции вида csr.matmul(vec) или torch.sparse.mm(csr, B). При умножении разреженная матрица плотная даётся обычная плотная матрица, а разрежённая разрежённая выдаёт разрежённую. Например:

# Умножение CSR на плотную матрицу
a = torch.tensor([[1.,0,2],[0,3,0]]).to_sparse_csr()
b = torch.tensor([[0., 1.], [2., 0.], [0., 0.]])
y = torch.sparse.mm(a, b)
print(y)
# tensor([[0., 1.],
#         [6., 0.]])

Здесь a – CSR тензор 2×3, b – плотная 3×2, результат – плотная 2×2 (как и Dense×Dense). При этом torch.sparse.mm поддерживает и CSR, и COO форматы. Учтите, что бэпроп (градиент) пока НЕ считается по CSR-матрице.

CSC – аналогично CSR, но сжатие по столбцам. Формат: тензоры ccol_indices (размер (ncols+1,)), row_indices (длины nnz) и values (длины nnz). CSR и CSC – транспонированные друг друга версии. PyTorch позволяет создать CSC напрямую через torch.sparse_csc_tensor. Вот пример:

ccol = torch.tensor([0, 2, 4], dtype=torch.int64)  # ccol_indices
rows = torch.tensor([0, 1, 0, 1], dtype=torch.int64)  # row_indices
vals = torch.tensor([1.,2.,3.,4.], dtype=torch.float32)
csc = torch.sparse_csc_tensor(ccol, rows, vals)
print(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.float32, layout=torch.sparse_csc)
print(csc.to_dense())
# tensor([[1., 3.],
#         [2., 4.]])

CSC выигрывает там же, где CSR – особенно при операциях по столбцам. Во многом, если нужна разреженная матричная операция, имеет смысл использовать CSR или CSC, а не COO напрямую.

BSR/BSC – разрежённые тензоры с блочной структурой. Аналогично CSR/CSC, но values – двухмерные блоки. PyTorch поддерживает torch.sparse_bsr_tensor и torch.sparse_bsc_tensor (с параметром blocksize). Блочные форматы экономят индексы, когда ненулевые элементы сгруппированы в равные блоки. Их использование специализируется на некоторых оптимизациях, но применяются реже. Словом, они есть, но для основного ознакомления хватит знаний про CSR/CSC.

Как конвертировать и проверять

Из dense в sparse: любой обычный (плотный) тензор X можно перевести в разрежённый формат. По умолчанию X.to_sparse() делает COO-сжатие (считаются все ненули). Похожее: X.to_sparse_coo(), X.to_sparse_csr(), X.to_sparse_csc() и т.д. Например:

dense = torch.tensor([[1,0,2],[0,3,0]])
sparse_coo = dense.to_sparse_coo()
print(sparse_coo)  # COO-версия
sparse_csr = dense.to_sparse_csr()
print(sparse_csr)  # CSR-версия

Метод to_sparse_csr() автоматически учитывает, сколько у матрицы разреженных измерений. Обратите внимание, что при конверсии все нулевые элементы теряются: например, у a = torch.tensor([[0,0,1],[2,3,0]]) при .to_sparse_csr() будут сохранены только ненулевые значения 1, 2, 3, а сама позиция других – отсутствует (предполагается 0).

Из sparse в dense: оазрежённый тензор всегда можно превратить обратно в обычный через метод .to_dense(). Обычно это используется для отладки или когда потеря преимуществ sparse больше, чем затраты на конвертацию.

Методы проверки: чтобы не ошибиться с размерностями, есть у разрежённого тензора свойства sparse_dim(), dense_dim(), s.is_sparse, s.is_sparse_csr и т.д. И например s.layout покажет формат (torch.sparse_coo, torch.sparse_csr и т.п.). Проверки инвариантов индексов можно включать (через аргумент check_invariants=True при создании) – это помогает отловить несогласованность форм, но по умолчанию они отключены.

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

В большинстве случаев операции с разрежёнными тензорами семантически аналогичны обычным. Но обратите внимание:

  • Умножение матриц: как мы уже видели, torch.sparse.mm(A, B) позволяет умножать разреженную матрицу A на любую другую B. Если оба исходника разрежены, ответ тоже разрежен; если второй плотный, ответ – плотный. Бэпроп работает по COO-операндам, но с CSR-матрицами градиент не считается (ограничение текущей реализации). Также PyTorch поддерживает torch.sparse.addmm(C, A, B), torch.sparse.sum(sparse, dim) и ряд других операций из модуля torch.sparse.

  • Сумма и добавление: разрежённые тензоры можно складывать: A + B для двух COO- или CSR-матриц (с одним и тем же размером). Если есть совпадения индексов, то после операции получается некоалесцированный результат – PyTorch просто конкатенирует индексы и значения. Иногда полезно потом вызвать .coalesce(), чтобы объединить дубликаты и не раздувать структуру. Добавление с плотной матрицей по дефолту приводит к плотному результату – это фича, если хочется жесткой экономии, нужно заранее перевести вторую матрицу тоже в sparse и пользоваться torch.sparse.mm.

  • Другие операции: множество базовых операций (transpose, reshape, sum, max и т.д.) над sparse-тензорами поддерживаются, но не все.

В общем важно помнить: сохраняйте тензор в разреженном формате как можно дольше и используйте специальные функции для sparse.

Заключение

Итак, torch.sparse имеет смысл там, где нулей много, а математика простая и регулярная. Форматы выбирают так: COO для построения и обновлений, CSR/CSC для линейной алгебры, BSR — когда есть устойчивые блоки. Если плотность растёт — возвращаемся к dense.

Машинное обучение выглядит красиво в презентациях, но в продакшене всё быстро упирается в вопросы: как заставить алгоритм работать на реальных данных, как не сломаться на масштабе и как объяснить результат бизнесу. Чтобы ответить на эти боли, мы готовим практические разборы сложных тем — без воды и с кодом. Участие бесплатное, записывайтесь:

Научиться работать с важнейшими моделями машинного обучения, NLP, DL, рекомендательными системами на практике с реальными данными можно на онлайн-курсе "Machine Learning. Professional".

Пройдите вступительный тест, чтобы оценить свой уровень и узнать, подойдет ли вам программа курса.

Комментарии (0)