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