def compute_neighbours_and_distancesq(self, coordinates, row_splits):
idx, dist = SelectKnn(self.n_neighbours,
coordinates,
row_splits,
max_radius=-1.0,
tf_compatible=False)
return idx, dist
def selectNeighbours_CUDA(K, coords, row_splits, return_distances=False):
with tf.GradientTape(persistent=True,
watch_accessed_variables=True) as t_newop:
t_newop.watch(coords)
out = SelectKnn(K=K,
coords=coords,
row_splits=row_splits,
max_radius=-1.,
tf_compatible=True)
return out, t_newop
def applyClustering(K, incoords, hier, row_splits):
neighs, _ = SelectKnn(K=K,
coords=incoords,
row_splits=row_splits,
tf_compatible=False,
max_radius=radius)
sel, rs, gscatidx = LocalClustering()([neighs, hier, row_splits])
return rs, sel, gscatidx
def compute_neighbours_and_distancesq(self, coordinates, row_splits,
masking_values):
idx, dist = SelectKnn(self.n_neighbours,
coordinates,
row_splits,
max_radius=1.0,
tf_compatible=False,
masking_values=masking_values,
threshold=self.threshold,
mask_mode=self.direction,
mask_logic=self.ex_mode)
return idx, dist
def applyClustering(K, incoords, hier, row_splits):
hierarchy_idxs = []
for i in range(len(row_splits.numpy()) - 1):
a = tf.argsort(hier[row_splits[i]:row_splits[i + 1]], axis=0)
hierarchy_idxs.append(a + row_splits[i])
hierarchy_idxs = tf.concat(hierarchy_idxs, axis=0)
print('hierarchy_idxs', hierarchy_idxs.shape)
print('incoords', incoords.shape)
assert incoords.shape[0] == hierarchy_idxs.shape[0]
neighs, _ = SelectKnn(K=K,
coords=incoords,
row_splits=row_splits,
tf_compatible=False,
max_radius=radius)
rs, sel, gscatidx = lc(neighs, hierarchy_idxs, row_splits)
return rs, sel, gscatidx
def compute_neighbours_and_distancesq(self, coordinates, row_splits):
ragged_split_added_indices, _ = SelectKnn(self.n_neighbours,
coordinates,
row_splits,
max_radius=1.0,
tf_compatible=True)
#ragged_split_added_indices = ragged_split_added_indices[:,1:]
ragged_split_added_indices = ragged_split_added_indices[...,
tf.newaxis]
distancesq = tf.reduce_sum(
(coordinates[:, tf.newaxis, :] -
tf.gather_nd(coordinates, ragged_split_added_indices))**2,
axis=-1) # [SV, N]
return ragged_split_added_indices, distancesq
radius = 0.05
def getHierIdx(hier, row_splits):
hierarchy_idxs = []
for i in range(len(row_splits.numpy()) - 1):
a = tf.argsort(hier[row_splits[i]:row_splits[i + 1]], axis=0)
hierarchy_idxs.append(a + row_splits[i])
return tf.concat(hierarchy_idxs, axis=0)
for v in [10000, 50000, 100000, 400000]:
row_splits, hier, coords, truth, glidxs = createData(v)
neighs, _ = SelectKnn(K=K,
coords=coords,
row_splits=row_splits,
tf_compatible=False,
max_radius=radius)
#warmup
hierarchy_idxs = getHierIdx(hier, row_splits)
start = time.time()
hierarchy_idxs = getHierIdx(hier, row_splits)
sortt.append(time.time() - start)
#warm up
rs, sel, gscatidx = lc(neighs, hierarchy_idxs, row_splits)
start = time.time()
for _ in range(10):
rs, sel, gscatidx = lc(neighs, hierarchy_idxs, row_splits)
totime = (time.time() - start) / 10.
if return_distances:
return (idcs, distances), t_newop
return tf.concat(out_indices, axis=0), t_newop
#visual inspection
coords, row_splits = createData(20, 2)
masking_values = tf.random.uniform((20, 1), dtype='float32', seed=1)
idx, dist = SelectKnn(K=5,
coords=coords,
row_splits=row_splits,
max_radius=-.3,
tf_compatible=False,
masking_values=masking_values,
threshold=0.5,
mask_mode='scat',
mask_logic='and')
print('masking_values', masking_values)
print('idx', idx)
print('dist', dist)
exit()
reldiff = []
for it in range(1):
coords, row_splits = createData(4000, 4)
K = 200
from neighbour_covariance_op import NeighbourCovariance
import numpy as np
import tensorflow as tf
from select_knn_op import SelectKnn
import time
n_vert = 300000
n_coords = 3
n_feats = 2
n_neigh = 64
coords = tf.random.uniform((n_vert, n_coords), dtype='float32', seed=2)
feats = tf.random.uniform((n_vert, n_feats), dtype='float32', seed=2)
row_splits = tf.constant([0, n_vert], dtype='int32')
nidx, _ = SelectKnn(n_neigh, coords, row_splits)
print('launching op')
cov, means = NeighbourCovariance(coords, feats, nidx)
t0 = time.time()
for _ in range(20):
cov, means = NeighbourCovariance(coords, feats, nidx)
print((time.time() - t0) / 20)
print(cov.shape, cov)
print(means.shape)
def benchmark(self,
nvert=30000,
nfeat=64,
nneigh=128,
ncoords=4,
dogradient=False,
do_tf=True):
coords = tf.constant(np.random.rand(nvert, ncoords), dtype='float32')
feats = tf.constant(np.random.rand(nvert, nfeat), dtype='float32')
row_splits = tf.constant([0, nvert], dtype='int32')
indices, distances = SelectKnn(K=nneigh,
coords=coords,
row_splits=row_splits)
if self.use_distances_direct:
coords = distances
tf_failed = False
if not dogradient:
#each gets one dry run to compile
meanmax = self.customimpl(coords,
features=feats,
indices=indices,
mean_and_max=self.mean_and_max)
t0 = time.time()
for i in range(0, 50):
meanmax = self.customimpl(coords,
features=feats,
indices=indices,
mean_and_max=self.mean_and_max)
op_time = (time.time() - t0) / 50.
print('op_time', op_time)
tf_time = 0
if do_tf:
try:
meanmax = self.tfimp(coords,
features=feats,
indices=indices,
mean_and_max=self.mean_and_max)
t0 = time.time()
for i in range(0, 50):
meanmax = self.tfimp(coords,
features=feats,
indices=indices,
mean_and_max=self.mean_and_max)
tf_time = (time.time() - t0) / 50.
except:
tf_failed = True
print('tf_time', tf_time)
return op_time, tf_time
else:
with tf.GradientTape(persistent=True,
watch_accessed_variables=True) as t_newop:
t_newop.watch(coords)
t_newop.watch(feats)
meanmax = self.customimpl(coords,
features=feats,
indices=indices,
mean_and_max=self.mean_and_max)
#once to get it compiled in case needed
feat_grad = t_newop.gradient(meanmax, feats)
coord_grad = t_newop.gradient(meanmax, coords)
t0 = time.time()
for i in range(5):
feat_grad = t_newop.gradient(meanmax, feats)
coord_grad = t_newop.gradient(meanmax, coords)
op_time = (time.time() - t0) / 5.
tf_time = 0
if do_tf:
try:
with tf.GradientTape(persistent=True) as t_tfop:
t_tfop.watch(coords)
t_tfop.watch(feats)
meanmax = self.tfimp(coords,
features=feats,
indices=indices,
mean_and_max=self.mean_and_max)
feat_grad = t_tfop.gradient(meanmax, feats)
coord_grad = t_tfop.gradient(meanmax, coords)
t0 = time.time()
for i in range(5):
feat_grad = t_tfop.gradient(meanmax, feats)
coord_grad = t_tfop.gradient(meanmax, coords)
tf_time = (time.time() - t0) / 5.
except:
tf_failed = True
return op_time, tf_time
def difference(self,
nvert=300,
nfeat=64,
nneigh=32,
ncoords=4,
onlyForward=False,
assert_error=True):
coords = tf.constant(np.random.rand(nvert, ncoords), dtype='float32')
feats = np.random.rand(nvert, nfeat)
#to make the max unambiguous
frange = np.arange(nvert)
np.random.shuffle(frange)
toadd = np.expand_dims(frange, axis=1)
feats = tf.constant(feats + toadd, dtype='float32')
row_splits = tf.constant([0, nvert], dtype='int32')
#print('building indices')
with tf.device("/cpu:0"):
indices, distances = SelectKnn(K=nneigh,
coords=coords,
row_splits=row_splits)
#indices = indices[:,1:]
#distances = distances[:,1:]
#print('process custom op')
if self.use_distances_direct:
coords = distances
op_time = 0
tfdevstring = "/gpu:0"
if self.customoncpu:
tfdevstring = "/cpu:0"
tfdev = tf.device(tfdevstring)
t0 = time.time()
with tfdev:
t0 = time.time()
with tf.GradientTape(persistent=True,
watch_accessed_variables=True) as t_newop:
t_newop.watch(coords)
t_newop.watch(feats)
meanmax = self.customimpl(coords,
features=feats,
indices=indices,
mean_and_max=self.mean_and_max)
t1 = time.time()
op_time = t1 - t0
#print('op time',op_time)
with tfdev:
coord_grad = t_newop.gradient(meanmax, coords)
feat_grad = t_newop.gradient(meanmax, feats)
if self.debugout:
print('coords', coords, '\n')
print('feats', feats, '\n')
print('custom output', meanmax, '\n')
print('indices', indices)
### tf op implementation
print('TFTFTF')
tf_feat_grad = None
tf_coord_grad = None
#print('process TF op')
tfdevstring = "/gpu:0"
if self.tfoncpu:
tfdevstring = "/cpu:0"
tfdev = tf.device(tfdevstring)
t0 = time.time()
with tfdev:
with tf.GradientTape(persistent=True) as t_tfop:
t_tfop.watch(coords)
t_tfop.watch(feats)
tf_meanmax = self.tfimp(coords,
features=feats,
indices=indices,
mean_and_max=self.mean_and_max)
tf_time = time.time() - t0
if self.debugout:
print('TF output', tf_meanmax, '\n')
with tfdev:
tf_feat_grad = t_tfop.gradient(tf_meanmax, feats)
tf_coord_grad = t_tfop.gradient(tf_meanmax, coords)
with tf.device("/cpu:0"):
difference = meanmax - tf_meanmax
max_rel_difference = tf.reduce_max(
tf.abs(difference / (tf.abs(tf_meanmax) + 1e-3))).numpy()
max_difference = tf.reduce_max(tf.abs(difference)).numpy()
#print('max rel difference',max_rel_difference)
#print('max difference',max_difference)
#print('op time',op_time)
#print('tf time',tf_time)
if assert_error:
assert max_difference < 1e-2
if onlyForward:
return
## gradients
#print('tf_feat_grad',tf_feat_grad)
#print('tf_coord_grad',tf_coord_grad)
#print('feat_grad',feat_grad)
#print('coord_grad',coord_grad)
feat_grad_diff = feat_grad - tf_feat_grad
coord_grad_diff = coord_grad - tf_coord_grad
#print('feat_grad_diff',feat_grad_diff)
#print('coord_grad_diff',coord_grad_diff)
#print('relative feat_grad_diff',feat_grad_diff/tf_feat_grad)
#print('relative coord_grad_diff',coord_grad_diff/tf_coord_grad)
maxfeatgraddiff = tf.reduce_max(tf.abs(feat_grad_diff))
maxcoordgraddiff = tf.reduce_max(tf.abs(coord_grad_diff))
rel_feat_grad_diff = (feat_grad_diff) / (tf.abs(tf_feat_grad) +
1e-2)
rel_coord_grad_diff = coord_grad_diff / (tf.abs(tf_coord_grad) +
1e-2)
maxrelfeatgraddiff = tf.reduce_max(tf.abs(rel_feat_grad_diff))
maxrelcoordgraddiff = tf.reduce_max(tf.abs(rel_coord_grad_diff))
#print('\nmax relative feature grad diff', maxrelfeatgraddiff)
#print('max relative coordinate grad diff', maxrelcoordgraddiff)
def check_indices():
idx_ok = True
for i in tf.range(indices.shape[0]):
y, idx, c = tf.unique_with_counts(indices[i])
if (c.numpy() > 1).any() or (indices[i].numpy() >=
indices.shape[0]).any():
idx_ok = False
print("indices not unique", indices[i])
if idx_ok:
print('indices ok')
if self.debugout:
print('custom feature grad ', feat_grad)
print('TF feature grad', tf_feat_grad)
print('difference', feat_grad_diff)
print('custom coord grad', coord_grad)
print('TF coord grad', tf_coord_grad)
print('Difference', coord_grad_diff)
if maxrelfeatgraddiff > 1e-2:
print('Feature gradient off:')
print('max rel diff', maxrelfeatgraddiff)
print('max diff', maxfeatgraddiff)
print('min,max feat', tf.reduce_min(feats),
tf.reduce_max(feats))
print('min,max coords', tf.reduce_min(coords),
tf.reduce_max(coords))
check_indices()
if maxrelcoordgraddiff > 1e-2:
print('Coordinate gradient off:')
print('max rel diff', maxrelcoordgraddiff)
print('max diff', maxcoordgraddiff)
print('min,max feat', tf.reduce_min(feats),
tf.reduce_max(feats))
print('min,max coords', tf.reduce_min(coords),
tf.reduce_max(coords))
check_indices()
if maxfeatgraddiff > 1e-2:
print('Feature gradient off:')
print('max rel diff', maxrelfeatgraddiff)
print('max diff', maxfeatgraddiff)
print('min,max feat', tf.reduce_min(feats),
tf.reduce_max(feats))
print('min,max coords', tf.reduce_min(coords),
tf.reduce_max(coords))
check_indices()
if maxcoordgraddiff > 1e-2:
print('Coordinate gradient off:')
print('max rel diff', maxrelcoordgraddiff)
print('max diff', maxcoordgraddiff)
print('min,max feat', tf.reduce_min(feats),
tf.reduce_max(feats))
print('min,max coords', tf.reduce_min(coords),
tf.reduce_max(coords))
check_indices()
if assert_error:
assert maxrelfeatgraddiff < 5e-2
assert maxrelcoordgraddiff < 5e-2
reldifference = tf.reshape(
difference / (tf.abs(tf_meanmax) + 1e-4), [-1])
difference = tf.reshape(difference, [-1])
rel_feat_grad_diff = tf.reshape(rel_feat_grad_diff, [-1])
rel_coord_grad_diff = tf.reshape(rel_coord_grad_diff, [-1])
feat_grad_diff = tf.reshape(feat_grad_diff, [-1])
coord_grad_diff = tf.reshape(coord_grad_diff, [-1])
return difference, reldifference, rel_feat_grad_diff, rel_coord_grad_diff, feat_grad_diff, coord_grad_diff
def createData(n_vert, n_coords):
coords = tf.random.uniform((n_vert, n_coords), dtype='float32', seed=2)
row_splits = tf.constant([0, n_vert], dtype='int32')
return coords, row_splits #row splits just for other ops, no need to test
n_neigh = 5
max_radius = -1
coords, row_splits = createData(10, 3)
neigh, dst = SelectKnn(n_neigh,
coords,
row_splits,
max_radius=max_radius,
tf_compatible=True)
dst = None
dst2 = None
with tf.GradientTape(persistent=True,
watch_accessed_variables=True) as t_newop:
t_newop.watch(coords)
dst2 = LocalDistance(coords, neigh)
coord_grad = t_newop.gradient(dst2, coords)
with tf.GradientTape(persistent=True,
watch_accessed_variables=True) as t_compareop:
t_compareop.watch(coords)
_, dst = SelectKnn(n_neigh,
coords,
def call(self, inputs):
x = inputs[0]
row_splits = inputs[1]
thresh_values = inputs[2]
coordinates = self.input_spatial_transform(x)
up_neighbour_indices, up_distancesq = SelectKnn(
self.n_neighbours,
coordinates,
row_splits,
max_radius=self.max_radius,
masking_values=thresh_values,
tf_compatible=False,
threshold=self.threshold,
mask_mode='acc',
mask_logic='xor')
down_neighbour_indices, down_distancesq = SelectKnn(
self.n_neighbours,
coordinates,
row_splits,
max_radius=self.max_radius,
masking_values=thresh_values,
tf_compatible=False,
threshold=self.threshold,
mask_mode='scat',
mask_logic='xor')
allfeat = []
tfet = []
features = x
#tf.print(down_neighbour_indices)
for i in range(len(self.input_feature_transform)):
t = self.input_feature_transform[i]
neighbour_indices = up_neighbour_indices
distancesq = up_distancesq
if i % 2:
neighbour_indices = down_neighbour_indices
distancesq = down_distancesq
features = t(features)
orig_tf = features
prev_shape = features.shape
features = self.collect_neighbours(features, neighbour_indices,
distancesq)
features = tf.reshape(features, [-1, prev_shape[1] * 2])
#tf.print(i, tf.concat([thresh_values,features,orig_tf],axis=-1), summarize=10)
allfeat.append(features)
if self.parse_lin and i > 0:
up_distancesq *= 0. #the remaining parsing operations will be at zero distance
down_distancesq *= 0.
#tf.print(thresh_values.shape, allfeat[0].shape)
#tf.print(tf.concat([thresh_values]+allfeat+tfet,axis=-1), summarize=120)
features = tf.concat(allfeat + [x], axis=-1)
return self.output_feature_transform(features), coordinates