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