def vis(): import numpy as np import trimesh tf.compat.v1.enable_eager_execution() split = 'train' batch_size = 2 problem = PartnetProblem() with problem: dataset = problem.get_base_dataset(split).map(pre_batch_map, -1) dataset = dataset.batch(batch_size) dataset = dataset.map( functools.partial(post_batch_map, return_coords=True)) for features, labels, weights in dataset: del labels, weights def vis_clouds(in_coords, out_coords, target_index, neighbors): out_colors = np.ones((out_coords.shape[0], 3), dtype=np.uint8) out_colors *= 128 out_colors[target_index] = [0, 255, 0] out_cloud = trimesh.PointCloud(out_coords, out_colors) in_colors = np.zeros((in_coords.shape[0], 3), dtype=np.uint8) in_colors[:, 2] = 255 neigh_indices = neighbors[neighbors[:, 0] == target_index, 1] in_colors[neigh_indices] = [255, 0, 0] in_cloud = trimesh.PointCloud(in_coords, in_colors) scene = in_cloud.scene() scene.add_geometry(out_cloud) scene.show(background=(0, 0, 0)) features = tf.nest.map_structure(lambda x: x.numpy(), features) all_coords = features['all_coords'] # sample_indices = features['sample_indices'] in_place_indices = features['in_place_indices'] in_place_rel_coords = features['in_place_rel_coords'] down_sample_indices = features['down_sample_indices'] down_sample_rel_coords = features['down_sample_rel_coords'] row_splits = features['row_splits'] depth = len(all_coords) for i in range(depth - 1): # down sample print(np.max(np.linalg.norm(down_sample_rel_coords[i], axis=0))) print(np.max(np.linalg.norm(in_place_rel_coords[i], axis=0))) vis_clouds(all_coords[i], all_coords[i + 1], row_splits[i + 1][-2], down_sample_indices[i]) # in place vis_clouds(all_coords[i + 1], all_coords[i + 1], row_splits[i + 1][-2], in_place_indices[i])
def run_explicit(): tf.compat.v1.enable_eager_execution() tf.config.experimental_run_functions_eagerly(True) split = 'train' batch_size = 16 problem = PartnetProblem() with problem: dataset = problem.get_base_dataset(split).map(pre_batch_map, -1) dataset = dataset.batch(batch_size) for ex in dataset: post_batch_map(*ex) break print('Finished single run of explicit post_batch_map')
def run_benchmark(): split = 'train' batch_size = 16 num_warmup = 5 num_batches = 10 problem = PartnetProblem() with problem: dataset = problem.get_base_dataset(split).map(pre_batch_map, -1) dataset = dataset.batch(batch_size) dataset = dataset.map(post_batch_map, -1).prefetch(-1) for i, _ in enumerate( tqdm.tqdm(dataset.take(num_warmup + num_batches), total=num_warmup + num_batches, desc='benchmarking')): if i == num_warmup: t = time() if i == num_warmup + num_batches - 1: dt = time() - t print('{} batches in {} s: {} ms / batch'.format( num_batches, dt, dt * 1000 / num_batches))
@gin.configurable(blacklist=['input_spec', 'output_spec']) def res_conv_semantic_segmenter(input_spec, output_spec, dense_factory=mk_layers.Dense, batch_norm_impl=mk_layers.BatchNormalization, activation='relu', filters0=32): pass if __name__ == '__main__': import tensorflow_datasets as tfds from deep_cloud.problems.partnet import PartnetProblem tf.compat.v1.enable_v2_tensorshape() problem = PartnetProblem() dataset = problem.get_base_dataset(split='validation') def test_eager(): for positions, _ in tfds.as_numpy(dataset): compute_edges_eager(positions) break print('test_pre_batch_map_fn passed') def test_pre_batch_map_fn(): ds = dataset.map(pre_batch_map) for _ in tfds.as_numpy(ds): break print('test_pre_batch_map_fn passed') def test_pipeline():
def run_full_benchmark(): from deep_cloud.ops.np_utils.tree_utils import pykd from deep_cloud.ops.np_utils.tree_utils import core from deep_cloud.problems.partnet import PartnetProblem from tqdm import tqdm from time import time depth = 6 num_warmup = 5 num_iterations = 10 k0 = 16 tree_impl = pykd.KDTree SQRT_2 = np.sqrt(2.) context = PyFuncContext() problem = PartnetProblem() def f(coords, labels): tree = context.py_function(tree_impl, None, coords) dists, indices = context.py_function( lambda tree: tree.query(tree.data, 2, return_distance=True), (tf.TensorSpec(shape=(None, 2), dtype=tf.float32), tf.TensorSpec(shape=(None, ), dtype=tf.int64)), tree) scale = tf.reduce_mean(dists[:, 1]) coords = coords * (2. / scale) all_coords = [coords] trees = [tree] radii = 4 * np.power(2, np.arange(depth)) flat_indices = [] row_splits = [] rel_coords = [] sample_indices = [] def add_conv(tree, coords, radius, k0): def fn(tree, coords): indices = tree.query_ball_point(coords, radius, approx_neighbors=k0) rc = np.repeat(coords, indices.row_lengths, axis=0) - coords[indices.flat_values] rc /= radius return flat_indices, row_splits, rel_coords fi, rs, rc = context.py_function( fn, (tf.TensorSpec(shape=(None, ), dtype=tf.int64), tf.TensorSpec(shape=(None, ), dtype=tf.int64), tf.TensorSpec(shape=(None, 3), dtype=tf.float32)), (tree, coords)) flat_indices.append(fi) row_splits.append(rs) rel_coords.append(rc) # n = tree.n # m = coords.shape[0] # e = indices.row_splits[-1] # lines.append(str((e, n, m, e / n, e / m, radius))) return fi, rs # initial query in order to do initial rejection sample # indices = tree.query_ball_point(coords, radii[0], approx_neighbors=k0) indices = context.py_function( lambda tree, coords: np.array( core.rejection_sample_active(tree, coords.numpy(), radii[0], k0)), tf.TensorSpec(shape=(None, ), dtype=tf.int64), tree, coords) # indices = np.array(core.rejection_sample_lazy(tree, coords, radii[0], k0)) sample_indices.append(indices) out_coords = tf.gather(coords, indices) all_coords.append(out_coords) tree = context.py_function(tree_impl, None, out_coords) trees.append(tree) # initial large down-sample conv add_conv(tree, coords, radii[0] * 2, k0 * 4) coords = out_coords def rejection_sample(flat_indices, row_splits): from more_keras.ragged.np_impl import RaggedArray ra = RaggedArray.from_row_splits(flat_indices, row_splits) return np.array(core.rejection_sample_precomputed(ra), dtype=np.int64) for i in range(1, depth - 1): # in place indices_comp = add_conv(tree, coords, radii[i], k0) indices = context.py_function( rejection_sample, tf.TensorSpec(shape=(None, ), dtype=tf.int64), *indices_comp) sample_indices.append(indices) out_coords = tf.gather(coords, indices) all_coords.append(out_coords) tree = context.py_function(tree_impl, None, out_coords) trees.append(tree) # down sample # larger radius means number of edges remains roughly constant # number of ops remains constant if number of filters doubles # also makes more neighbors for unsampling (~4 on average vs ~2) add_conv(tree, coords, radii[i] * SQRT_2, k0) coords = out_coords # final in_place add_conv(tree, coords, radii[-1], k0) # lines.append('***') # print('\n'.join(lines)) # DEBUG return ( tuple(flat_indices), tuple(rel_coords), tuple(row_splits), tuple(all_coords), tuple(sample_indices), ) # out = dists, indices dataset = problem.get_base_dataset('validation') dataset = dataset.map(f, -1) out = tf.compat.v1.data.make_one_shot_iterator(dataset).get_next() with tf.Session() as sess: for _ in tqdm(range(num_warmup), desc='warming up'): sess.run(out) t = time() for _ in tqdm(range(num_iterations), desc='profiling'): sess.run(out) dt = time() - t print('Ran {} iterations in {} s: {} ms / iteration'.format( num_iterations, dt, dt * 1000 / num_iterations))
# print(sess.run(tf.reduce_max(tf.abs(conv0 - conv2)))) def grid_conv_benchmark(): grid_features = tf.Variable(np.random.normal(size=(8, 64, 64, 64)), dtype=tf.float32) grid_conv = tf.keras.layers.Conv2D(64, 3)(grid_features) grid_grad, = tf.gradients(grid_conv, grid_features) with tf.Session() as sess: sess.run(tf.global_variables_initializer()) bm = tf.test.Benchmark() bm.run_op_benchmark(sess, (grid_conv, grid_grad)) for coords, labels in tfds.as_numpy( PartnetProblem().get_base_dataset(split='validation').take(1)): break coords = coords[:4096] tree = pykd.KDTree(coords) dists, indices = tree.query(tree.data, 2, return_distance=True) scale = np.mean(dists[:, 1]) assert (scale > 0) coords *= (2 / scale) print('in_place') in_place_conv_benchmark(coords) print('subsample') subsample_conv_benchmark(coords) print('grid')
from __future__ import absolute_import from __future__ import division from __future__ import print_function import numpy as np from time import time from tqdm import tqdm import tensorflow_datasets as tfds from deep_cloud.problems.partnet import PartnetProblem from deep_cloud.ops.np_utils.tree_utils import pykd problem = PartnetProblem() warm_up = 5 benchmark = 10 total = warm_up + benchmark tree_impl = pykd.KDTree all_coords = [] for coords, _ in tqdm(tfds.as_numpy( problem.get_base_dataset('validation').take(total)), total=total, desc='getting base data...'): tree = tree_impl(coords) dists, indices = tree.query(tree.data, 2, return_distance=True) del indices scale = np.mean(dists[:, 1]) coords *= 2 / scale all_coords.append(coords)
from __future__ import absolute_import from __future__ import division from __future__ import print_function import numpy as np import tensorflow as tf import tensorflow_datasets as tfds from more_keras.ops import utils as op_utils from deep_cloud.problems.partnet import PartnetProblem from deep_cloud.ops.np_utils.tree_utils import pykd for coords, labels in tfds.as_numpy( PartnetProblem().get_base_dataset().take(1)): break tree = pykd.KDTree(coords) dists, indices = tree.query(tree.data, 2, return_distance=True) scale = np.mean(dists[:, 1]) assert (scale > 0) coords *= (2 / scale) tree = pykd.KDTree(coords) indices = tree.query_ball_point(coords, 4, approx_neighbors=16) c0 = coords.shape[0] # nf = 64 nf = 16 nd = 3 features = np.random.uniform(size=(c0, nf, nd))
if __name__ == '__main__': from time import time import functools from deep_cloud.problems.partnet import PartnetProblem from deep_cloud.models.sparse import preprocess as pp import tqdm # tf.compat.v1.enable_eager_execution() # tf.config.experimental_run_functions_eagerly(True) split = 'train' batch_size = 16 # batch_size = 2 num_classes = 4 num_warmup = 5 num_batches = 10 problem = PartnetProblem() with problem: dataset = problem.get_base_dataset(split).map(pp.pre_batch_map, -1) dataset = dataset.batch(batch_size) dataset = dataset.map(pp.post_batch_map, -1).prefetch(-1) # for args in dataset.take(1): # if len(args) == 3: # inputs, labels, weights = args # else: # inputs, labels = args # logits = semantic_segmenter_logits(inputs, # problem.output_spec.shape[-1]) # print('Successs!') # exit()