torch.sparse_coo_tensor

以**COO格式**构造一个稀疏张量，给定索引一个对应值。其他地方的值为0。

• COO格式（coordinate格式）：稀疏矩阵被表示为三个数组：
  • row：行索引数组，表示非零元素所在的行。
    col：列索引数组，表示非零元素所在的列。
    data：值数组，包含非零元素的值。
  • 假设我们有一个5x5的矩阵，其中只有四个非零元素，位置和值如下：
    元素1位于第1行第1列，值为3。（Python中索引从0开始，所以第1行实际上是索引0）
    元素2位于第2行第3列，值为4。
    元素3位于第4行第2列，值为5。
    元素4位于第5行第5列，值为6。
    COO格式中，这个矩阵可以表示为：
    row = [0, 1, 3, 4]
    col = [0, 2, 1, 4]
    data = [3, 4, 5, 6]

torch.sparse_coo_tensor(indices, values, size=None, *, dtype=None, device=None, pin_memory=False, requires_grad=False, check_invariants=None, is_coalesced=None) → Tensor

参数

• indices （array_like） – 张量的初始数据。可以是列表、元组、NumPy ndarray、标量和其他类型。但在内部会被转换为torch.LongTensor(PyTorch中的长整型张量)。indices 是矩阵中非零值的坐标，因此 应该是二维的，其中第一个维度是张量维度的数量，并且 第二个维度是非零值的数量。
• values （array_like） – 张量的初始值。可以是列表、元组、NumPy ndarray、标量和其他类型。
• size （列表、元组或 torch.Size，可选） – 稀疏张量的大小。如果未提供，则大小将被推断为足够大的最小大小，以容纳所有非零元素。

关键字参数

• dtype （torch.dtype，可选） – 返回张量的所需数据类型。默认值：如果为 None，则从值推断数据类型。
• device （torch.device，可选） – 返回的张量的所需设备。Default：如果为 None，则使用当前设备作为默认张量类型（请参阅 torch.set_default_device（） ）。device 将是 CPU 张量类型的 CPU，CUDA 张量类型的当前 CUDA 设备。
• pin_memory （bool，可选） – 如果设置，则返回的张量将在固定内存中分配。仅适用于 CPU 张量。默认值：False。可以使用 .to(‘cuda’) 方法将数据从固定内存移动到 GPU，通常比从常规内存移动到 GPU 更快。
• requires_grad （bool， 可选） – 如果 autograd 应记录对返回的张量的操作。默认值：False。
• check_invariants （bool， 可选） – 是否选中稀疏张量不变性。默认值：False（由 torch.sparse.check_sparse_tensor_invariants.is_enabled() 得到）
  • 不变性检查是一种验证机制，用于确保稀疏张量的数据结构和逻辑是正确的。如：压缩稀疏行（CSR）格式中的一个不变性条件，即累积索引数组的最后一个元素必须等于非零元素的数量（nnz）。
• is_coalesced （bool， 可选） – 默认值：None：除了琐碎的情况（例如 nnz(非零元素的数量) < 2）外，生成的 Tensor is_coalesced设置为False。当 ‘‘True’’ 时，调用方负责提供与合并张量相对应的张量索引。如果 check_invariants 标志为 False，则如果不满足先决条件，则不会引发错误，这将导致无提示错误的结果。要强制合并，请在结果的 Tensor 上使用 coalesce（） 。
  • 未合并的稀疏COO张量：假设我们有一个未合并的1-D稀疏COO张量，其索引和值如下所示：
    索引: [0, 1, 1, 2, 2, 2]
    值: [1, 2, 3, 4, 5, 6]
    在这个例子中，索引1和2都有重复的值。当我们将这个未合并的稀疏张量合并时，索引1处的值2和3会被加起来，得到5；索引2处的值4、5和6会被加起来，得到15。
  • 合并后的稀疏张量的索引和值如下
    索引: [0, 1, 2]
    值: [1, 5, 15]
    这样，合并后的稀疏张量就具有唯一的索引，并且每个索引处的值是所有重复值条目的总和。

例程

import torch

# 定义indices和values
# indices是二维的，其中第一维是张量维度的数量，第二维是非零值的数量
# values是一维的，包含了非零值
indices = torch.tensor([[0, 1, 1],
                        [2, 0, 2]])
values = torch.tensor([3, 4, 5], dtype=torch.float32)

# 创建一个稀疏张量，指定其大小
# size参数是可选的，如果不提供，将根据indices和values推断大小
sparse_tensor = torch.sparse_coo_tensor(indices, values, size=[2, 4])

# 打印稀疏张量
print("指定大小的稀疏张量:")
print(sparse_tensor)

# 如果不指定size，将根据indices和values推断大小
# 这里我们没有指定size，所以大小将被推断为(2, 3)
inferred_sparse_tensor = torch.sparse_coo_tensor(indices, values)

# 打印推断大小的稀疏张量
print("\n根据indices和values推断大小的稀疏张量:")
print(inferred_sparse_tensor)

# 创建一个稀疏张量，并指定数据类型和设备
# dtype参数是可选的，如果不提供，将从values推断数据类型
# device参数是可选的，如果不提供，将使用当前设备的默认张量类型
specific_dtype_device_sparse_tensor = torch.sparse_coo_tensor(
    indices, values, [2, 4],
    dtype=torch.float64,  # 指定数据类型为double
    device=torch.device('cuda:0')  # 指定设备为CUDA的第一个设备
)

# 打印指定数据类型和设备的稀疏张量
print("\n指定数据类型和设备的稀疏张量:")
print(specific_dtype_device_sparse_tensor)

# 创建一个空的稀疏张量，不包含任何非零元素
# 这里我们使用空的indices和values来创建一个空的稀疏张量
empty_indices = torch.empty([1, 0])  # 空的indices
empty_values = []  # 空的values

empty_sparse_tensor = torch.sparse_coo_tensor(empty_indices, empty_values, [1])

# 打印空的稀疏张量
print("\n空的稀疏张量:")
print(empty_sparse_tensor)

# 创建一个包含一个非零元素的稀疏张量
# 这里我们指定了非零元素的位置和值
indices_with_value = torch.tensor([[0],  # 包含一个非零元素的indices
                                   [1]])  # 注意这里的形状是(1, 2)
values_with_value = torch.tensor([10])  # 非零元素的值

sparse_tensor_with_value = torch.sparse_coo_tensor(
    indices_with_value, values_with_value, [2, 1]
)

# 打印包含一个非零元素的稀疏张量
print("\n包含一个非零元素的稀疏张量:")
print(sparse_tensor_with_value)
'''output
指定大小的稀疏张量:
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([3., 4., 5.]),
       size=(2, 4), nnz=3, layout=torch.sparse_coo)

根据indices和values推断大小的稀疏张量:
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([3., 4., 5.]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo)

指定数据类型和设备的稀疏张量:
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([3., 4., 5.]),
       device='cuda:0', size=(2, 4), nnz=3, dtype=torch.float64,
       layout=torch.sparse_coo)

空的稀疏张量:
tensor(indices=tensor([], size=(1, 0)),
       values=tensor([], size=(0,)),
       size=(1,), nnz=0, layout=torch.sparse_coo)

包含一个非零元素的稀疏张量:
tensor(indices=tensor([[0],
                       [1]]),
       values=tensor([10]),
       size=(2, 1), nnz=1, layout=torch.sparse_coo)
'''