Python PyTorch sparse_coo_tensor用法及代碼示例

本文簡要介紹python語言中 torch.sparse_coo_tensor 的用法。

用法:

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

參數：

• indices(array_like) -張量的初始數據。可以是列表、元組、NumPy ndarray、標量和其他類型。將在內部轉換為 torch.LongTensor。索引是矩陣中非零值的坐標，因此應該是二維的，其中第一維是張量維度的數量，第二維是非零值的數量。

• values(array_like) -張量的初始值。可以是列表、元組、NumPy ndarray、標量和其他類型。

• size(列表、元組或torch.Size， 可選的) -稀疏張量的大小。如果未提供，則大小將被推斷為足以容納所有非零元素的最小大小。

關鍵字參數：

• dtype(torch.dtype， 可選的) -返回張量的所需數據類型。默認值：如果無，則從 values 推斷數據類型。

• device(torch.device， 可選的) -返回張量的所需設備。默認值：如果沒有，則使用當前設備作為默認張量類型(請參閱 torch.set_default_tensor_type() )。對於 CPU 張量類型，device 將是 CPU；對於 CUDA 張量類型，device 將是當前的 CUDA 設備。

• requires_grad(bool,可選的) -如果 autograd 應該在返回的張量上記錄操作。默認值：False。

在給定的 indices 處構造一個具有指定值的 COO(rdinate) 格式的稀疏張量。

注意

此函數返回一個未合並的張量。

例子：

>>> i = torch.tensor([[0, 1, 1],
...                   [2, 0, 2]])
>>> v = torch.tensor([3, 4, 5], dtype=torch.float32)
>>> torch.sparse_coo_tensor(i, v, [2, 4])
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([3., 4., 5.]),
       size=(2, 4), nnz=3, layout=torch.sparse_coo)

>>> torch.sparse_coo_tensor(i, v)  # Shape inference
tensor(indices=tensor([[0, 1, 1],
                       [2, 0, 2]]),
       values=tensor([3., 4., 5.]),
       size=(2, 3), nnz=3, layout=torch.sparse_coo)

>>> torch.sparse_coo_tensor(i, v, [2, 4],
...                         dtype=torch.float64,
...                         device=torch.device('cuda:0'))
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)

# Create an empty sparse tensor with the following invariants:
#   1. sparse_dim + dense_dim = len(SparseTensor.shape)
#   2. SparseTensor._indices().shape = (sparse_dim, nnz)
#   3. SparseTensor._values().shape = (nnz, SparseTensor.shape[sparse_dim:])
#
# For instance, to create an empty sparse tensor with nnz = 0, dense_dim = 0 and
# sparse_dim = 1 (hence indices is a 2D tensor of shape = (1, 0))
>>> S = torch.sparse_coo_tensor(torch.empty([1, 0]), [], [1])
tensor(indices=tensor([], size=(1, 0)),
       values=tensor([], size=(0,)),
       size=(1,), nnz=0, layout=torch.sparse_coo)

# and to create an empty sparse tensor with nnz = 0, dense_dim = 1 and
# sparse_dim = 1
>>> S = torch.sparse_coo_tensor(torch.empty([1, 0]), torch.empty([0, 2]), [1, 2])
tensor(indices=tensor([], size=(1, 0)),
       values=tensor([], size=(0, 2)),
       size=(1, 2), nnz=0, layout=torch.sparse_coo)