def offset_batched_neighbors(self):
"""
[B, n?, k?] offset index into flattened input batched cloud features.
For point cloud features of size [B, n?, ...],
`tf.gather(features.values, offset_batched_neighbors)` will give the
same result as
```
tf.map(lambda features, batched_neighbors:
tf.gather(features, batched_neighbors))
```
The returned value is a model input.
"""
if self._offset_batched_neighbors is None:
row_length = utils.leading_dim(self.in_coords)
row_lengths = b.batched(row_length)
batched_neighbors = b.batched(self.neighbors)
# row_lengths = b.as_batched_model_input(row_length)
# batched_neighbors = b.as_batched_model_input(self.neighbors)
offset = tf.keras.layers.Lambda(
tf.math.cumsum, arguments=dict(exclusive=True))(row_lengths)
offset_batched_neighbors = utils.apply_row_offset(
batched_neighbors, offset)
self._offset_batched_neighbors = b.as_model_input(
offset_batched_neighbors)
return self._offset_batched_neighbors
def test_batched_tensor(self):
b = builder.MetaNetworkBuilder()
with b:
x = b.prebatch_input(shape=(4, 5), dtype=tf.float32)
batched = builder.batched(x)
self.assertEqual(batched.shape.as_list(), [None, 4, 5])
b.preprocessor(())
def test_batched_ragged(self):
b = builder.MetaNetworkBuilder()
with b:
x = b.prebatch_input(shape=(None, 5), dtype=tf.float32)
batched = builder.batched(x)
self.assertTrue(isinstance(batched, tf.RaggedTensor))
self.assertEqual(batched.shape.as_list(), [None, None, 5])
b.preprocessor(())
def batch_labels_and_weights(self, labels, weights=None):
from more_keras.meta_models import builder as b
return tf.nest.map_structure(
lambda tensor: None
if tensor is None else b.batched(tensor), (labels, weights))
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))