def offset_batched_neighbors(self): if self._offset_batched_neighbors is None: batched_sample_indices = b.as_batched_model_input( self._sample_indices) self._offset_batched_neighbors = utils.map_gather( self._base.offset_batched_neighbors, batched_sample_indices) return self._offset_batched_neighbors
def offset_batched_neighbors(self): if self._offset_batched_neighbors is None: batched_neighbors = b.as_batched_model_input(self.neighbors) offset = utils.get_row_offsets(batched_neighbors) self._offset_batched_neighbors = utils.apply_row_offset( batched_neighbors, offset) return self._offset_batched_neighbors
def __call__(self, x, max_radius2): key = (x, max_radius2) if key not in self._cache: xb = b.as_batched_model_input(x) value = tf.ragged.map_flat_values( lambda x: self._fn(x, max_radius2), xb) self._cache[key] = value return self._cache[key]
def pointnet2_head(coords, normals=None, specs=None, global_mlp=None, coords_as_features=True): if specs is None: specs = ( SetAbstractionSpec(0.2, 512, pointnet2_mlp((64, 64, 128))), SetAbstractionSpec(0.4, 128, pointnet2_mlp((128, 128, 256))), ) if global_mlp is None: global_mlp = pointnet2_mlp((256, 512, 1024)) features = None if normals is None else b.as_batched_model_input(normals) for radius, n_out, mlp in specs: neighbors, sample_rate = core.query_pairs(coords, radius) neighborhood = neigh.InPlaceNeighborhood(coords, neighbors) sample_indices = sample.sample(sample_rate, n_out) neighborhood = neigh.SampledNeighborhood(neighbors, sample_indices) features = set_abstraction(features, neighborhood, mlp, coords_as_features=coords_as_features) coords = neighborhood.out_coords
def flat_rel_coords(): return b.as_batched_model_input( neighborhood.rel_coords.flat_values).flat_values
def __call__(self, x, max_radius2): batched_x = b.as_batched_model_input(x) x_features = tf.ragged.map_flat_values( lambda f: self._fn(f, max_radius2), batched_x) return x_features
def _test_transposed_consistent(self): batch_size = 3 with tf.device('/cpu:0'): np_data = [ np.random.uniform(size=(100, 3)).astype(np.float32), np.random.uniform(size=(200, 3)).astype(np.float32), np.random.uniform(size=(50, 3)).astype(np.float32), ] def generator(): labels = np.array([0, 1, 2], dtype=np.int64) indices = [ np.array([0, 5, 2, 7, 10], dtype=np.int64), np.array([1, 4, 3, 2], dtype=np.int64), np.array([10, 15], dtype=np.int64), ] yield (np_data[0], indices[0]), labels[0] yield (np_data[1], indices[1]), labels[1] yield (np_data[2], indices[2]), labels[2] dataset = tf.data.Dataset.from_generator( generator, output_types=((tf.float32, tf.int64), tf.int64), output_shapes=((tf.TensorShape( (None, 3)), tf.TensorShape((None,))), tf.TensorShape(()))) coords = b.prebatch_input((None, 3), tf.float32) sample_indices = b.prebatch_input((None,), dtype=tf.int64) labels = b.prebatch_input((), dtype=tf.int64) neighbors = q.query_pairs(coords, 0.1) in_place_neighborhood = n.InPlaceNeighborhood(coords, neighbors) sampled_neighborhood = n.SampledNeighborhood( in_place_neighborhood, sample_indices) transposed = sampled_neighborhood.transpose simple = n.Neighborhood(sampled_neighborhood.out_coords, sampled_neighborhood.in_coords, transposed.neighbors) simple_obn = simple.offset_batched_neighbors trans_obn = transposed.offset_batched_neighbors simple_out = b.as_batched_model_input(simple.in_coords) trans_out = b.as_batched_model_input(transposed.in_coords) out = [[o.flat_values, o.nested_row_splits] for o in (simple_obn, trans_obn, simple_out, trans_out)] flat_out = tf.nest.flatten(out) model = b.model(flat_out) preprocessor = b.preprocessor(b.batched(labels)) dataset = preprocessor.map_and_batch(dataset, batch_size) # if tf.executing_eagerly(): # for data, label in dataset.take(1): # pass # else: data, label = tf.compat.v1.data.make_one_shot_iterator( dataset).get_next() flat_out = model(tf.nest.flatten(data)) if not isinstance(flat_out, (list, tuple)): flat_out = flat_out, out = tf.nest.pack_sequence_as(out, flat_out) out = [tf.RaggedTensor.from_nested_row_splits(*o) for o in out] out, label = self.evaluate((out, label)) simple_obn, trans_obn, simple_coords, trans_coords = out self.assertRaggedEqual(simple_obn, trans_obn) self.assertRaggedEqual(simple_coords, trans_coords) self.assertEqual(simple_obn.nested_row_splits[-1][-1], sum(d.shape[0] for d in np_data))
def cls_head(coords, normals=None, r0=0.1, initial_filters=(16, ), initial_activation=cls_head_activation, filters=(32, 64, 128, 256), global_units='combined', query_fn=core.query_pairs, radii_fn=core.constant_radii, coords_transform=None, weights_transform=None, convolver=None): if convolver is None: convolver = c.ExpandingConvolver(activation=cls_head_activation) if coords_transform is None: coords_transform = t.polynomial_transformer(max_order=1) if weights_transform is None: weights_transform = t.ctg_transformer() # weights_transform = lambda *args, **kwargs: None n_res = len(filters) unscaled_radii2 = radii_fn(n_res) if isinstance(unscaled_radii2, tf.Tensor): assert (unscaled_radii2.shape == (n_res, )) radii2 = utils.lambda_call(tf.math.scalar_mul, r0**2, unscaled_radii2) radii2 = tf.keras.layers.Lambda(tf.unstack, arguments=dict(axis=0))(radii2) for i, radius2 in enumerate(radii2): tb.add_custom_scalar('radius{}'.format(i), tf.sqrt(radius2)) # tf.compat.v1.summary.scalar('r%d' % i, # tf.sqrt(radius2), # family='radii') else: radii2 = unscaled_radii2 * (r0**2) def maybe_feed(r2, r20): if isinstance(r2, (tf.Tensor, tf.Variable)): r = tf.keras.layers.Lambda(tf.sqrt)(radius2) return cache.get_cached(r, r20) else: return np.sqrt(r2) features = b.as_batched_model_input(normals) for f in initial_filters: layer = Dense(f) features = tf.ragged.map_flat_values(layer, features) features = tf.ragged.map_flat_values(initial_activation, features) features = utils.flatten_leading_dims(features, 2) global_features = [] default_r0 = r0 for i, radius2 in enumerate(radii2): neighbors, sample_rate = query_fn(coords, maybe_feed(radius2, default_r0**2), name='query%d' % i) default_r0 *= 2 if not isinstance(radius2, tf.Tensor): radius2 = source.constant(radius2, dtype=tf.float32) neighborhood = n.InPlaceNeighborhood(coords, neighbors) features, nested_row_splits = core.convolve(features, radius2, filters[i], neighborhood, coords_transform, weights_transform, convolver.in_place_conv) if global_units == 'combined': coord_features = coords_transform(neighborhood.out_coords, None) global_features.append( convolver.global_conv(features, coord_features, nested_row_splits[-2], filters[i])) if i < n_res - 1: sample_indices = sample.sample( sample_rate, tf.keras.layers.Lambda(lambda s: tf.size(s) // 4)(sample_rate)) neighborhood = n.SampledNeighborhood(neighborhood, sample_indices) features, nested_row_splits = core.convolve( features, radius2, filters[i + 1], neighborhood, coords_transform, weights_transform, convolver.resample_conv) coords = neighborhood.out_coords # global_conv if global_units == 'combined': features = tf.keras.layers.Lambda( tf.concat, arguments=dict(axis=-1))(global_features) else: coord_features = coords_transform(coords, None) features = convolver.global_conv(features, coord_features, nested_row_splits[-2], global_units) return features