Всё что нужно знать про torch.sparse. data science.. data science. ml.. data science. ml. PyTorch.. data science. ml. PyTorch. torch.sparse.. data science. ml. PyTorch. torch.sparse. матричное умножение.. data science. ml. PyTorch. torch.sparse. матричное умножение. оптимизация памяти.. data science. ml. PyTorch. torch.sparse. матричное умножение. оптимизация памяти. разрежённые тензоры.

Разработчики 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”.

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

Автор: badcasedaily1

Источник

Версия для печати Версия для печати
Rambler's Top100