Skip to content

GitLab

Commit 19374a65 authored by rusty1s's avatar rusty1s
Browse files

new strucutre

parent 6a01b87b
Hide whitespace changes
Inline Side-by-side
Showing with 82 additions and 330 deletions
aten/TH/THGraclus.c 0 → 100644
View file @ 19374a65
#include <TH/TH.h>
#define TH_TENSOR_GRACLUS(self, row, col, PRESELECT, SELECT) { \
THLongTensor_fill(self, -1); \
int64_t *selfData = THLongTensor_data(self); \
int64_t *rowData = THLongTensor_data(row); \
int64_t *colData = THLongTensor_data(col); \
ptrdiff_t e = 0, nEdges = THLongTensor_nElement(row); \
int64_t rowValue, colValue, matchedValue, value; \
while(e < nEdges) { \
rowValue = rowData[e]; \
matchedValue = rowValue; \
PRESELECT \
if (selfData[rowValue] < 0) { \
do { \
colValue = colData[e]; \
SELECT \
e++; \
} while(e < nEdges && rowData[e] == rowValue); \
value = rowValue < matchedValue ? rowValue : matchedValue; \
selfData[rowValue] = value; \
selfData[matchedValue] = value; \
} \
while(e < nEdges && rowData[e] == rowValue) e++; \
} \
}
void THTensor_graclus(THLongTensor *self, THLongTensor *row, THLongTensor *col) {
TH_TENSOR_GRACLUS(self, row, col, {},
if (selfData[colValue] < 0) { matchedValue = colValue; break; }
)
}
#include "generic/THGraclus.c"
#include "THGenerateAllTypes.h"
aten/TH/THGraclus.h 0 → 100644
View file @ 19374a65
void THTensor_graclus(THLongTensor *self, THLongTensor *row, THLongTensor *col);
void THByteTensor_graclus(THLongTensor *self, THLongTensor *row, THLongTensor *col, THByteTensor *weight);
void THCharTensor_graclus(THLongTensor *self, THLongTensor *row, THLongTensor *col, THCharTensor *weight);
void THShortTensor_graclus(THLongTensor *self, THLongTensor *row, THLongTensor *col, THShortTensor *weight);
void THIntTensor_graclus(THLongTensor *self, THLongTensor *row, THLongTensor *col, THIntTensor *weight);
void THLongTensor_graclus(THLongTensor *self, THLongTensor *row, THLongTensor *col, THLongTensor *weight);
void THFloatTensor_graclus(THLongTensor *self, THLongTensor *row, THLongTensor *col, THFloatTensor *weight);
void THDoubleTensor_graclus(THLongTensor *self, THLongTensor *row, THLongTensor *col, THDoubleTensor *weight);
aten/TH/THGreedy.c deleted 100644 → 0
View file @ 6a01b87b
#include <TH/TH.h>
#include "common.h"
#define THGreedy_ TH_CONCAT_3(TH,Real,Greedy)
#define TH_GREEDY_CLUSTER(cluster, row, col, SELECT) { \
THLongTensor_fill(cluster, -1); \
int64_t *clusterData = THTensor_getData(cluster); \
int64_t *rowData = THTensor_getData(row); \
int64_t *colData = THTensor_getData(col); \
ptrdiff_t idx = 0, nEdges = THLongTensor_nElement(row); \
int64_t rowValue, pairValue, colValue, clusterValue; \
while(idx < nEdges) { \
rowValue = rowData[idx]; \
pairValue = rowValue; \
if (clusterData[rowValue] < 0) { \
while(idx < nEdges && rowData[idx] == rowValue) { \
colValue = colData[idx]; \
if (clusterData[colValue] < 0) { \
SELECT \
} \
idx++; \
} \
} \
clusterValue = rowValue < pairValue ? rowValue : pairValue; \
clusterData[rowValue] = clusterValue; \
clusterData[pairValue] = clusterValue; \
while(idx < nEdges && rowData[idx] == rowValue) idx++; \
} \
}
void THGreedy(THLongTensor *cluster, THLongTensor *row, THLongTensor *col) {
TH_GREEDY_CLUSTER(cluster, row, col,
pairValue = colValue;
break;
)
}
#include "generic/THGreedy.c"
#include "THGenerateAllTypes.h"
aten/TH/THGreedy.h deleted 100644 → 0
View file @ 6a01b87b
void THGreedy(THLongTensor *cluster, THLongTensor *row, THLongTensor *col);
void THByteGreedy(THLongTensor *cluster, THLongTensor *row, THLongTensor *col, THByteTensor *weight);
void THCharGreedy(THLongTensor *cluster, THLongTensor *row, THLongTensor *col, THCharTensor *weight);
void THShortGreedy(THLongTensor *cluster, THLongTensor *row, THLongTensor *col, THShortTensor *weight);
void THIntGreedy(THLongTensor *cluster, THLongTensor *row, THLongTensor *col, THIntTensor *weight);
void THLongGreedy(THLongTensor *cluster, THLongTensor *row, THLongTensor *col, THLongTensor *weight);
void THFloatGreedy(THLongTensor *cluster, THLongTensor *row, THLongTensor *col, THFloatTensor *weight);
void THDoubleGreedy(THLongTensor *cluster, THLongTensor *row, THLongTensor *col, THDoubleTensor *weight);
aten/TH/THGrid.c
View file @ 19374a65
#include <TH/TH.h>
#include "common.h"
#define THGrid_ TH_CONCAT_3(TH,Real,Grid)
#include "generic/THGrid.c"
#include "THGenerateAllTypes.h"
aten/TH/THGrid.h
View file @ 19374a65
void THByteGrid(THLongTensor *cluster, THByteTensor *pos, THByteTensor *size, THLongTensor *count);
void THCharGrid(THLongTensor *cluster, THCharTensor *pos, THCharTensor *size, THLongTensor *count);
void THShortGrid(THLongTensor *cluster, THShortTensor *pos, THShortTensor *size, THLongTensor *count);
void THIntGrid(THLongTensor *cluster, THIntTensor *pos, THIntTensor *size, THLongTensor *count);
void THLongGrid(THLongTensor *cluster, THLongTensor *pos, THLongTensor *size, THLongTensor *count);
void THFloatGrid(THLongTensor *cluster, THFloatTensor *pos, THFloatTensor *size, THLongTensor *count);
void THDoubleGrid(THLongTensor *cluster, THDoubleTensor *pos, THDoubleTensor *size, THLongTensor *count);
void THByteTensor_grid(THLongTensor *self, THByteTensor *pos, THByteTensor *size, THLongTensor *count);
void THCharTensor_grid(THLongTensor *self, THCharTensor *pos, THCharTensor *size, THLongTensor *count);
void THShortTensor_grid(THLongTensor *self, THShortTensor *pos, THShortTensor *size, THLongTensor *count);
void THIntTensor_grid(THLongTensor *self, THIntTensor *pos, THIntTensor *size, THLongTensor *count);
void THLongTensor_grid(THLongTensor *self, THLongTensor *pos, THLongTensor *size, THLongTensor *count);
void THFloatTensor_grid(THLongTensor *self, THFloatTensor *pos, THFloatTensor *size, THLongTensor *count);
void THDoubleTensor_grid(THLongTensor *self, THDoubleTensor *pos, THDoubleTensor *size, THLongTensor *count);
aten/TH/common.h deleted 100644 → 0
View file @ 6a01b87b
#ifndef TH_COMMON_INC
#define TH_COMMON_INC
#define THTensor_getData(TENSOR) TENSOR->storage->data + TENSOR->storageOffset
#endif // TH_COMMON_INC
aten/TH/generic/THGraclus.c 0 → 100644
View file @ 19374a65
#ifndef TH_GENERIC_FILE
#define TH_GENERIC_FILE "generic/THGraclus.c"
#else
void THTensor_(graclus)(THLongTensor *self, THLongTensor *row, THLongTensor *col, THTensor *weight) {
real *weightData = THTensor_(data)(weight);
real maxWeight, tmp;
TH_TENSOR_GRACLUS(self, row, col, maxWeight = 0;,
tmp = weightData[e];
if (selfData[colValue] < 0 && tmp > maxWeight) { matchedValue = colValue; maxWeight = tmp; }
)
}
#endif // TH_GENERIC_FILE
aten/TH/generic/THGreedy.c deleted 100644 → 0
View file @ 6a01b87b
#ifndef TH_GENERIC_FILE
#define TH_GENERIC_FILE "generic/THGreedy.c"
#else
void THGreedy_(THLongTensor *cluster, THLongTensor *row, THLongTensor *col, THTensor *weight) {
real *weightData = THTensor_getData(weight);
real maxWeight = 0, tmpWeight;
TH_GREEDY_CLUSTER(cluster, row, col,
tmpWeight = weightData[idx];
if (tmpWeight > maxWeight) {
pairValue = colValue;
maxWeight = tmpWeight;
}
)
}
#endif // TH_GENERIC_FILE
aten/TH/generic/THGrid.c
View file @ 19374a65
......@@ -2,11 +2,11 @@
#define TH_GENERIC_FILE "generic/THGrid.c"
#else
void THGrid_(THLongTensor *cluster, THTensor *pos, THTensor *size, THLongTensor *count) {
int64_t *clusterData = THTensor_getData(cluster);
real *posData = THTensor_getData(pos);
real *sizeData = THTensor_getData(size);
int64_t *countData = THTensor_getData(count);
void THTensor_(grid)(THLongTensor *self, THTensor *pos, THTensor *size, THLongTensor *count) {
int64_t *selfData = THLongTensor_data(self);
real *posData = THTensor_(data)(pos);
real *sizeData = THTensor_(data)(size);
int64_t *countData = THLongTensor_data(count);
ptrdiff_t n, d; int64_t coef, value;
for (n = 0; n < THTensor_(size)(pos, 0); n++) {
......@@ -16,7 +16,7 @@ void THGrid_(THLongTensor *cluster, THTensor *pos, THTensor *size, THLongTensor
coef *= countData[d];
}
posData += pos->stride[0];
clusterData[n] = value;
selfData[n] = value;
}
}
......
build.py
View file @ 19374a65
......@@ -8,7 +8,7 @@ from torch.utils.ffi import create_extension
if osp.exists('build'):
shutil.rmtree('build')
files = ['Greedy', 'Grid']
files = ['Graclus', 'Grid']
headers = ['aten/TH/TH{}.h'.format(f) for f in files]
sources = ['aten/TH/TH{}.c'.format(f) for f in files]
......
test/dense_grid.json deleted 100644 → 0
View file @ 6a01b87b
[
{
"name": "One-dimensional positions without start/end",
"position": [2, 16],
"size": [5],
"expected": [0, 2],
"expected_C": 3
},
{
"name": "Start parameter",
"position": [2, 16],
"size": [5],
"start": [-1],
"expected": [0, 3],
"expected_C": 4
},
{
"name": "End parameter",
"position": [2, 16],
"size": [5],
"start": [0],
"end": [30],
"expected": [0, 3],
"expected_C": 6
},
{
"name": "Two-dimensional positions",
"position": [[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]],
"size": [5, 5],
"expected": [0, 5, 1, 0, 2],
"expected_C": 6
},
{
"name": "Batch",
"position": [[0, 0], [11, 9], [2, 8], [2, 2], [8, 3], [1, 1], [6, 6]],
"size": [5, 5],
"batch": [0, 0, 0, 0, 0, 1, 1],
"expected": [0, 5, 1, 0, 2, 6, 9],
"expected_C": 12
},
{
"name": "Batch with start/end parameter",
"position": [[0, 0], [11, 9], [2, 8], [2, 2], [8, 3], [1, 1], [6, 6]],
"size": [5, 5],
"batch": [0, 0, 0, 0, 0, 1, 1],
"start": [0, 0],
"end": [20, 20],
"expected": [0, 9, 1, 0, 4, 16, 21],
"expected_C": 32
}
]
test/sparse_grid.json deleted 100644 → 0
View file @ 6a01b87b
[
{
"name": "One-dimensional positions",
"position": [2, 16],
"size": [5],
"expected": [0, 1]
},
{
"name": "Start parameter",
"position": [2, 6],
"size": [5],
"start": [0],
"expected": [0, 1]
},
{
"name": "Consecutive cluster",
"position": [0, 17, 2, 8, 3],
"size": [5],
"expected": [0, 2, 0, 1, 0]
},
{
"name": "Two-dimensional positions",
"position": [[0, 0], [11, 9], [2, 8], [2, 2], [8, 3]],
"size": [5, 5],
"expected": [0, 3, 1, 0, 2]
},
{
"name": "Batch",
"position": [[0, 0], [11, 9], [2, 8], [2, 2], [8, 3], [1, 1], [6, 6]],
"size": [5, 5],
"batch": [0, 0, 0, 0, 0, 1, 1],
"expected": [0, 3, 1, 0, 2, 4, 5],
"expected_batch": [0, 0, 0, 0, 1, 1]
},
{
"name": "Position tensor",
"position": [[[0, 0], [9, 9]], [[0, 0], [9, 0]]],
"size": [5, 5],
"expected": [[0, 2], [0, 1]]
}
]
test/test_dense_grid.py deleted 100644 → 0
View file @ 6a01b87b
from os import path as osp
from itertools import product
import pytest
import json
import torch
from torch_cluster import dense_grid_cluster
from .utils import tensors, Tensor
f = open(osp.join(osp.dirname(__file__), 'dense_grid.json'), 'r')
data = json.load(f)
f.close()
@pytest.mark.parametrize('tensor,i', product(tensors, range(len(data))))
def test_dense_grid_cluster_cpu(tensor, i):
position = Tensor(tensor, data[i]['position'])
size = torch.LongTensor(data[i]['size'])
batch = data[i].get('batch')
batch = None if batch is None else torch.LongTensor(batch)
start = data[i].get('start')
start = None if start is None else torch.LongTensor(start)
end = data[i].get('end')
end = None if end is None else torch.LongTensor(end)
expected = torch.LongTensor(data[i]['expected'])
expected_C = data[i]['expected_C']
output = dense_grid_cluster(position, size, batch, start, end)
assert output[0].tolist() == expected.tolist()
assert output[1] == expected_C
@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
@pytest.mark.parametrize('tensor,i', product(tensors, range(len(data))))
def test_dense_grid_cluster_gpu(tensor, i): # pragma: no cover
position = Tensor(tensor, data[i]['position']).cuda()
size = torch.cuda.LongTensor(data[i]['size'])
batch = data[i].get('batch')
batch = None if batch is None else torch.cuda.LongTensor(batch)
start = data[i].get('start')
start = None if start is None else torch.cuda.LongTensor(start)
end = data[i].get('end')
end = None if end is None else torch.cuda.LongTensor(end)
expected = torch.LongTensor(data[i]['expected'])
expected_C = data[i]['expected_C']
output = dense_grid_cluster(position, size, batch, start, end)
assert output[0].cpu().tolist() == expected.tolist()
assert output[1] == expected_C
test/test_graclus.py 0 → 100644
View file @ 19374a65
import torch
from torch_cluster import graclus_cluster
def test_graclus_cluster_cpu():
row = torch.LongTensor([0, 0, 1, 1, 1, 2, 2, 2, 3, 3])
col = torch.LongTensor([1, 2, 0, 2, 3, 0, 1, 3, 1, 2])
weight_ = torch.Tensor([1, 2, 1, 3, 2, 2, 3, 1, 2, 1])
cluster = graclus_cluster(row, col)
cluster = graclus_cluster(row, col, weight_)
test/test_serial.py deleted 100644 → 0
View file @ 6a01b87b
# import torch
# from torch_cluster import serial_cluster
def test_serial():
pass
# ed_index = torch.LongTensor([[0, 0, 0, 1, 2, 3, 3, 3, 4, 5, 5, 5, 6, 6],
# [2, 3, 6, 5, 0, 0, 4, 5, 3, 1, 3, 6, 0, 3]])
# edge_attr = torch.Tensor([2, 2, 2, 1, 2, 2, 2, 2, 2, 1, 2, 2, 2, 2])
# rid = torch.arange(edge_index.max() + 1, out=edge_index.new())
# output = random_cluster(edge_index, rid=rid, perm_edges=False)
# expected_output = [0, 1, 2, 0, 3, 1, 4]
# assert output.tolist() == expected_output
# TODO: Test only if conditions are met:
# * at most two pairs with the same cluster
# * pairs need to be neighbors of each other
# TODO: Rename to serial
test/test_sparse_grid.py deleted 100644 → 0
View file @ 6a01b87b
from os import path as osp
from itertools import product
import pytest
import json
import torch
from torch_cluster import sparse_grid_cluster
from .utils import tensors, Tensor
f = open(osp.join(osp.dirname(__file__), 'sparse_grid.json'), 'r')
data = json.load(f)
f.close()
@pytest.mark.parametrize('tensor,i', product(tensors, range(len(data))))
def test_sparse_grid_cluster_cpu(tensor, i):
position = Tensor(tensor, data[i]['position'])
size = torch.LongTensor(data[i]['size'])
batch = data[i].get('batch')
batch = None if batch is None else torch.LongTensor(batch)
start = data[i].get('start')
start = None if start is None else torch.LongTensor(start)
expected = torch.LongTensor(data[i]['expected'])
output = sparse_grid_cluster(position, size, batch, start)
if batch is None:
assert output.tolist() == expected.tolist()
else:
# expected_batch = torch.LongTensor(data[i]['expected_batch'])
assert output[0].tolist() == expected.tolist()
# assert output[1].tolist() == expected_batch.tolist()
@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
@pytest.mark.parametrize('tensor,i', product(tensors, range(len(data))))
def test_sparse_grid_cluster_gpu(tensor, i): # pragma: no cover
position = Tensor(tensor, data[i]['position']).cuda()
size = torch.cuda.LongTensor(data[i]['size'])
batch = data[i].get('batch')
batch = None if batch is None else torch.cuda.LongTensor(batch)
start = data[i].get('start')
start = None if start is None else torch.cuda.LongTensor(start)
expected = torch.LongTensor(data[i]['expected'])
output = sparse_grid_cluster(position, size, batch, start)
if batch is None:
assert output.cpu().tolist() == expected.tolist()
else:
# expected_batch = torch.LongTensor(data[i]['expected_batch'])
assert output[0].cpu().tolist() == expected.tolist()
# assert output[1].cpu().tolist() == expected_batch.tolist()
test/utils.py deleted 100644 → 0
View file @ 6a01b87b
import torch
from torch._tensor_docs import tensor_classes
tensors = [t[:-4] for t in tensor_classes]
def Tensor(str, x):
tensor = getattr(torch, str)
return tensor(x)
test/utils/test_consecutive.py deleted 100644 → 0
View file @ 6a01b87b
import pytest
import torch
from torch_cluster.functions.utils.consecutive import consecutive
def test_consecutive_cpu():
x = torch.LongTensor([[0, 3, 2], [2, 3, 0]])
assert consecutive(x).tolist() == [[0, 2, 1], [1, 2, 0]]
@pytest.mark.skipif(not torch.cuda.is_available(), reason='no CUDA')
def test_consecutive_gpu(): # pragma: no cover
x = torch.cuda.LongTensor([[0, 3, 2], [2, 3, 0]])
assert consecutive(x).cpu().tolist() == [[0, 2, 1], [1, 2, 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