該軟件包由一個小型擴展庫組成,其中包括具有自動射擊支持的優化稀疏矩陣操作。該軟件包當前包含以下方法:
所有包括在不同數據類型上的操作工作,並針對CPU和GPU實施。為了避免創建torch.sparse_coo_tensor ,此軟件包通過簡單地將index和value張量作為參數(具有與Pytorch中定義的相同形狀)來定義稀疏張量的操作。請注意,僅具有自動踢的value ,因為index是離散的,因此沒有可區分。
更新:您現在可以通過Anaconda安裝pytorch-sparse以用於所有主要的OS/Pytorch/CUDA組合?鑑於您安裝了pytorch >= 1.8.0 ,只需運行
conda install pytorch-sparse -c pyg
或者,我們為所有主要的OS/Pytorch/CUDA組合提供PIP輪轂,請參見此處。
要安裝Pytorch 2.5.0的二進製文件,只需運行
pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-2.5.0+${CUDA}.html
${CUDA}應取代cpu , cu118 , cu121或cu124 ,具體取決於您的pytorch安裝。
cpu | cu118 | cu121 | cu124 | |
|---|---|---|---|---|
| Linux | ✅ | ✅ | ✅ | ✅ |
| 視窗 | ✅ | ✅ | ✅ | ✅ |
| macos | ✅ |
要安裝Pytorch 2.4.0的二進製文件,只需運行
pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-2.4.0+${CUDA}.html
${CUDA}應取代cpu , cu118 , cu121或cu124 ,具體取決於您的pytorch安裝。
cpu | cu118 | cu121 | cu124 | |
|---|---|---|---|---|
| Linux | ✅ | ✅ | ✅ | ✅ |
| 視窗 | ✅ | ✅ | ✅ | ✅ |
| macos | ✅ |
Note: Binaries of older versions are also provided for PyTorch 1.4.0, PyTorch 1.5.0, PyTorch 1.6.0, PyTorch 1.7.0/1.7.1, PyTorch 1.8.0/1.8.1, PyTorch 1.9.0, PyTorch 1.10.0/1.10.1/1.10.2, PyTorch 1.11.0, PyTorch 1.12.0/1.12.1,Pytorch 1.13.0/1.13.1,Pytorch 2.0.0/2.0.1,Pytorch 2.1.0/2.1.1.1/2.1.2,Pytorch 2.2.0/2.2.1.2.1.2.1/2.2.2,以及Pytorch 2.3.0/2.3.0/2.3.0/2.3.1/2.3.1(遵循同一過程)。對於較舊的版本,您需要明確指定最新的支持版本號,或通過pip install --no-index安裝,以防止源頭安裝手動安裝。您可以在此處查找最新支持的版本號。
確保安裝了至少Pytorch 1.7.0,並驗證cuda/bin and cuda/include分別在您的$PATH和$CPATH中,例如:
$ python -c "import torch; print(torch.__version__)"
>>> 1.7.0
$ echo $PATH
>>> /usr/local/cuda/bin:...
$ echo $CPATH
>>> /usr/local/cuda/include:...
如果您想額外構建具有METIS支持的torch-sparse ,例如用於分區,請按照Install.txt文件中的說明下載並安裝METIS庫。請注意,METIS需要通過更改include/metis.h安裝64個位IDXTYPEWIDTH 。之後, WITH_METIS=1設置環境變量。
然後運行:
pip install torch-scatter torch-sparse
當在沒有NVIDIA驅動程序的Docker容器中運行時,Pytorch需要評估計算功能,並且可能會失敗。在這種情況下,請確保通過TORCH_CUDA_ARCH_LIST設置計算功能,例如:
export TORCH_CUDA_ARCH_LIST="6.0 6.1 7.2+PTX 7.5+PTX"
torch_sparse.coalesce(index, value, m, n, op="add") -> (torch.LongTensor, torch.Tensor)
劃分的index並刪除重複的條目。重複的條目通過將它們散佈在一起來刪除。為了散射,可以使用torch_scatter的任何操作。
"add" ) import torch
from torch_sparse import coalesce
index = torch . tensor ([[ 1 , 0 , 1 , 0 , 2 , 1 ],
[ 0 , 1 , 1 , 1 , 0 , 0 ]])
value = torch . Tensor ([[ 1 , 2 ], [ 2 , 3 ], [ 3 , 4 ], [ 4 , 5 ], [ 5 , 6 ], [ 6 , 7 ]])
index , value = coalesce ( index , value , m = 3 , n = 2 ) print(index)
tensor([[0, 1, 1, 2],
[1, 0, 1, 0]])
print(value)
tensor([[6.0, 8.0],
[7.0, 9.0],
[3.0, 4.0],
[5.0, 6.0]])
torch_sparse.transpose(index, value, m, n) -> (torch.LongTensor, torch.Tensor)
稀疏基質的尺寸為0和1。
False ,則不會融合產出。 (默認: True ) import torch
from torch_sparse import transpose
index = torch . tensor ([[ 1 , 0 , 1 , 0 , 2 , 1 ],
[ 0 , 1 , 1 , 1 , 0 , 0 ]])
value = torch . Tensor ([[ 1 , 2 ], [ 2 , 3 ], [ 3 , 4 ], [ 4 , 5 ], [ 5 , 6 ], [ 6 , 7 ]])
index , value = transpose ( index , value , 3 , 2 ) print(index)
tensor([[0, 0, 1, 1],
[1, 2, 0, 1]])
print(value)
tensor([[7.0, 9.0],
[5.0, 6.0],
[6.0, 8.0],
[3.0, 4.0]])
torch_sparse.spmm(index, value, m, n, matrix) -> torch.Tensor
具有緻密基質的稀疏基質的基質產物。
import torch
from torch_sparse import spmm
index = torch . tensor ([[ 0 , 0 , 1 , 2 , 2 ],
[ 0 , 2 , 1 , 0 , 1 ]])
value = torch . Tensor ([ 1 , 2 , 4 , 1 , 3 ])
matrix = torch . Tensor ([[ 1 , 4 ], [ 2 , 5 ], [ 3 , 6 ]])
out = spmm ( index , value , 3 , 3 , matrix ) print(out)
tensor([[7.0, 16.0],
[8.0, 20.0],
[7.0, 19.0]])
torch_sparse.spspmm(indexA, valueA, indexB, valueB, m, k, n) -> (torch.LongTensor, torch.Tensor)
兩個稀疏張量的基質產物。兩種輸入稀疏矩陣都需要合併(使用coalesced屬性來強制)。
True ,則將合併兩個輸入稀疏矩陣。 (默認: False ) import torch
from torch_sparse import spspmm
indexA = torch . tensor ([[ 0 , 0 , 1 , 2 , 2 ], [ 1 , 2 , 0 , 0 , 1 ]])
valueA = torch . Tensor ([ 1 , 2 , 3 , 4 , 5 ])
indexB = torch . tensor ([[ 0 , 2 ], [ 1 , 0 ]])
valueB = torch . Tensor ([ 2 , 4 ])
indexC , valueC = spspmm ( indexA , valueA , indexB , valueB , 3 , 3 , 2 ) print(indexC)
tensor([[0, 1, 2],
[0, 1, 1]])
print(valueC)
tensor([8.0, 6.0, 8.0])
pytest
torch-sparse還提供了一個C ++ API,其中包含C ++等效的Python模型。為此,我們需要將TorchLib添加到-DCMAKE_PREFIX_PATH (例如,它可能存在於{CONDA}/lib/python{XX}/site-packages/torch中,如果通過conda安裝):
mkdir build
cd build
# Add -DWITH_CUDA=on support for CUDA support
cmake -DCMAKE_PREFIX_PATH="..." ..
make
make install