pytorch_scatter

Dispersión de pytorch

---

Documentación

Este paquete consiste en una pequeña biblioteca de extensión de operaciones de actualización escasa altamente optimizada (dispersión y segmento) para el uso en Pytorch, que faltan en el paquete principal. Las operaciones de dispersión y segmento se pueden describir aproximadamente como operaciones reducidas basadas en un tensor de "índice de grupo" dado. Las operaciones de segmento requieren que se clasifique el tensor de "índice de grupo", mientras que las operaciones de dispersión no están sujetas a estos requisitos.

El paquete consta de las siguientes operaciones con tipos de reducción "sum"|"mean"|"min"|"max" :

• dispersión basada en índices arbitrarios
• segment_coo basado en índices ordenados
• segment_csr basado en índices comprimidos a través de punteros

Además, proporcionamos las siguientes funciones compuestas que utilizan operaciones scatter_* debajo del capó: scatter_std , scatter_logsumexp , scatter_softmax y scatter_log_softmax .

Todas las operaciones incluidas son transmitibles, funcionan en diferentes tipos de datos, se implementan tanto para CPU como para GPU con implementaciones atrasadas correspondientes y son totalmente rastreables.

Instalación

Anaconda

ACTUALIZACIÓN: ¿Ahora puede instalar pytorch-scatter a través de Anaconda para todas las principales combinaciones de OS/Pytorch/Cuda? Dado que tiene pytorch >= 1.8.0 instalado, simplemente ejecute

 conda install pytorch-scatter -c pyg

Binarios

Alternativamente, proporcionamos ruedas PIP para todas las principales combinaciones de OS/Pytorch/CUDA, ver aquí.

Pytorch 2.5

Para instalar los binarios para Pytorch 2.5.0, simplemente ejecute

 pip install torch-scatter -f https://data.pyg.org/whl/torch-2.5.0+${CUDA}.html

Donde ${CUDA} debe ser reemplazado por cpu , cu118 , cu121 o cu124 dependiendo de su instalación de Pytorch.

	cpu	cu118	cu121	cu124
Linux	✅	✅	✅	✅
Windows	✅	✅	✅	✅
macosa	✅

Pytorch 2.4

Para instalar los binarios para Pytorch 2.4.0, simplemente ejecute

 pip install torch-scatter -f https://data.pyg.org/whl/torch-2.4.0+${CUDA}.html

Donde ${CUDA} debe ser reemplazado por cpu , cu118 , cu121 o cu124 dependiendo de su instalación de Pytorch.

	cpu	cu118	cu121	cu124
Linux	✅	✅	✅	✅
Windows	✅	✅	✅	✅
macosa	✅

Nota: También se proporcionan binarios de las versiones más antiguas para 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/2.1.2, Pytorch 2.2.0/2.2.1/2.2.2 y Pytorch 2.3.0/2.3.1 (siguiente el mismo procedimiento). Para versiones anteriores, debe especificar explícitamente el último número de versión compatible o instalar a través de pip install --no-index para evitar una instalación manual desde la fuente. Puede buscar el último número de versión compatible aquí.

De la fuente

Asegúrese de que al menos Pytorch 1.4.0 esté instalado y verifique que cuda/bin y cuda/include estén en su $PATH y $CPATH respectivamente, por ejemplo :

 $ python -c "import torch; print(torch.__version__)"
>>> 1.4.0

$ echo $PATH
>>> /usr/local/cuda/bin:...

$ echo $CPATH
>>> /usr/local/cuda/include:...

Luego corre:

 pip install torch-scatter

Cuando se ejecuta en un contenedor Docker sin controlador NVIDIA, Pytorch necesita evaluar las capacidades de cálculo y puede fallar. En este caso, asegúrese de que las capacidades de cálculo se establezcan a través de TORCH_CUDA_ARCH_LIST , EG :

 export TORCH_CUDA_ARCH_LIST = "6.0 6.1 7.2+PTX 7.5+PTX"

Ejemplo

 import torch
from torch_scatter import scatter_max

src = torch . tensor ([[ 2 , 0 , 1 , 4 , 3 ], [ 0 , 2 , 1 , 3 , 4 ]])
index = torch . tensor ([[ 4 , 5 , 4 , 2 , 3 ], [ 0 , 0 , 2 , 2 , 1 ]])

out , argmax = scatter_max ( src , index , dim = - 1 )

 print(out)
tensor([[0, 0, 4, 3, 2, 0],
        [2, 4, 3, 0, 0, 0]])

print(argmax)
tensor([[5, 5, 3, 4, 0, 1]
        [1, 4, 3, 5, 5, 5]])

Pruebas de ejecución

 pytest

API C ++

torch-scatter también ofrece una API C ++ que contiene C ++ equivalente de modelos Python. Para esto, necesitamos agregar TorchLib al -DCMAKE_PREFIX_PATH ( por ejemplo , puede existe en {CONDA}/lib/python{XX}/site-packages/torch si está instalado a través de conda )::

 mkdir build
cd build
# Add -DWITH_CUDA=on support for CUDA support
cmake -DCMAKE_PREFIX_PATH="..." ..
make
make install