PyTorch sparse

torch.sparse은 PyTorch에서 희소 텐서(sparse tensor)를 다루기 위한 기능 집합이다.

개요

희소 텐서는 대부분의 원소가 0인 고차원 배열을 효율적으로 표현하기 위한 자료구조이다. PyTorch에서는 여러 희소 표현(layout) 방식을 지원하며, torch.sparse 모듈 및 관련 함수들이 이 기능을 제공한다.

표현 방식 및 레이아웃

PyTorch는 여러 희소 레이아웃을 지원한다. 주요 레이아웃과 특징은 다음과 같다.

COO (좌표형, Coordinate)

• torch.sparse_coo_tensor(indices, values, size) 함수를 통해 생성할 수 있다.
• indices는 크기 (ndim × nnz)의 정수 텐서로, 각 열이 하나의 비제로(nonzero) 원소의 좌표를 나타낸다.
• values는 크기 (nnz, *dense_dims)의 텐서로 비제로 원소 값을 저장한다.
• 희소 텐서는 일반적으로 비정합(uncoalesced) 상태일 수 있으며, .coalesce() 메서드를 통해 동일한 좌표들이 합쳐진 상태로 만들 수 있다.
• 텐서의 sparse_dim()과 dense_dim() 메서드로 희소 차원과 밀집 차원을 구분할 수 있다.

CSR / CSC (압축 행 / 압축 열)

• torch.sparse_csr_tensor(crow_indices, col_indices, values, size)
• torch.sparse_csc_tensor(ccol_indices, row_indices, values, size)
• CSR 방식은 행 인덱스를 압축 형태(crow_indices)로 저장하고, 각 행 내의 열 인덱스(col_indices)와 값 배열을 저장하는 방식이다.
• CSC는 열 중심의 압축 표현 방식이다.
• CSR/CSC 방식은 특히 행 또는 열 기반 연산 (예: 행 단위 순회)에 효율적이다.

BSR / BSC (블록 압축형)

• torch.sparse_bsr_tensor(crow_indices, col_indices, values, size)
• torch.sparse_bsc_tensor(ccol_indices, row_indices, values, size)
• 이 방식은 텐서를 일정 블록(block) 단위로 나누어 블록 단위의 압축 표현을 허용한다.
• 대형 희소 행렬을 블록 단위로 처리할 때 유리하다.

일반 압축 희소 텐서 생성 함수

• torch.sparse_compressed_tensor(...) 함수를 통해 layout 인자를 지정하는 방식으로 CSR, CSC, BSR, BSC 등을 생성할 수 있다.

주요 메서드 및 연산

• .to_dense() : 희소 텐서를 밀집(dense) 텐서로 변환
• .indices(), .values() : COO 방식의 희소 텐서에서 좌표 인덱스와 값 반환
• .coalesce() : 중복 좌표의 항목을 병합
• torch.sparse.mm(mat1, mat2) : 희소 행렬과 밀집 또는 희소 행렬 간의 곱셈 연산. COO 또는 CSR 포맷 지원
• torch.sparse.addmm(mat, mat1, mat2, *, beta, alpha) : mat + alpha * (mat1 @ mat2)를 계산하며, mat1이 희소인 경우 backward를 지원
• is_sparse, is_sparse_csr 등 속성 메서드로 현재 텐서가 어떠한 희소 레이아웃을 사용하는지 판단 가능

제약 및 주의점

• 모든 연산이 모든 희소 레이아웃에 대해 지원되는 것은 아니다. 예를 들어, CSR 방식에 대한 미분(gradient) 지원은 제한적일 수 있다.
• torch.sparse_csr_tensor()로 생성된 텐서는 역전파 시 gradient가 차단(stopped)되는 문제가 보고된 바 있다.
• 희소 방식은 메모리 절약과 연산 효율을 제공하지만, 빈번한 인덱스 접근, 슬라이싱, 변경 등의 작업에는 오버헤드가 클 수 있다.
• 일부 고성능 희소 연산을 위해 외부 라이브러리 (예: torch_sparse 패키지)가 사용된다. 이 패키지는 sparse-dense 곱, sparse-sparse 곱, 전치(transpose) 등의 최적화된 연산을 제공한다.

활용 예시

import torch

# COO 희소 텐서 생성 예
i = torch.tensor([[0, 1, 1], [2, 0, 2]])  # 좌표 인덱스 (2×3)
v = torch.tensor([3, 4, 5], dtype=torch.float32)  # 값
s = torch.sparse_coo_tensor(i, v, [2, 3])
dense = s.to_dense()  # 밀집 텐서로 변환

# CSR 희소 텐서 생성 예
crow = torch.tensor([0, 2, 3])
cols = torch.tensor([0, 2, 1])
vals = torch.tensor([3, 4, 5], dtype=torch.float32)
s_csr = torch.sparse_csr_tensor(crow, cols, vals, size=(2, 3))
out = s_csr.matmul(torch.randn(3, 4))  # 희소 × 밀집 연산

관련 기술 및 라이브러리

• torch_sparse 패키지는 PyTorch의 희소 텐서 연산을 확장한 외부 라이브러리로, 자동 미분(sparse autograd) 및 더 많은 커널을 제공한다.
• PyTorch 내부의 희소 텐서 지원은 계속 개발 중이며, 차후 더 많은 레이아웃과 연산이 확대될 가능성이 있다.

같이 보기

• PyTorch
• 희소 행렬
• 희소 기법 (딥러닝)
• torch_sparse
• Autograd

참고 문헌

• PyTorch 공식 문서, “torch.sparse”