Skip to content

GitLab

Commit 1cb25232 authored by limm's avatar limm
Browse files

push 0.6.15 version

parent e8309f27
Hide whitespace changes
Inline Side-by-side
Showing with 1011 additions and 126 deletions
CMakeLists.txt 0 → 100644
View file @ 1cb25232
cmake_minimum_required(VERSION 3.0)
project(torchsparse)
set(CMAKE_CXX_STANDARD 14)
set(TORCHSPARSE_VERSION 0.6.15)
option(WITH_CUDA "Enable CUDA support" OFF)
option(WITH_PYTHON "Link to Python when building" ON)
if(WITH_CUDA)
enable_language(CUDA)
add_definitions(-D__CUDA_NO_HALF_OPERATORS__)
add_definitions(-DWITH_CUDA)
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-relaxed-constexpr")
endif()
if (WITH_PYTHON)
add_definitions(-DWITH_PYTHON)
find_package(Python3 COMPONENTS Development)
endif()
find_package(Torch REQUIRED)
file(GLOB HEADERS csrc/*.h)
file(GLOB OPERATOR_SOURCES csrc/*.* csrc/cpu/*.*)
if(WITH_CUDA)
file(GLOB OPERATOR_SOURCES ${OPERATOR_SOURCES} csrc/cuda/*.h csrc/cuda/*.cu)
endif()
add_library(${PROJECT_NAME} SHARED ${OPERATOR_SOURCES})
target_link_libraries(${PROJECT_NAME} PRIVATE ${TORCH_LIBRARIES})
if (WITH_PYTHON)
target_link_libraries(${PROJECT_NAME} PRIVATE Python3::Python)
endif()
set_target_properties(${PROJECT_NAME} PROPERTIES EXPORT_NAME TorchSparse)
target_include_directories(${PROJECT_NAME} INTERFACE
$<BUILD_INTERFACE:${HEADERS}>
$<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>)
include(GNUInstallDirs)
include(CMakePackageConfigHelpers)
set(PHMAP_DIR third_party/parallel-hashmap)
target_include_directories(${PROJECT_NAME} PRIVATE ${PHMAP_DIR})
set(TORCHSPARSE_CMAKECONFIG_INSTALL_DIR "share/cmake/TorchSparse" CACHE STRING "install path for TorchSparseConfig.cmake")
configure_package_config_file(cmake/TorchSparseConfig.cmake.in
"${CMAKE_CURRENT_BINARY_DIR}/TorchSparseConfig.cmake"
INSTALL_DESTINATION ${TORCHSPARSE_CMAKECONFIG_INSTALL_DIR})
write_basic_package_version_file(${CMAKE_CURRENT_BINARY_DIR}/TorchSparseConfigVersion.cmake
VERSION ${TORCHSPARSE_VERSION}
COMPATIBILITY AnyNewerVersion)
install(FILES ${CMAKE_CURRENT_BINARY_DIR}/TorchSparseConfig.cmake
${CMAKE_CURRENT_BINARY_DIR}/TorchSparseConfigVersion.cmake
DESTINATION ${TORCHSPARSE_CMAKECONFIG_INSTALL_DIR})
install(TARGETS ${PROJECT_NAME}
EXPORT TorchSparseTargets
LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
)
install(EXPORT TorchSparseTargets
NAMESPACE TorchSparse::
DESTINATION ${TORCHSPARSE_CMAKECONFIG_INSTALL_DIR})
install(FILES ${HEADERS} DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME})
install(FILES
csrc/cpu/convert_cpu.h
csrc/cpu/diag_cpu.h
csrc/cpu/metis_cpu.h
csrc/cpu/rw_cpu.h
csrc/cpu/saint_cpu.h
csrc/cpu/sample_cpu.h
csrc/cpu/spmm_cpu.h
csrc/cpu/spspmm_cpu.h
DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}/cpu)
if(WITH_CUDA)
install(FILES
csrc/cuda/convert_cuda.h
csrc/cuda/diag_cuda.h
csrc/cuda/rw_cuda.h
csrc/cuda/spmm_cuda.h
csrc/cuda/spspmm_cuda.h
DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/${PROJECT_NAME}/cuda)
endif()
if(WITH_CUDA)
set_property(TARGET torch_cuda PROPERTY INTERFACE_COMPILE_OPTIONS "")
set_property(TARGET torch_cpu PROPERTY INTERFACE_COMPILE_OPTIONS "")
endif()
README.md
View file @ 1cb25232
# <div align="center"><strong>torch-sparse-0.6.13</strong></div> [pypi-image]: https://badge.fury.io/py/torch-sparse.svg
## 简介 [pypi-url]: https://pypi.python.org/pypi/torch-sparse
torch-sparse是PyTorch的一个扩展库，用于处理稀疏数据。它提供了一些功能强大的稀疏矩阵操作，以及能够高效执行稀疏计算的函数和工具。 [testing-image]: https://github.com/rusty1s/pytorch_sparse/actions/workflows/testing.yml/badge.svg
[testing-url]: https://github.com/rusty1s/pytorch_sparse/actions/workflows/testing.yml
[linting-image]: https://github.com/rusty1s/pytorch_sparse/actions/workflows/linting.yml/badge.svg
[linting-url]: https://github.com/rusty1s/pytorch_sparse/actions/workflows/linting.yml
[coverage-image]: https://codecov.io/gh/rusty1s/pytorch_sparse/branch/master/graph/badge.svg
[coverage-url]: https://codecov.io/github/rusty1s/pytorch_sparse?branch=master
## 依赖安装 # PyTorch Sparse
+ pytorch1.10或者pytorch1.13 以及对应的torchvision（建议dtk-22.04.2、dtk-23.04与dtk-23.10）
+ python 3.7-3.10
### 1、使用源码编译方式安装 [![PyPI Version][pypi-image]][pypi-url]
[![Testing Status][testing-image]][testing-url]
[![Linting Status][linting-image]][linting-url]
[![Code Coverage][coverage-image]][coverage-url]
#### 编译环境准备 --------------------------------------------------------------------------------
提供2种环境准备方式：
1. 基于光源pytorch基础镜像环境：镜像下载地址：[https://sourcefind.cn/#/image/dcu/pytorch](https://sourcefind.cn/#/image/dcu/pytorch)，根据pytorch、python、dtk及系统下载对应的镜像版本。 This package consists of a small extension library of optimized sparse matrix operations with autograd support.
This package currently consists of the following methods:
2. 基于现有python环境：安装pytorch和torchvision，whl包下载目录：[https://cancon.hpccube.com:65024/4/main/pytorch](https://cancon.hpccube.com:65024/4/main/pytorch)，[https://cancon.hpccube.com:65024/4/main/vision](https://cancon.hpccube.com:65024/4/main/vision)，根据python、dtk版本,下载对应pytorch和torchvision的whl包。安装命令如下： * **[Coalesce](#coalesce)**
```shell * **[Transpose](#transpose)**
pip install torch* (下载的torch的whl包) * **[Sparse Dense Matrix Multiplication](#sparse-dense-matrix-multiplication)**
pip install torchvision* (下载的torchvision的whl包) * **[Sparse Sparse Matrix Multiplication](#sparse-sparse-matrix-multiplication)**
pip install setuptools==59.5.0 wheel
All included operations work on varying data types and are implemented both for CPU and GPU.
To avoid the hazzle of creating [`torch.sparse_coo_tensor`](https://pytorch.org/docs/stable/torch.html?highlight=sparse_coo_tensor#torch.sparse_coo_tensor), this package defines operations on sparse tensors by simply passing `index` and `value` tensors as arguments ([with same shapes as defined in PyTorch](https://pytorch.org/docs/stable/sparse.html)).
Note that only `value` comes with autograd support, as `index` is discrete and therefore not differentiable.
## Installation
### Anaconda
**Update:** You can now install `pytorch-sparse` via [Anaconda](https://anaconda.org/pyg/pytorch-sparse) for all major OS/PyTorch/CUDA combinations 🤗
Given that you have [`pytorch >= 1.8.0` installed](https://pytorch.org/get-started/locally/), simply run
```
conda install pytorch-sparse -c pyg
``` ```
#### Note: Conda packages are not published for PyTorch 1.12 yet
### Binaries
We alternatively provide pip wheels for all major OS/PyTorch/CUDA combinations, see [here](https://data.pyg.org/whl).
#### PyTorch 1.12
To install the binaries for PyTorch 1.12.0, simply run
#### 源码编译安装
- 代码下载
```shell
git clone http://developer.hpccube.com/codes/aicomponent/torch-sparce # 根据编译需要切换分支
``` ```
- 源码编译（进入torch-sparse目录）： pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-1.12.0+${CUDA}.html
``` ```
export C_INCLUDE_PATH=/public/software/apps/DeepLearning/PyTorch_Lib/gflags-2.1.2-build/include:$C_INCLUDE_PATH
export CPLUS_INCLUDE_PATH=/public/software/apps/DeepLearning/PyTorch_Lib/gflags-2.1.2-build/include:$CPLUS_INCLUDE_PATH
export C_INCLUDE_PATH=/public/software/apps/DeepLearning/PyTorch_Lib/glog-build/include:$C_INCLUDE_PATH
export CPLUS_INCLUDE_PATH=/public/software/apps/DeepLearning/PyTorch_Lib/glog-build/include:$CPLUS_INCLUDE_PATH
export C_INCLUDE_PATH=$ROCM_PATH/rocrand/include:$C_INCLUDE_PATH
export CPLUS_INCLUDE_PATH=$ROCM_PATH/rocrand/include:$CPLUS_INCLUDE_PATH
export LD_LIBRARY_PATH=$ROCM_PATH/rocrand/lib:$LD_LIBRARY_PATH
export FORCE_ONLY_HIP=1
export CC=hipcc
export CXX=hipcc
python setup.py install where `${CUDA}` should be replaced by either `cpu`, `cu102`, `cu113`, or `cu116` depending on your PyTorch installation.
| | `cpu` | `cu102` | `cu113` | `cu116` |
|-------------|-------|---------|---------|---------|
| **Linux** | ✅ | ✅ | ✅ | ✅ |
| **Windows** | ✅ | | ✅ | ✅ |
| **macOS** | ✅ | | | |
#### PyTorch 1.11
To install the binaries for PyTorch 1.11.0, simply run
```
pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-1.11.0+${CUDA}.html
```
where `${CUDA}` should be replaced by either `cpu`, `cu102`, `cu113`, or `cu115` depending on your PyTorch installation.
| | `cpu` | `cu102` | `cu113` | `cu115` |
|-------------|-------|---------|---------|---------|
| **Linux** | ✅ | ✅ | ✅ | ✅ |
| **Windows** | ✅ | | ✅ | ✅ |
| **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, and PyTorch 1.10.0/1.10.1/1.10.2 (following the same procedure).
For older versions, you might need to explicitly specify the latest supported version number in order to prevent a manual installation from source.
You can look up the latest supported version number [here](https://data.pyg.org/whl).
### From source
Ensure that at least PyTorch 1.7.0 is installed and verify that `cuda/bin` and `cuda/include` are in your `$PATH` and `$CPATH` respectively, *e.g.*:
```
$ python -c "import torch; print(torch.__version__)"
>>> 1.7.0
$ echo $PATH
>>> /usr/local/cuda/bin:...
$ echo $CPATH
>>> /usr/local/cuda/include:...
```
If you want to additionally build `torch-sparse` with METIS support, *e.g.* for partioning, please download and install the [METIS library](https://web.archive.org/web/20211119110155/http://glaros.dtc.umn.edu/gkhome/metis/metis/download) by following the instructions in the `Install.txt` file.
Note that METIS needs to be installed with 64 bit `IDXTYPEWIDTH` by changing `include/metis.h`.
Afterwards, set the environment variable `WITH_METIS=1`.
Then run:
```
pip install torch-scatter torch-sparse
``` ```
#### 注意事项
+ 若使用pip install下载安装过慢，可添加pypi清华源：-i https://pypi.tuna.tsinghua.edu.cn/simple/
+ ROCM_PATH为dtk的路径，默认为/opt/dtk
## 验证 When running in a docker container without NVIDIA driver, PyTorch needs to evaluate the compute capabilities and may fail.
In this case, ensure that the compute capabilities are set via `TORCH_CUDA_ARCH_LIST`, *e.g.*:
```
export TORCH_CUDA_ARCH_LIST="6.0 6.1 7.2+PTX 7.5+PTX"
```
## Functions
### Coalesce
```
torch_sparse.coalesce(index, value, m, n, op="add") -> (torch.LongTensor, torch.Tensor)
```
Row-wise sorts `index` and removes duplicate entries.
Duplicate entries are removed by scattering them together.
For scattering, any operation of [`torch_scatter`](https://github.com/rusty1s/pytorch_scatter) can be used.
#### Parameters
* **index** *(LongTensor)* - The index tensor of sparse matrix.
* **value** *(Tensor)* - The value tensor of sparse matrix.
* **m** *(int)* - The first dimension of sparse matrix.
* **n** *(int)* - The second dimension of sparse matrix.
* **op** *(string, optional)* - The scatter operation to use. (default: `"add"`)
#### Returns
* **index** *(LongTensor)* - The coalesced index tensor of sparse matrix.
* **value** *(Tensor)* - The coalesced value tensor of sparse matrix.
#### Example
```python ```python
import torch import torch
...@@ -68,10 +160,156 @@ tensor([[6.0, 8.0], ...@@ -68,10 +160,156 @@ tensor([[6.0, 8.0],
[5.0, 6.0]]) [5.0, 6.0]])
``` ```
## Known Issue ### Transpose
- 该库没有基于cpu环境修改，仅支持dcu，请在有dcu卡的环境运行。
- 如需完整使用所有pyg功能，请pip install torch-geometric
## 参考资料 ```
- [README_ORIGIN](README_ORIGIN.md) torch_sparse.transpose(index, value, m, n) -> (torch.LongTensor, torch.Tensor)
- [https://pypi.org/project/torch-sparse/0.6.13/](https://pypi.org/project/torch-sparse/0.6.13/) ```
Transposes dimensions 0 and 1 of a sparse matrix.
#### Parameters
* **index** *(LongTensor)* - The index tensor of sparse matrix.
* **value** *(Tensor)* - The value tensor of sparse matrix.
* **m** *(int)* - The first dimension of sparse matrix.
* **n** *(int)* - The second dimension of sparse matrix.
* **coalesced** *(bool, optional)* - If set to `False`, will not coalesce the output. (default: `True`)
#### Returns
* **index** *(LongTensor)* - The transposed index tensor of sparse matrix.
* **value** *(Tensor)* - The transposed value tensor of sparse matrix.
#### Example
```python
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]])
```
### Sparse Dense Matrix Multiplication
```
torch_sparse.spmm(index, value, m, n, matrix) -> torch.Tensor
```
Matrix product of a sparse matrix with a dense matrix.
#### Parameters
* **index** *(LongTensor)* - The index tensor of sparse matrix.
* **value** *(Tensor)* - The value tensor of sparse matrix.
* **m** *(int)* - The first dimension of sparse matrix.
* **n** *(int)* - The second dimension of sparse matrix.
* **matrix** *(Tensor)* - The dense matrix.
#### Returns
* **out** *(Tensor)* - The dense output matrix.
#### Example
```python
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]])
```
### Sparse Sparse Matrix Multiplication
```
torch_sparse.spspmm(indexA, valueA, indexB, valueB, m, k, n) -> (torch.LongTensor, torch.Tensor)
```
Matrix product of two sparse tensors.
Both input sparse matrices need to be **coalesced** (use the `coalesced` attribute to force).
#### Parameters
* **indexA** *(LongTensor)* - The index tensor of first sparse matrix.
* **valueA** *(Tensor)* - The value tensor of first sparse matrix.
* **indexB** *(LongTensor)* - The index tensor of second sparse matrix.
* **valueB** *(Tensor)* - The value tensor of second sparse matrix.
* **m** *(int)* - The first dimension of first sparse matrix.
* **k** *(int)* - The second dimension of first sparse matrix and first dimension of second sparse matrix.
* **n** *(int)* - The second dimension of second sparse matrix.
* **coalesced** *(bool, optional)*: If set to `True`, will coalesce both input sparse matrices. (default: `False`)
#### Returns
* **index** *(LongTensor)* - The output index tensor of sparse matrix.
* **value** *(Tensor)* - The output value tensor of sparse matrix.
#### Example
```python
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])
```
## Running tests
```
pytest
```
## C++ API
`torch-sparse` also offers a C++ API that contains C++ equivalent of python models.
For this, we need to add `TorchLib` to the `-DCMAKE_PREFIX_PATH` (*e.g.*, it may exists in `{CONDA}/lib/python{X.X}/site-packages/torch` if installed via `conda`):
```
mkdir build
cd build
# Add -DWITH_CUDA=on support for CUDA support
cmake -DCMAKE_PREFIX_PATH="..." ..
make
make install
```
benchmark/main.py 0 → 100644
View file @ 1cb25232
import time
import os.path as osp
import itertools
import argparse
import wget
import torch
from scipy.io import loadmat
from torch_scatter import scatter_add
from torch_sparse.tensor import SparseTensor
short_rows = [
('DIMACS10', 'citationCiteseer'),
('SNAP', 'web-Stanford'),
]
long_rows = [
('Janna', 'StocF-1465'),
('GHS_psdef', 'ldoor'),
]
def download(dataset):
url = 'https://sparse.tamu.edu/mat/{}/{}.mat'
for group, name in itertools.chain(long_rows, short_rows):
if not osp.exists(f'{name}.mat'):
print(f'Downloading {group}/{name}:')
wget.download(url.format(group, name))
print('')
def bold(text, flag=True):
return f'\033[1m{text}\033[0m' if flag else text
@torch.no_grad()
def correctness(dataset):
group, name = dataset
mat_scipy = loadmat(f'{name}.mat')['Problem'][0][0][2].tocsr()
row = torch.from_numpy(mat_scipy.tocoo().row).to(args.device, torch.long)
col = torch.from_numpy(mat_scipy.tocoo().col).to(args.device, torch.long)
mat = SparseTensor(row=row, col=col, sparse_sizes=mat_scipy.shape)
mat.fill_cache_()
mat_pytorch = mat.to_torch_sparse_coo_tensor().coalesce()
for size in sizes:
try:
x = torch.randn((mat.size(1), size), device=args.device)
out1 = mat @ x
out2 = mat_pytorch @ x
assert torch.allclose(out1, out2, atol=1e-4)
except RuntimeError as e:
if 'out of memory' not in str(e):
raise RuntimeError(e)
torch.cuda.empty_cache()
def time_func(func, x):
try:
if torch.cuda.is_available():
torch.cuda.synchronize()
t = time.perf_counter()
if not args.with_backward:
with torch.no_grad():
for _ in range(iters):
func(x)
else:
x = x.requires_grad_()
for _ in range(iters):
out = func(x)
out = out[0] if isinstance(out, tuple) else out
torch.autograd.grad(out, x, out, only_inputs=True)
if torch.cuda.is_available():
torch.cuda.synchronize()
return time.perf_counter() - t
except RuntimeError as e:
if 'out of memory' not in str(e):
raise RuntimeError(e)
torch.cuda.empty_cache()
return float('inf')
def timing(dataset):
group, name = dataset
mat_scipy = loadmat(f'{name}.mat')['Problem'][0][0][2].tocsr()
row = torch.from_numpy(mat_scipy.tocoo().row).to(args.device, torch.long)
col = torch.from_numpy(mat_scipy.tocoo().col).to(args.device, torch.long)
mat = SparseTensor(row=row, col=col, sparse_sizes=mat_scipy.shape)
mat.fill_cache_()
mat_pytorch = mat.to_torch_sparse_coo_tensor().coalesce()
mat_scipy = mat.to_scipy(layout='csr')
def scatter(x):
return scatter_add(x[col], row, dim=0, dim_size=mat_scipy.shape[0])
def spmm_scipy(x):
if x.is_cuda:
raise RuntimeError('out of memory')
return mat_scipy @ x
def spmm_pytorch(x):
return mat_pytorch @ x
def spmm(x):
return mat @ x
t1, t2, t3, t4 = [], [], [], []
for size in sizes:
try:
x = torch.randn((mat.size(1), size), device=args.device)
t1 += [time_func(scatter, x)]
t2 += [time_func(spmm_scipy, x)]
t3 += [time_func(spmm_pytorch, x)]
t4 += [time_func(spmm, x)]
del x
except RuntimeError as e:
if 'out of memory' not in str(e):
raise RuntimeError(e)
torch.cuda.empty_cache()
for t in (t1, t2, t3, t4):
t.append(float('inf'))
ts = torch.tensor([t1, t2, t3, t4])
winner = torch.zeros_like(ts, dtype=torch.bool)
winner[ts.argmin(dim=0), torch.arange(len(sizes))] = 1
winner = winner.tolist()
name = f'{group}/{name}'
print(f'{bold(name)} (avg row length: {mat.avg_row_length():.2f}):')
print('\t'.join([' '] + [f'{size:>5}' for size in sizes]))
print('\t'.join([bold('Scatter ')] +
[bold(f'{t:.5f}', f) for t, f in zip(t1, winner[0])]))
print('\t'.join([bold('SPMM SciPy ')] +
[bold(f'{t:.5f}', f) for t, f in zip(t2, winner[1])]))
print('\t'.join([bold('SPMM PyTorch')] +
[bold(f'{t:.5f}', f) for t, f in zip(t3, winner[2])]))
print('\t'.join([bold('SPMM Own ')] +
[bold(f'{t:.5f}', f) for t, f in zip(t4, winner[3])]))
print()
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--with_backward', action='store_true')
parser.add_argument('--device', type=str, default='cuda')
args = parser.parse_args()
iters = 1 if args.device == 'cpu' else 20
sizes = [1, 16, 32, 64, 128, 256, 512]
sizes = sizes[:4] if args.device == 'cpu' else sizes
for _ in range(10): # Warmup.
torch.randn(100, 100, device=args.device).sum()
for dataset in itertools.chain(short_rows, long_rows):
download(dataset)
correctness(dataset)
timing(dataset)
benchmark/ptr2ind.py 0 → 100644
View file @ 1cb25232
import os
import wget
import time
import errno
import argparse
import os.path as osp
import torch
from scipy.io import loadmat
parser = argparse.ArgumentParser()
parser.add_argument('--root', type=str, default='/tmp/test_ptr2ind')
args = parser.parse_args()
matrices = [
('DIMACS10', 'citationCiteseer'),
('SNAP', 'web-Stanford'),
('Janna', 'StocF-1465'),
('GHS_psdef', 'ldoor'),
]
def get_torch_sparse_coo_tensor(root, group, name):
path = osp.join(root, f'{name}.mat')
if not osp.exists(path):
try:
os.makedirs(root)
except OSError as e:
if e.errno != errno.EEXIST and osp.isdir(path):
raise e
url = f'https://sparse.tamu.edu/mat/{group}/{name}.mat'
print(f'Downloading {group}/{name}:')
wget.download(url, path)
matrix = loadmat(path)['Problem'][0][0][2].tocoo()
row = torch.from_numpy(matrix.row).to(torch.long)
col = torch.from_numpy(matrix.col).to(torch.long)
index = torch.stack([row, col], dim=0)
value = torch.from_numpy(matrix.data).to(torch.float)
print(f'{name}.mat: shape={matrix.shape} nnz={row.numel()}')
return torch.sparse_coo_tensor(index, value, matrix.shape).coalesce()
def time_func(matrix, op, duration=5.0, warmup=1.0):
t = time.time()
while (time.time() - t) < warmup:
op(matrix)
torch.cuda.synchronize()
count = 0
t = time.time()
while (time.time() - t) < duration:
op(matrix)
count += 1
torch.cuda.synchronize()
return (time.time() - t) / count
def bucketize(matrix):
row_indices = matrix.indices()[0]
arange = torch.arange(matrix.size(0) + 1, device=row_indices.device)
return torch.bucketize(arange, row_indices)
def convert_coo_to_csr(matrix):
row_indices = matrix.indices()[0]
return torch._convert_coo_to_csr(row_indices, matrix.size(0))
for device in ['cpu', 'cuda']:
print('DEVICE:', device)
for group, name in matrices:
matrix = get_torch_sparse_coo_tensor(args.root, group, name)
matrix = matrix.to(device)
out1 = bucketize(matrix)
out2 = convert_coo_to_csr(matrix)
assert out1.tolist() == out2.tolist()
t = time_func(matrix, bucketize, duration=5.0, warmup=1.0)
print('old impl:', t)
t = time_func(matrix, convert_coo_to_csr, duration=5.0, warmup=1.0)
print('new impl:', t)
print()
cmake/TorchSparseConfig.cmake.in 0 → 100644
View file @ 1cb25232
# TorchSparseConfig.cmake
# --------------------
#
# Exported targets:: Sparse
#
@PACKAGE_INIT@
set(PN TorchSparse)
set(${PN}_INCLUDE_DIR "${PACKAGE_PREFIX_DIR}/@CMAKE_INSTALL_INCLUDEDIR@")
set(${PN}_LIBRARY "")
set(${PN}_DEFINITIONS USING_${PN})
check_required_components(${PN})
if(NOT (CMAKE_VERSION VERSION_LESS 3.0))
#-----------------------------------------------------------------------------
# Don't include targets if this file is being picked up by another
# project which has already built this as a subproject
#-----------------------------------------------------------------------------
if(NOT TARGET ${PN}::TorchSparse)
include("${CMAKE_CURRENT_LIST_DIR}/${PN}Targets.cmake")
if(NOT TARGET torch_library)
find_package(Torch REQUIRED)
endif()
if(NOT TARGET Python3::Python)
find_package(Python3 COMPONENTS Development)
endif()
target_link_libraries(TorchSparse::TorchSparse INTERFACE ${TORCH_LIBRARIES} Python3::Python)
if(@WITH_CUDA@)
target_compile_definitions(TorchSparse::TorchSparse INTERFACE WITH_CUDA)
endif()
endif()
endif()
conda/pytorch-sparse/README.md 0 → 100644
View file @ 1cb25232
```
./build_conda.sh 3.9 1.12.0 cu113 # python, pytorch and cuda version
```
conda/pytorch-sparse/bld.bat 0 → 100644
View file @ 1cb25232
copy "C:\\Program Files (x86)\\Microsoft Visual Studio\\2019\\Enterprise\\VC\\Tools\\MSVC\\14.29.30133\\lib\\x64\\metis.lib" %LIBRARY_LIB%
if errorlevel 1 exit 1
copy "C:\\Program Files (x86)\\Microsoft Visual Studio\\2019\\Enterprise\\VC\\Tools\\MSVC\\14.29.30133\\include\\metis.h" %LIBRARY_INC%
if errorlevel 1 exit 1
"%PYTHON%" -m pip install .
if errorlevel 1 exit 1
conda/pytorch-sparse/build.sh 0 → 100644
View file @ 1cb25232
$PYTHON -m pip install .
conda/pytorch-sparse/build_conda.sh 0 → 100755
View file @ 1cb25232
#!/bin/bash
export PYTHON_VERSION=$1
export TORCH_VERSION=$2
export CUDA_VERSION=$3
export CONDA_PYTORCH_CONSTRAINT="pytorch==${TORCH_VERSION%.*}.*"
if [ "${CUDA_VERSION}" = "cpu" ]; then
export CONDA_CUDATOOLKIT_CONSTRAINT="cpuonly # [not osx]"
else
case $CUDA_VERSION in
cu116)
export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit==11.6.*"
;;
cu115)
export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit==11.5.*"
;;
cu113)
export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit==11.3.*"
;;
cu111)
export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit==11.1.*"
;;
cu102)
export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit==10.2.*"
;;
cu101)
export CONDA_CUDATOOLKIT_CONSTRAINT="cudatoolkit==10.1.*"
;;
*)
echo "Unrecognized CUDA_VERSION=$CUDA_VERSION"
exit 1
;;
esac
fi
echo "PyTorch $TORCH_VERSION+$CUDA_VERSION"
echo "- $CONDA_PYTORCH_CONSTRAINT"
echo "- $CONDA_CUDATOOLKIT_CONSTRAINT"
if [ "${CUDA_VERSION}" = "cu116" ]; then
conda build . -c pytorch -c default -c nvidia -c conda-forge --output-folder "$HOME/conda-bld"
else
conda build . -c pytorch -c default -c nvidia --output-folder "$HOME/conda-bld"
fi
conda/pytorch-sparse/meta.yaml 0 → 100644
View file @ 1cb25232
package:
name: pytorch-sparse
version: 0.6.15
source:
path: ../..
requirements:
build:
- {{ compiler('c') }} # [win]
host:
- pip
- python {{ environ.get('PYTHON_VERSION') }}
- {{ environ.get('CONDA_PYTORCH_CONSTRAINT') }}
- {{ environ.get('CONDA_CUDATOOLKIT_CONSTRAINT') }}
run:
- scipy
- pytorch-scatter
- python {{ environ.get('PYTHON_VERSION') }}
- {{ environ.get('CONDA_PYTORCH_CONSTRAINT') }}
- {{ environ.get('CONDA_CUDATOOLKIT_CONSTRAINT') }}
build:
string: py{{ environ.get('PYTHON_VERSION').replace('.', '') }}_torch_{{ environ['TORCH_VERSION'] }}_{{ environ['CUDA_VERSION'] }}
script_env:
- FORCE_CUDA
- TORCH_CUDA_ARCH_LIST
- WITH_METIS=1
preserve_egg_dir: True
test:
imports:
- torch_sparse
about:
home: https://github.com/rusty1s/pytorch_sparse
license: MIT
summary: PyTorch Extension Library of Optimized Autograd Sparse Matrix Operations
csrc/convert.cpp
View file @ 1cb25232
...@@ -5,13 +5,13 @@ ...@@ -5,13 +5,13 @@
#include "cpu/convert_cpu.h" #include "cpu/convert_cpu.h"
#ifdef WITH_HIP #ifdef WITH_CUDA
#include "hip/convert_hip.h" #include "cuda/convert_cuda.h"
#endif #endif
#ifdef _WIN32 #ifdef _WIN32
#ifdef WITH_PYTHON #ifdef WITH_PYTHON
#ifdef WITH_HIP #ifdef WITH_CUDA
PyMODINIT_FUNC PyInit__convert_cuda(void) { return NULL; } PyMODINIT_FUNC PyInit__convert_cuda(void) { return NULL; }
#else #else
PyMODINIT_FUNC PyInit__convert_cpu(void) { return NULL; } PyMODINIT_FUNC PyInit__convert_cpu(void) { return NULL; }
...@@ -21,7 +21,7 @@ PyMODINIT_FUNC PyInit__convert_cpu(void) { return NULL; } ...@@ -21,7 +21,7 @@ PyMODINIT_FUNC PyInit__convert_cpu(void) { return NULL; }
SPARSE_API torch::Tensor ind2ptr(torch::Tensor ind, int64_t M) { SPARSE_API torch::Tensor ind2ptr(torch::Tensor ind, int64_t M) {
if (ind.device().is_cuda()) { if (ind.device().is_cuda()) {
#ifdef WITH_HIP #ifdef WITH_CUDA
return ind2ptr_cuda(ind, M); return ind2ptr_cuda(ind, M);
#else #else
AT_ERROR("Not compiled with CUDA support"); AT_ERROR("Not compiled with CUDA support");
...@@ -33,7 +33,7 @@ SPARSE_API torch::Tensor ind2ptr(torch::Tensor ind, int64_t M) { ...@@ -33,7 +33,7 @@ SPARSE_API torch::Tensor ind2ptr(torch::Tensor ind, int64_t M) {
SPARSE_API torch::Tensor ptr2ind(torch::Tensor ptr, int64_t E) { SPARSE_API torch::Tensor ptr2ind(torch::Tensor ptr, int64_t E) {
if (ptr.device().is_cuda()) { if (ptr.device().is_cuda()) {
#ifdef WITH_HIP #ifdef WITH_CUDA
return ptr2ind_cuda(ptr, E); return ptr2ind_cuda(ptr, E);
#else #else
AT_ERROR("Not compiled with CUDA support"); AT_ERROR("Not compiled with CUDA support");
......
csrc/cpu/diag_cpu.cpp
View file @ 1cb25232
...@@ -13,7 +13,7 @@ torch::Tensor non_diag_mask_cpu(torch::Tensor row, torch::Tensor col, int64_t M, ...@@ -13,7 +13,7 @@ torch::Tensor non_diag_mask_cpu(torch::Tensor row, torch::Tensor col, int64_t M,
auto row_data = row.data_ptr<int64_t>(); auto row_data = row.data_ptr<int64_t>();
auto col_data = col.data_ptr<int64_t>(); auto col_data = col.data_ptr<int64_t>();
auto mask = torch::zeros(E + num_diag, row.options().dtype(torch::kBool)); auto mask = torch::zeros({E + num_diag}, row.options().dtype(torch::kBool));
auto mask_data = mask.data_ptr<bool>(); auto mask_data = mask.data_ptr<bool>();
int64_t r, c; int64_t r, c;
......
csrc/cpu/ego_sample_cpu.cpp
View file @ 1cb25232
...@@ -19,8 +19,6 @@ ego_k_hop_sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, ...@@ -19,8 +19,6 @@ ego_k_hop_sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col,
torch::Tensor idx, int64_t depth, torch::Tensor idx, int64_t depth,
int64_t num_neighbors, bool replace) { int64_t num_neighbors, bool replace) {
srand(time(NULL) + 1000 * getpid()); // Initialize random seed.
std::vector<torch::Tensor> out_rowptrs(idx.numel() + 1); std::vector<torch::Tensor> out_rowptrs(idx.numel() + 1);
std::vector<torch::Tensor> out_cols(idx.numel()); std::vector<torch::Tensor> out_cols(idx.numel());
std::vector<torch::Tensor> out_n_ids(idx.numel()); std::vector<torch::Tensor> out_n_ids(idx.numel());
...@@ -56,14 +54,14 @@ ego_k_hop_sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, ...@@ -56,14 +54,14 @@ ego_k_hop_sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col,
} }
} else if (replace) { } else if (replace) {
for (int64_t j = 0; j < num_neighbors; j++) { for (int64_t j = 0; j < num_neighbors; j++) {
w = col_data[row_start + (rand() % row_count)]; w = col_data[row_start + uniform_randint(row_count)];
n_id_set.insert(w); n_id_set.insert(w);
n_ids.push_back(w); n_ids.push_back(w);
} }
} else { } else {
std::unordered_set<int64_t> perm; std::unordered_set<int64_t> perm;
for (int64_t j = row_count - num_neighbors; j < row_count; j++) { for (int64_t j = row_count - num_neighbors; j < row_count; j++) {
if (!perm.insert(rand() % j).second) { if (!perm.insert(uniform_randint(j)).second) {
perm.insert(j); perm.insert(j);
} }
} }
......
csrc/cpu/hgt_sample_cpu.cpp
View file @ 1cb25232
...@@ -105,8 +105,6 @@ hgt_sample_cpu(const c10::Dict<rel_t, torch::Tensor> &colptr_dict, ...@@ -105,8 +105,6 @@ hgt_sample_cpu(const c10::Dict<rel_t, torch::Tensor> &colptr_dict,
const c10::Dict<node_t, vector<int64_t>> &num_samples_dict, const c10::Dict<node_t, vector<int64_t>> &num_samples_dict,
const int64_t num_hops) { const int64_t num_hops) {
srand(time(NULL) + 1000 * getpid()); // Initialize random seed.
// Create a mapping to convert single string relations to edge type triplets: // Create a mapping to convert single string relations to edge type triplets:
unordered_map<rel_t, edge_t> to_edge_type; unordered_map<rel_t, edge_t> to_edge_type;
for (const auto &kv : colptr_dict) { for (const auto &kv : colptr_dict) {
...@@ -226,12 +224,10 @@ hgt_sample_cpu(const c10::Dict<rel_t, torch::Tensor> &colptr_dict, ...@@ -226,12 +224,10 @@ hgt_sample_cpu(const c10::Dict<rel_t, torch::Tensor> &colptr_dict,
} }
} }
} }
if (rows.size() > 0) { out_row_dict.insert(rel_type, from_vector<int64_t>(rows));
out_row_dict.insert(rel_type, from_vector<int64_t>(rows)); out_col_dict.insert(rel_type, from_vector<int64_t>(cols));
out_col_dict.insert(rel_type, from_vector<int64_t>(cols)); out_edge_dict.insert(rel_type, from_vector<int64_t>(edges));
out_edge_dict.insert(rel_type, from_vector<int64_t>(edges)); }
}
}
// Generate tensor-valued output node dictionary (line 20): // Generate tensor-valued output node dictionary (line 20):
for (const auto &kv : nodes_dict) { for (const auto &kv : nodes_dict) {
......
csrc/cpu/metis_cpu.cpp
View file @ 1cb25232
...@@ -44,7 +44,7 @@ torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col, ...@@ -44,7 +44,7 @@ torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col,
vwgt = optional_node_weight.value().data_ptr<int64_t>(); vwgt = optional_node_weight.value().data_ptr<int64_t>();
int64_t objval = -1; int64_t objval = -1;
auto part = torch::empty(nvtxs, rowptr.options()); auto part = torch::empty({nvtxs}, rowptr.options());
auto part_data = part.data_ptr<int64_t>(); auto part_data = part.data_ptr<int64_t>();
if (recursive) { if (recursive) {
...@@ -99,7 +99,7 @@ mt_partition_cpu(torch::Tensor rowptr, torch::Tensor col, ...@@ -99,7 +99,7 @@ mt_partition_cpu(torch::Tensor rowptr, torch::Tensor col,
mtmetis_pid_type nparts = num_parts; mtmetis_pid_type nparts = num_parts;
mtmetis_wgt_type objval = -1; mtmetis_wgt_type objval = -1;
auto part = torch::empty(nvtxs, rowptr.options()); auto part = torch::empty({nvtxs}, rowptr.options());
mtmetis_pid_type *part_data = (mtmetis_pid_type *)part.data_ptr<int64_t>(); mtmetis_pid_type *part_data = (mtmetis_pid_type *)part.data_ptr<int64_t>();
double *opts = mtmetis_init_options(); double *opts = mtmetis_init_options();
......
csrc/cpu/neighbor_sample_cpu.cpp
View file @ 1cb25232
...@@ -10,16 +10,16 @@ using namespace std; ...@@ -10,16 +10,16 @@ using namespace std;
namespace { namespace {
typedef phmap::flat_hash_map<pair<int64_t, int64_t>, int64_t> temporarl_edge_dict;
template <bool replace, bool directed> template <bool replace, bool directed>
tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor> tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
sample(const torch::Tensor &colptr, const torch::Tensor &row, sample(const torch::Tensor &colptr, const torch::Tensor &row,
const torch::Tensor &input_node, const vector<int64_t> num_neighbors) { const torch::Tensor &input_node, const vector<int64_t> num_neighbors) {
srand(time(NULL) + 1000 * getpid()); // Initialize random seed.
// Initialize some data structures for the sampling process: // Initialize some data structures for the sampling process:
vector<int64_t> samples; vector<int64_t> samples;
unordered_map<int64_t, int64_t> to_local_node; phmap::flat_hash_map<int64_t, int64_t> to_local_node;
auto *colptr_data = colptr.data_ptr<int64_t>(); auto *colptr_data = colptr.data_ptr<int64_t>();
auto *row_data = row.data_ptr<int64_t>(); auto *row_data = row.data_ptr<int64_t>();
...@@ -59,7 +59,7 @@ sample(const torch::Tensor &colptr, const torch::Tensor &row, ...@@ -59,7 +59,7 @@ sample(const torch::Tensor &colptr, const torch::Tensor &row,
} }
} else if (replace) { } else if (replace) {
for (int64_t j = 0; j < num_samples; j++) { for (int64_t j = 0; j < num_samples; j++) {
const int64_t offset = col_start + rand() % col_count; const int64_t offset = col_start + uniform_randint(col_count);
const int64_t &v = row_data[offset]; const int64_t &v = row_data[offset];
const auto res = to_local_node.insert({v, samples.size()}); const auto res = to_local_node.insert({v, samples.size()});
if (res.second) if (res.second)
...@@ -73,7 +73,7 @@ sample(const torch::Tensor &colptr, const torch::Tensor &row, ...@@ -73,7 +73,7 @@ sample(const torch::Tensor &colptr, const torch::Tensor &row,
} else { } else {
unordered_set<int64_t> rnd_indices; unordered_set<int64_t> rnd_indices;
for (int64_t j = col_count - num_samples; j < col_count; j++) { for (int64_t j = col_count - num_samples; j < col_count; j++) {
int64_t rnd = rand() % j; int64_t rnd = uniform_randint(j);
if (!rnd_indices.insert(rnd).second) { if (!rnd_indices.insert(rnd).second) {
rnd = j; rnd = j;
rnd_indices.insert(j); rnd_indices.insert(j);
...@@ -95,7 +95,7 @@ sample(const torch::Tensor &colptr, const torch::Tensor &row, ...@@ -95,7 +95,7 @@ sample(const torch::Tensor &colptr, const torch::Tensor &row,
} }
if (!directed) { if (!directed) {
unordered_map<int64_t, int64_t>::iterator iter; phmap::flat_hash_map<int64_t, int64_t>::iterator iter;
for (int64_t i = 0; i < (int64_t)samples.size(); i++) { for (int64_t i = 0; i < (int64_t)samples.size(); i++) {
const auto &w = samples[i]; const auto &w = samples[i];
const auto &col_start = colptr_data[w]; const auto &col_start = colptr_data[w];
...@@ -116,7 +116,20 @@ sample(const torch::Tensor &colptr, const torch::Tensor &row, ...@@ -116,7 +116,20 @@ sample(const torch::Tensor &colptr, const torch::Tensor &row,
from_vector<int64_t>(cols), from_vector<int64_t>(edges)); from_vector<int64_t>(cols), from_vector<int64_t>(edges));
} }
template <bool replace, bool directed> inline bool satisfy_time(const c10::Dict<node_t, torch::Tensor> &node_time_dict,
const node_t &src_node_type, int64_t dst_time,
int64_t src_node) {
try {
// Check whether src -> dst obeys the time constraint
const torch::Tensor &src_node_time = node_time_dict.at(src_node_type);
return src_node_time.data_ptr<int64_t>()[src_node] <= dst_time;
} catch (const std::out_of_range& e) {
// If no time is given, fall back to normal sampling
return true;
}
}
template <bool replace, bool directed, bool temporal>
tuple<c10::Dict<node_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>, tuple<c10::Dict<node_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>,
c10::Dict<rel_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>> c10::Dict<rel_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>>
hetero_sample(const vector<node_t> &node_types, hetero_sample(const vector<node_t> &node_types,
...@@ -125,24 +138,29 @@ hetero_sample(const vector<node_t> &node_types, ...@@ -125,24 +138,29 @@ hetero_sample(const vector<node_t> &node_types,
const c10::Dict<rel_t, torch::Tensor> &row_dict, const c10::Dict<rel_t, torch::Tensor> &row_dict,
const c10::Dict<node_t, torch::Tensor> &input_node_dict, const c10::Dict<node_t, torch::Tensor> &input_node_dict,
const c10::Dict<rel_t, vector<int64_t>> &num_neighbors_dict, const c10::Dict<rel_t, vector<int64_t>> &num_neighbors_dict,
const c10::Dict<node_t, torch::Tensor> &node_time_dict,
const int64_t num_hops) { const int64_t num_hops) {
srand(time(NULL) + 1000 * getpid()); // Initialize random seed.
// Create a mapping to convert single string relations to edge type triplets: // Create a mapping to convert single string relations to edge type triplets:
unordered_map<rel_t, edge_t> to_edge_type; phmap::flat_hash_map<rel_t, edge_t> to_edge_type;
for (const auto &k : edge_types) for (const auto &k : edge_types)
to_edge_type[get<0>(k) + "__" + get<1>(k) + "__" + get<2>(k)] = k; to_edge_type[get<0>(k) + "__" + get<1>(k) + "__" + get<2>(k)] = k;
// Initialize some data structures for the sampling process: // Initialize some data structures for the sampling process:
unordered_map<node_t, vector<int64_t>> samples_dict; phmap::flat_hash_map<node_t, vector<int64_t>> samples_dict;
unordered_map<node_t, unordered_map<int64_t, int64_t>> to_local_node_dict; phmap::flat_hash_map<node_t, vector<pair<int64_t, int64_t>>> temp_samples_dict;
phmap::flat_hash_map<node_t, phmap::flat_hash_map<int64_t, int64_t>> to_local_node_dict;
phmap::flat_hash_map<node_t, temporarl_edge_dict> temp_to_local_node_dict;
phmap::flat_hash_map<node_t, vector<int64_t>> root_time_dict;
for (const auto &node_type : node_types) { for (const auto &node_type : node_types) {
samples_dict[node_type]; samples_dict[node_type];
temp_samples_dict[node_type];
to_local_node_dict[node_type]; to_local_node_dict[node_type];
temp_to_local_node_dict[node_type];
root_time_dict[node_type];
} }
unordered_map<rel_t, vector<int64_t>> rows_dict, cols_dict, edges_dict; phmap::flat_hash_map<rel_t, vector<int64_t>> rows_dict, cols_dict, edges_dict;
for (const auto &kv : colptr_dict) { for (const auto &kv : colptr_dict) {
const auto &rel_type = kv.key(); const auto &rel_type = kv.key();
rows_dict[rel_type]; rows_dict[rel_type];
...@@ -156,41 +174,82 @@ hetero_sample(const vector<node_t> &node_types, ...@@ -156,41 +174,82 @@ hetero_sample(const vector<node_t> &node_types,
const torch::Tensor &input_node = kv.value(); const torch::Tensor &input_node = kv.value();
const auto *input_node_data = input_node.data_ptr<int64_t>(); const auto *input_node_data = input_node.data_ptr<int64_t>();
int64_t *node_time_data;
if (temporal) {
const torch::Tensor &node_time = node_time_dict.at(node_type);
node_time_data = node_time.data_ptr<int64_t>();
}
auto &samples = samples_dict.at(node_type); auto &samples = samples_dict.at(node_type);
auto &temp_samples = temp_samples_dict.at(node_type);
auto &to_local_node = to_local_node_dict.at(node_type); auto &to_local_node = to_local_node_dict.at(node_type);
auto &temp_to_local_node = temp_to_local_node_dict.at(node_type);
auto &root_time = root_time_dict.at(node_type);
for (int64_t i = 0; i < input_node.numel(); i++) { for (int64_t i = 0; i < input_node.numel(); i++) {
const auto &v = input_node_data[i]; const auto &v = input_node_data[i];
samples.push_back(v); if (temporal) {
to_local_node.insert({v, i}); temp_samples.push_back({v, i});
temp_to_local_node.insert({{v, i}, i});
} else {
samples.push_back(v);
to_local_node.insert({v, i});
}
if (temporal)
root_time.push_back(node_time_data[v]);
} }
} }
unordered_map<node_t, pair<int64_t, int64_t>> slice_dict; phmap::flat_hash_map<node_t, pair<int64_t, int64_t>> slice_dict;
for (const auto &kv : samples_dict) if (temporal) {
slice_dict[kv.first] = {0, kv.second.size()}; for (const auto &kv : temp_samples_dict) {
slice_dict[kv.first] = {0, kv.second.size()};
}
} else {
for (const auto &kv : samples_dict)
slice_dict[kv.first] = {0, kv.second.size()};
}
vector<rel_t> all_rel_types;
for (const auto &kv : num_neighbors_dict) {
all_rel_types.push_back(kv.key());
}
std::sort(all_rel_types.begin(), all_rel_types.end());
for (int64_t ell = 0; ell < num_hops; ell++) { for (int64_t ell = 0; ell < num_hops; ell++) {
for (const auto &kv : num_neighbors_dict) { for (const auto &rel_type : all_rel_types) {
const auto &rel_type = kv.key();
const auto &edge_type = to_edge_type[rel_type]; const auto &edge_type = to_edge_type[rel_type];
const auto &src_node_type = get<0>(edge_type); const auto &src_node_type = get<0>(edge_type);
const auto &dst_node_type = get<2>(edge_type); const auto &dst_node_type = get<2>(edge_type);
const auto num_samples = kv.value()[ell]; const auto num_samples = num_neighbors_dict.at(rel_type)[ell];
const auto &dst_samples = samples_dict.at(dst_node_type); const auto &dst_samples = samples_dict.at(dst_node_type);
const auto &temp_dst_samples = temp_samples_dict.at(dst_node_type);
auto &src_samples = samples_dict.at(src_node_type); auto &src_samples = samples_dict.at(src_node_type);
auto &temp_src_samples = temp_samples_dict.at(src_node_type);
auto &to_local_src_node = to_local_node_dict.at(src_node_type); auto &to_local_src_node = to_local_node_dict.at(src_node_type);
auto &temp_to_local_src_node = temp_to_local_node_dict.at(src_node_type);
const auto *colptr_data = ((torch::Tensor)colptr_dict.at(rel_type)).data_ptr<int64_t>(); const torch::Tensor &colptr = colptr_dict.at(rel_type);
const auto *row_data = ((torch::Tensor)row_dict.at(rel_type)).data_ptr<int64_t>(); const auto *colptr_data = colptr.data_ptr<int64_t>();
const torch::Tensor &row = row_dict.at(rel_type);
const auto *row_data = row.data_ptr<int64_t>();
auto &rows = rows_dict.at(rel_type); auto &rows = rows_dict.at(rel_type);
auto &cols = cols_dict.at(rel_type); auto &cols = cols_dict.at(rel_type);
auto &edges = edges_dict.at(rel_type); auto &edges = edges_dict.at(rel_type);
// For temporal sampling, sampled nodes cannot have a timestamp greater
// than the timestamp of the root nodes:
const auto &dst_root_time = root_time_dict.at(dst_node_type);
auto &src_root_time = root_time_dict.at(src_node_type);
const auto &begin = slice_dict.at(dst_node_type).first; const auto &begin = slice_dict.at(dst_node_type).first;
const auto &end = slice_dict.at(dst_node_type).second; const auto &end = slice_dict.at(dst_node_type).second;
for (int64_t i = begin; i < end; i++) { for (int64_t i = begin; i < end; i++) {
const auto &w = dst_samples[i]; const auto &w = temporal ? temp_dst_samples[i].first : dst_samples[i];
const int64_t root_w = temporal ? temp_dst_samples[i].second : -1;
int64_t dst_time = 0;
if (temporal)
dst_time = dst_root_time[i];
const auto &col_start = colptr_data[w]; const auto &col_start = colptr_data[w];
const auto &col_end = colptr_data[w + 1]; const auto &col_end = colptr_data[w + 1];
const auto col_count = col_end - col_start; const auto col_count = col_end - col_start;
...@@ -199,60 +258,125 @@ hetero_sample(const vector<node_t> &node_types, ...@@ -199,60 +258,125 @@ hetero_sample(const vector<node_t> &node_types,
continue; continue;
if ((num_samples < 0) || (!replace && (num_samples >= col_count))) { if ((num_samples < 0) || (!replace && (num_samples >= col_count))) {
// Select all neighbors:
for (int64_t offset = col_start; offset < col_end; offset++) { for (int64_t offset = col_start; offset < col_end; offset++) {
const int64_t &v = row_data[offset]; const int64_t &v = row_data[offset];
const auto res = to_local_src_node.insert({v, src_samples.size()}); if (temporal) {
if (res.second) if (!satisfy_time(node_time_dict, src_node_type, dst_time, v))
src_samples.push_back(v); continue;
if (directed) { // force disjoint of computation tree based on source batch idx.
// note that the sampling always needs to have directed=True
// for temporal case
// to_local_src_node is not used for temporal / directed case
const auto res = temp_to_local_src_node.insert({{v, root_w}, (int64_t)temp_src_samples.size()});
if (res.second) {
temp_src_samples.push_back({v, root_w});
src_root_time.push_back(dst_time);
}
cols.push_back(i); cols.push_back(i);
rows.push_back(res.first->second); rows.push_back(res.first->second);
edges.push_back(offset); edges.push_back(offset);
} else {
const auto res = to_local_src_node.insert({v, src_samples.size()});
if (res.second)
src_samples.push_back(v);
if (directed) {
cols.push_back(i);
rows.push_back(res.first->second);
edges.push_back(offset);
}
} }
} }
} else if (replace) { } else if (replace) {
for (int64_t j = 0; j < num_samples; j++) { // Sample with replacement:
const int64_t offset = col_start + rand() % col_count; int64_t num_neighbors = 0;
while (num_neighbors < num_samples) {
const int64_t offset = col_start + uniform_randint(col_count);
const int64_t &v = row_data[offset]; const int64_t &v = row_data[offset];
const auto res = to_local_src_node.insert({v, src_samples.size()}); if (temporal) {
if (res.second) // TODO Infinity loop if no neighbor satisfies time constraint:
src_samples.push_back(v); if (!satisfy_time(node_time_dict, src_node_type, dst_time, v))
if (directed) { continue;
// force disjoint of computation tree based on source batch idx.
// note that the sampling always needs to have directed=True
// for temporal case
const auto res = temp_to_local_src_node.insert({{v, root_w}, (int64_t)temp_src_samples.size()});
if (res.second) {
temp_src_samples.push_back({v, root_w});
src_root_time.push_back(dst_time);
}
cols.push_back(i); cols.push_back(i);
rows.push_back(res.first->second); rows.push_back(res.first->second);
edges.push_back(offset); edges.push_back(offset);
} else {
const auto res = to_local_src_node.insert({v, src_samples.size()});
if (res.second)
src_samples.push_back(v);
if (directed) {
cols.push_back(i);
rows.push_back(res.first->second);
edges.push_back(offset);
}
} }
num_neighbors += 1;
} }
} else { } else {
// Sample without replacement:
unordered_set<int64_t> rnd_indices; unordered_set<int64_t> rnd_indices;
for (int64_t j = col_count - num_samples; j < col_count; j++) { for (int64_t j = col_count - num_samples; j < col_count; j++) {
int64_t rnd = rand() % j; int64_t rnd = uniform_randint(j);
if (!rnd_indices.insert(rnd).second) { if (!rnd_indices.insert(rnd).second) {
rnd = j; rnd = j;
rnd_indices.insert(j); rnd_indices.insert(j);
} }
const int64_t offset = col_start + rnd; const int64_t offset = col_start + rnd;
const int64_t &v = row_data[offset]; const int64_t &v = row_data[offset];
const auto res = to_local_src_node.insert({v, src_samples.size()}); if (temporal) {
if (res.second) if (!satisfy_time(node_time_dict, src_node_type, dst_time, v))
src_samples.push_back(v); continue;
if (directed) { // force disjoint of computation tree based on source batch idx.
// note that the sampling always needs to have directed=True
// for temporal case
const auto res = temp_to_local_src_node.insert({{v, root_w}, (int64_t)temp_src_samples.size()});
if (res.second) {
temp_src_samples.push_back({v, root_w});
src_root_time.push_back(dst_time);
}
cols.push_back(i); cols.push_back(i);
rows.push_back(res.first->second); rows.push_back(res.first->second);
edges.push_back(offset); edges.push_back(offset);
} else {
const auto res = to_local_src_node.insert({v, src_samples.size()});
if (res.second)
src_samples.push_back(v);
if (directed) {
cols.push_back(i);
rows.push_back(res.first->second);
edges.push_back(offset);
}
} }
} }
} }
} }
} }
for (const auto &kv : samples_dict) { if (temporal) {
slice_dict[kv.first] = {slice_dict.at(kv.first).second, kv.second.size()}; for (const auto &kv : temp_samples_dict) {
slice_dict[kv.first] = {0, kv.second.size()};
}
} else {
for (const auto &kv : samples_dict)
slice_dict[kv.first] = {0, kv.second.size()};
} }
} }
// Temporal sample disable undirected
assert(!(temporal && !directed));
if (!directed) { // Construct the subgraph among the sampled nodes: if (!directed) { // Construct the subgraph among the sampled nodes:
unordered_map<int64_t, int64_t>::iterator iter; phmap::flat_hash_map<int64_t, int64_t>::iterator iter;
for (const auto &kv : colptr_dict) { for (const auto &kv : colptr_dict) {
const auto &rel_type = kv.key(); const auto &rel_type = kv.key();
const auto &edge_type = to_edge_type[rel_type]; const auto &edge_type = to_edge_type[rel_type];
...@@ -262,7 +386,8 @@ hetero_sample(const vector<node_t> &node_types, ...@@ -262,7 +386,8 @@ hetero_sample(const vector<node_t> &node_types,
auto &to_local_src_node = to_local_node_dict.at(src_node_type); auto &to_local_src_node = to_local_node_dict.at(src_node_type);
const auto *colptr_data = ((torch::Tensor)kv.value()).data_ptr<int64_t>(); const auto *colptr_data = ((torch::Tensor)kv.value()).data_ptr<int64_t>();
const auto *row_data = ((torch::Tensor)row_dict.at(rel_type)).data_ptr<int64_t>(); const auto *row_data =
((torch::Tensor)row_dict.at(rel_type)).data_ptr<int64_t>();
auto &rows = rows_dict.at(rel_type); auto &rows = rows_dict.at(rel_type);
auto &cols = cols_dict.at(rel_type); auto &cols = cols_dict.at(rel_type);
...@@ -285,6 +410,18 @@ hetero_sample(const vector<node_t> &node_types, ...@@ -285,6 +410,18 @@ hetero_sample(const vector<node_t> &node_types,
} }
} }
// Construct samples dictionary from temporal sample dictionary.
if (temporal) {
for (const auto &kv : temp_samples_dict) {
const auto &node_type = kv.first;
const auto &samples = kv.second;
samples_dict[node_type].reserve(samples.size());
for (const auto &v : samples) {
samples_dict[node_type].push_back(v.first);
}
}
}
return make_tuple(from_vector<node_t, int64_t>(samples_dict), return make_tuple(from_vector<node_t, int64_t>(samples_dict),
from_vector<rel_t, int64_t>(rows_dict), from_vector<rel_t, int64_t>(rows_dict),
from_vector<rel_t, int64_t>(cols_dict), from_vector<rel_t, int64_t>(cols_dict),
...@@ -320,21 +457,51 @@ hetero_neighbor_sample_cpu( ...@@ -320,21 +457,51 @@ hetero_neighbor_sample_cpu(
const c10::Dict<rel_t, vector<int64_t>> &num_neighbors_dict, const c10::Dict<rel_t, vector<int64_t>> &num_neighbors_dict,
const int64_t num_hops, const bool replace, const bool directed) { const int64_t num_hops, const bool replace, const bool directed) {
c10::Dict<node_t, torch::Tensor> node_time_dict; // Empty dictionary.
if (replace && directed) { if (replace && directed) {
return hetero_sample<true, true>(node_types, edge_types, colptr_dict, return hetero_sample<true, true, false>(
row_dict, input_node_dict, node_types, edge_types, colptr_dict, row_dict, input_node_dict,
num_neighbors_dict, num_hops); num_neighbors_dict, node_time_dict, num_hops);
} else if (replace && !directed) { } else if (replace && !directed) {
return hetero_sample<true, false>(node_types, edge_types, colptr_dict, return hetero_sample<true, false, false>(
row_dict, input_node_dict, node_types, edge_types, colptr_dict, row_dict, input_node_dict,
num_neighbors_dict, num_hops); num_neighbors_dict, node_time_dict, num_hops);
} else if (!replace && directed) { } else if (!replace && directed) {
return hetero_sample<false, true>(node_types, edge_types, colptr_dict, return hetero_sample<false, true, false>(
row_dict, input_node_dict, node_types, edge_types, colptr_dict, row_dict, input_node_dict,
num_neighbors_dict, num_hops); num_neighbors_dict, node_time_dict, num_hops);
} else { } else {
return hetero_sample<false, false>(node_types, edge_types, colptr_dict, return hetero_sample<false, false, false>(
row_dict, input_node_dict, node_types, edge_types, colptr_dict, row_dict, input_node_dict,
num_neighbors_dict, num_hops); num_neighbors_dict, node_time_dict, num_hops);
} }
} }
tuple<c10::Dict<node_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>,
c10::Dict<rel_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>>
hetero_temporal_neighbor_sample_cpu(
const vector<node_t> &node_types, const vector<edge_t> &edge_types,
const c10::Dict<rel_t, torch::Tensor> &colptr_dict,
const c10::Dict<rel_t, torch::Tensor> &row_dict,
const c10::Dict<node_t, torch::Tensor> &input_node_dict,
const c10::Dict<rel_t, vector<int64_t>> &num_neighbors_dict,
const c10::Dict<node_t, torch::Tensor> &node_time_dict,
const int64_t num_hops, const bool replace, const bool directed) {
AT_ASSERTM(directed, "Temporal sampling requires 'directed' sampling");
if (replace) {
// We assume that directed = True for temporal sampling
// The current implementation uses disjoint computation trees
// to tackle the case of the same node sampled having different
// root time constraint.
// In future, we could extend to directed = False case,
// allowing additional edges within each computation tree.
return hetero_sample<true, true, true>(
node_types, edge_types, colptr_dict, row_dict, input_node_dict,
num_neighbors_dict, node_time_dict, num_hops);
} else {
return hetero_sample<false, true, true>(
node_types, edge_types, colptr_dict, row_dict, input_node_dict,
num_neighbors_dict, node_time_dict, num_hops);
}
}
\ No newline at end of file
csrc/cpu/neighbor_sample_cpu.h
View file @ 1cb25232
...@@ -22,3 +22,15 @@ hetero_neighbor_sample_cpu( ...@@ -22,3 +22,15 @@ hetero_neighbor_sample_cpu(
const c10::Dict<node_t, torch::Tensor> &input_node_dict, const c10::Dict<node_t, torch::Tensor> &input_node_dict,
const c10::Dict<rel_t, std::vector<int64_t>> &num_neighbors_dict, const c10::Dict<rel_t, std::vector<int64_t>> &num_neighbors_dict,
const int64_t num_hops, const bool replace, const bool directed); const int64_t num_hops, const bool replace, const bool directed);
std::tuple<c10::Dict<node_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>,
c10::Dict<rel_t, torch::Tensor>, c10::Dict<rel_t, torch::Tensor>>
hetero_temporal_neighbor_sample_cpu(
const std::vector<node_t> &node_types,
const std::vector<edge_t> &edge_types,
const c10::Dict<rel_t, torch::Tensor> &colptr_dict,
const c10::Dict<rel_t, torch::Tensor> &row_dict,
const c10::Dict<node_t, torch::Tensor> &input_node_dict,
const c10::Dict<rel_t, std::vector<int64_t>> &num_neighbors_dict,
const c10::Dict<node_t, torch::Tensor> &node_time_dict,
const int64_t num_hops, const bool replace, const bool directed);
csrc/cpu/relabel_cpu.cpp
View file @ 1cb25232
...@@ -64,7 +64,7 @@ relabel_one_hop_cpu(torch::Tensor rowptr, torch::Tensor col, ...@@ -64,7 +64,7 @@ relabel_one_hop_cpu(torch::Tensor rowptr, torch::Tensor col,
std::unordered_map<int64_t, int64_t> n_id_map; std::unordered_map<int64_t, int64_t> n_id_map;
std::unordered_map<int64_t, int64_t>::iterator it; std::unordered_map<int64_t, int64_t>::iterator it;
auto out_rowptr = torch::empty(idx.numel() + 1, rowptr.options()); auto out_rowptr = torch::empty({idx.numel() + 1}, rowptr.options());
auto out_rowptr_data = out_rowptr.data_ptr<int64_t>(); auto out_rowptr_data = out_rowptr.data_ptr<int64_t>();
out_rowptr_data[0] = 0; out_rowptr_data[0] = 0;
...@@ -76,12 +76,12 @@ relabel_one_hop_cpu(torch::Tensor rowptr, torch::Tensor col, ...@@ -76,12 +76,12 @@ relabel_one_hop_cpu(torch::Tensor rowptr, torch::Tensor col,
out_rowptr_data[i + 1] = offset; out_rowptr_data[i + 1] = offset;
} }
auto out_col = torch::empty(offset, col.options()); auto out_col = torch::empty({offset}, col.options());
auto out_col_data = out_col.data_ptr<int64_t>(); auto out_col_data = out_col.data_ptr<int64_t>();
torch::optional<torch::Tensor> out_value = torch::nullopt; torch::optional<torch::Tensor> out_value = torch::nullopt;
if (optional_value.has_value()) { if (optional_value.has_value()) {
out_value = torch::empty(offset, optional_value.value().options()); out_value = torch::empty({offset}, optional_value.value().options());
AT_DISPATCH_ALL_TYPES(optional_value.value().scalar_type(), "relabel", [&] { AT_DISPATCH_ALL_TYPES(optional_value.value().scalar_type(), "relabel", [&] {
auto value_data = optional_value.value().data_ptr<scalar_t>(); auto value_data = optional_value.value().data_ptr<scalar_t>();
......
csrc/cpu/sample_cpu.cpp
View file @ 1cb25232
...@@ -15,13 +15,11 @@ sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx, ...@@ -15,13 +15,11 @@ sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx,
CHECK_CPU(idx); CHECK_CPU(idx);
CHECK_INPUT(idx.dim() == 1); CHECK_INPUT(idx.dim() == 1);
srand(time(NULL) + 1000 * getpid()); // Initialize random seed.
auto rowptr_data = rowptr.data_ptr<int64_t>(); auto rowptr_data = rowptr.data_ptr<int64_t>();
auto col_data = col.data_ptr<int64_t>(); auto col_data = col.data_ptr<int64_t>();
auto idx_data = idx.data_ptr<int64_t>(); auto idx_data = idx.data_ptr<int64_t>();
auto out_rowptr = torch::empty(idx.numel() + 1, rowptr.options()); auto out_rowptr = torch::empty({idx.numel() + 1}, rowptr.options());
auto out_rowptr_data = out_rowptr.data_ptr<int64_t>(); auto out_rowptr_data = out_rowptr.data_ptr<int64_t>();
out_rowptr_data[0] = 0; out_rowptr_data[0] = 0;
...@@ -69,7 +67,7 @@ sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx, ...@@ -69,7 +67,7 @@ sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx,
if (row_count > 0) { if (row_count > 0) {
for (int64_t j = 0; j < num_neighbors; j++) { for (int64_t j = 0; j < num_neighbors; j++) {
e = row_start + rand() % row_count; e = row_start + uniform_randint(row_count);
c = col_data[e]; c = col_data[e];
if (n_id_map.count(c) == 0) { if (n_id_map.count(c) == 0) {
...@@ -96,7 +94,7 @@ sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx, ...@@ -96,7 +94,7 @@ sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx,
} else { // See: https://www.nowherenearithaca.com/2013/05/ } else { // See: https://www.nowherenearithaca.com/2013/05/
// robert-floyds-tiny-and-beautiful.html // robert-floyds-tiny-and-beautiful.html
for (int64_t j = row_count - num_neighbors; j < row_count; j++) { for (int64_t j = row_count - num_neighbors; j < row_count; j++) {
if (!perm.insert(rand() % j).second) if (!perm.insert(uniform_randint(j)).second)
perm.insert(j); perm.insert(j);
} }
} }
...@@ -119,9 +117,9 @@ sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx, ...@@ -119,9 +117,9 @@ sample_adj_cpu(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx,
auto out_n_id = torch::from_blob(n_ids.data(), {N}, col.options()).clone(); auto out_n_id = torch::from_blob(n_ids.data(), {N}, col.options()).clone();
int64_t E = out_rowptr_data[idx.numel()]; int64_t E = out_rowptr_data[idx.numel()];
auto out_col = torch::empty(E, col.options()); auto out_col = torch::empty({E}, col.options());
auto out_col_data = out_col.data_ptr<int64_t>(); auto out_col_data = out_col.data_ptr<int64_t>();
auto out_e_id = torch::empty(E, col.options()); auto out_e_id = torch::empty({E}, col.options());
auto out_e_id_data = out_e_id.data_ptr<int64_t>(); auto out_e_id_data = out_e_id.data_ptr<int64_t>();
i = 0; i = 0;
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment