def fixed_budget_radii(num_res):
"""
Learnable unscaled squared radii based on a fixed computational budget.
The budget assumes point cloud density quarters at each resolution, and
aims to keep the sum across all resolutions of the sum of all neighborhood
sizes roughly constante. Note there are generally more filters in later
resolutions, so the number of floating point operations may increase if the
model learns to expand final layer radii at the expense of lower
resolutions.
Args:
num_res: number of resolutions.
Returns:
[num_res] float32 tensor of squared unscaled radii values, initially
the same as `constant_radii(num_res)`.
"""
r20 = constant_radii(num_res)
radii = source.variable(shape=(num_res, ),
dtype=tf.float32,
initializer=tf.constant_initializer(np.sqrt(r20)))
radii2 = tf.keras.layers.Lambda(tf.square)(radii)
density = np.power(4.0, -np.arange(num_res)).astype(np.float32)
budget = utils.lambda_call(_cost, r20, density)
return utils.lambda_call(_rescale_for_budget, radii2, density, budget)
def mlp_global_deconv(global_features, coord_features, row_splits_or_k,
network_fn):
global_features = layer_utils.lambda_call(_expand_and_tile,
global_features, row_splits_or_k)
features = tf.keras.layers.Lambda(
tf.concat, arguments=dict(axis=-1))([global_features, coord_features])
return network_fn(features)
def query_pairs(coords,
radius,
name=None,
max_neighbors=None,
sample_rate_reciprocal_offset=0):
"""
Get neighbors and inverse density sample rates from the provided coords.
sample_rate is given by
sample_rate = 1. / (sum(num_neighbors) - sample_rate_reciprocal_offset)
Args:
coords: [n, num_dims] float32 point cloud coordinates
radius: float scalar ball search radius
name: used in layers
max_neighbors: if not None, neighborhoods are cropped to this size.
sample_rate_reciprocal_offset: affects sample rate. Values close to 1
result in sparse neighborhoods being selected more frequently
Returns:
neighbors: [n, ?] int32 ragged tensor of neighboring indices
sample_rate: sampling rate based roughly on inverse density.
"""
neighbors = _q.query_pairs(coords,
radius,
name=name,
max_neighbors=max_neighbors)
sample_rate_reciprocal = utils.cast(utils.row_lengths(neighbors),
tf.float32)
if sample_rate_reciprocal_offset != 0:
sample_rate_reciprocal = utils.lambda_call(
tf.math.subtract, sample_rate_reciprocal,
sample_rate_reciprocal_offset)
sample_rate = utils.reciprocal(sample_rate_reciprocal)
return neighbors, sample_rate
def flat_expanding_edge_conv(node_features,
coord_features,
indices,
row_splits_or_k,
weights=None):
features = layer_utils.lambda_call(conv_ops.flat_expanding_edge_conv,
node_features, coord_features, indices,
row_splits_or_k, weights)
return features
def prebatch_input(self, shape, dtype, name=None):
inp = tf.keras.layers.Input(shape=shape,
dtype=dtype,
batch_size=1,
name=name)
self._prebatch_inputs.append(inp)
inp = layer_utils.lambda_call(tf.squeeze, inp, axis=0)
self._mark(inp, Marks.PREBATCH)
return inp
def f(inputs, min_log_value=-10):
x = inputs
x = layer_utils.lambda_call(tf.abs, x)
x = layer_utils.lambda_call(tf.math.log, x)
x = layer_utils.lambda_call(tf.maximum, x, min_log_value)
x = dense(x)
x = layer_utils.lambda_call(tf.exp, x)
angle = layer_utils.lambda_call(_complex_angle, inputs)
angle = dense(angle)
components = layer_utils.lambda_call(_angle_to_unit_vector, angle)
x = layer_utils.lambda_call(tf.expand_dims, x, axis=-1)
x = layer_utils.lambda_call(tf.multiply, x, components)
x = layer_utils.flatten_final_dims(x, n=2)
x = layer_utils.lambda_call(_zeros_if_any_small, inputs, x)
if dropout_rate is not None:
x = tf.keras.layers.Dropout(dropout_rate)(x)
if use_batch_norm:
x = tf.keras.layers.BatchNormalization()(x)
return x
def reduce_flat_mean(x, row_splits_or_k, weights, eps=1e-7):
return layer_utils.lambda_call(conv_ops.reduce_flat_mean,
x,
row_splits_or_k,
weights,
eps=eps)
def flat_expanding_global_deconv(global_features, coord_features,
row_splits_or_k):
features = layer_utils.lambda_call(conv_ops.flat_expanding_global_deconv,
global_features, coord_features,
row_splits_or_k)
return features
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
