原文链接
常用的稀疏矩阵存储格式有COO,CSR/CSC,LIL
COO(Coordinate format )是最为简单的格式,以三元组的形式存储稀疏矩阵。记录矩阵中非零元素的数值和所在的行序号和列序号。形式为(行号,列号,数值)。这种存储方式的主要优点是灵活、简单。但是缺点是不可以直接进行矩阵的相关运算
CSR(Compressed Sparse Row)格式实现了用于存储二维张量的 CSR 格式。尽管不支持 N 维张量,但与 COO 格式相比的主要优势是更好地利用存储和更快的计算操作。目前尚不存在 CUDA 支持
LIL (List-of-List) 每行存储一个列表,每个条目包含列索引和值。通常,这些条目按列索引进行排序,以便更快地查找
Pytorch中,处理稀疏矩阵的有效工具torch.sparse。Torch 支持 COO(rdinate) 格式的稀疏张量,可以有效地存储和处理大多数元素为零的张量
torch.spares_coo_tensor(indices, values, siez=None,*, dtype=None, requires_grad=False)->Tensor
参数:
indices
: 一个2D张量,其中每一列都代表一个非零元素的坐标。values
: 一个1D张量,其中每个值都是与indices
对应坐标中的非零元素对应的值。size
: (可选) 一个表示稀疏张量大小的元组假设一个2D的tensor数据
0 0 0 5
0 0 3 0
0 2 0 0
非零元素的坐标(indices)和对应的值(values)为:
indices = [[0, 3],
[1, 2],
[2, 1]]
values = [5, 3, 2]
可以使用torch.sparse_coo_tensor
来创建这个稀疏张量:
import torch
indices = torch.tensor([[0, 1, 2],
[3, 2, 1]])
values = torch.tensor([5, 3, 2])
size = (3, 4)
sparse_tensor = torch.sparse_coo_tensor(indices, values, size)
import torch
# 示例张量
tensor = torch.tensor([[0, 2, 1], [0, 0, 3], [4, 0, 0]])
# 寻找非零元素的索引
non_zero_indices = torch.nonzero(tensor).t()
print(non_zero_indices[0]) # tensor([0, 0, 1, 2])
print(non_zero_indices[1]) # tensor([1, 2, 2, 0])
print(tensor.dim()) # 2
# 获取非零元素的值
values = tensor[tuple(non_zero_indices[i] for i in range(tensor.dim()))]
print(values) # tensor([2, 1, 3, 4])
print(tensor.size()) # torch.Size([3, 3])
# 创建稀疏张量
sparse_tensor = torch.sparse_coo_tensor(non_zero_indices, values, tensor.size())
print(sparse_tensor)
最后得到的稀疏矩阵
tensor(indices=tensor([[0, 0, 1, 2],
[1, 2, 2, 0]]),
values=tensor([2, 1, 3, 4]),
size=(3, 3), nnz=4, layout=torch.sparse_coo)
仍然可以使用torch.sparse_coo_tensor
函数将其转换为稀疏表示。与2D张量相似,你需要确定非零元素的位置和它们的值。对于3D张量,每个非零元素的坐标将由三个值表示
import torch
# 示例3D张量
tensor = torch.tensor([
[[0, 2, 0], [0, 0, 3], [4, 0, 0]],
[[0, 0, 0], [0, 5, 0], [0, 0, 6]]
])
print(tensor.dim()) # 3
# 寻找非零元素的索引
non_zero_indices = torch.nonzero(tensor).t()
print(non_zero_indices[0]) # tensor([0, 0, 0, 1, 1])
print(non_zero_indices[1]) # tensor([0, 1, 2, 1, 2])
print(non_zero_indices[2]) # tensor([1, 2, 0, 1, 2])
# 获取非零元素的值
values = tensor[tuple(non_zero_indices[i] for i in range(tensor.dim()))]
# 创建稀疏张量
sparse_tensor = torch.sparse_coo_tensor(non_zero_indices, values, tensor.size())
print(sparse_tensor)
"""
将tensor 数据转换为COO 稀疏张量函数
"""
def tensor_to_sparse(dense_tensor):
size = dense_tensor.size()
# 寻找非零元素的索引
non_zero_indices = torch.nonzero(dense_tensor).t()
# 获取非零元素的值
values = dense_tensor[tuple(non_zero_indices[i] for i in range(dense_tensor.dim()))]
# 创建稀疏张量
sparse_tensor = torch.sparse_coo_tensor(non_zero_indices, values, size)
return sparse_tensor,size
# 随机产生一个4D张量数据
dense_tensor = torch.randn((2,3,3,3))
print(dense_tensor)
print(dense_tensor.dim()) # 4
sparse_tensor,size = tensor_to_sparse(dense_tensor)
print(sparse_tensor)
print(size) # torch.Size([2, 3, 3, 3])
# 保存稀疏张量到硬盘
torch.save(sparse_tensor, "spare_tensor.npz")
# 从硬盘上加载稀疏张量
loaded_sparse_tensor = torch.load("spare_tensor.npz")
"""
COO 稀疏张量转换为密集张量
"""
def sparse_to_tensor(loaded_sparse_tensor):
# 将稀疏张量复原为原始的密集张量
dense_tensor = loaded_sparse_tensor.to_dense()
return dense_tensor
# 调用函数
sparse_to_tensor(loaded_sparse_tensor)
import scipy.sparse
import torch
import scipy.sparse
def tensor_to_sparse(dense_tensor):
# 将dense_tensor转化为2D
shape = dense_tensor.shape
tensor_2d = dense_tensor.view(-1, shape[-1])
# 将2D tensor转化为numpy array
array_2d = tensor_2d.numpy()
# 从numpy array创建sparse matrix
sparse_matrix = scipy.sparse.coo_matrix(array_2d)
return sparse_matrix,shape
def sparse_to_tensor(sparse_matrix, original_shape):
# 从稀疏矩阵转换为2D array
array_2d = sparse_matrix.toarray()
# 将2D array转换为original_shape_array
original_shape_array = array_2d.reshape(original_shape)
# 将3D array转换为3D tensor
dense_tensor = torch.from_numpy(original_shape_array)
return dense_tensor
# 随机产生一个dense_tensor
dense_tensor = torch.randn((2,3,3,3))
# 转化为sparse matrix
sparse_matrix,original_shape = tensor_to_sparse(dense_tensor)
print(original_shape)
# 将sparse matrix保存到硬盘上
scipy.sparse.save_npz('sparse_matrix.npz', sparse_matrix)
# 使用scipy.sparse.load_npz从硬盘加载保存的稀疏张量
loaded_sparse_matrix = scipy.sparse.load_npz('sparse_matrix.npz')
# 稀疏张量复原为原来的tensor数据
restored_tensor = sparse_to_tensor(loaded_sparse_matrix, original_shape)
print(restored_tensor)
# 判断restored_tensor与原来的tensor数据是否一致
print(dense_tensor==restored_tensor)