def _training_summary(self, training_info, loss_info, grads_and_vars):
if self._summarize_grads_and_vars:
summary_utils.add_variables_summaries(grads_and_vars,
self._train_step_counter)
summary_utils.add_gradients_summaries(grads_and_vars,
self._train_step_counter)
if self._debug_summaries:
common.add_action_summaries(training_info.action,
self.env.action_spec())
common.add_loss_summaries(loss_info)
if self._summarize_action_distributions:
summary_utils.summarize_action_dist(
training_info.action_distribution, self.env.action_spec())
if training_info.collect_action_distribution:
summary_utils.summarize_action_dist(
action_distributions=training_info.
collect_action_distribution,
action_specs=self.env.action_spec(),
name="collect_action_dist")
for metric in self.get_metrics():
metric.tf_summaries(
train_step=self._train_step_counter,
step_metrics=self.get_metrics()[:2])
mem = tf.py_function(
lambda: self._proc.memory_info().rss // 1e6, [],
tf.float32,
name='memory_usage')
if not tf.executing_eagerly():
mem.set_shape(())
tf.summary.scalar(name='memory_usage', data=mem)
def assign_variables(variable_mapping):
"""Assign variables according to the provided `variable_mapping`.
Args:
variable_mapping: An iterable of variable pairs, each corresponding to a
variable whose value is to be overwitten (destination) and a reference
variable (source).
Returns:
If running in TensorFlow Eager mode, returns None; otherwise, returns a list
of assignment operations.
"""
for variable, reference_variable in variable_mapping:
if tf.executing_eagerly():
# Just perform the assignment.
variable.assign(reference_variable)
else:
# Piggyback on the variable's initializer attribute, which is included in
# `tf.global_variables_initializer`.
initializer_ops = [variable._initializer_op] # pylint: disable=protected-access
if isinstance(reference_variable, tf.Variable):
initializer_ops += [reference_variable._initializer_op] # pylint: disable=protected-access
with tf.control_dependencies(initializer_ops):
assign_op = variable.assign(reference_variable)
variable._initializer_op = assign_op # pylint: disable=protected-access
def result(self):
def _result():
return np.sum(self._buffer) / self.buffer_size
result_value = tf.py_function(_result, [],
tf.float32,
name='metric_result_py_func')
if not tf.executing_eagerly():
return result_value.set_shape(())
return result_value
def summarize_metrics(self):
"""Generate summaries for metrics `AverageEpisodeLength`, `AverageReturn`..."""
if self._metrics:
for metric in self._metrics:
metric.tf_summaries(train_step=common.get_global_counter(),
step_metrics=self._metrics[:2])
mem = tf.py_function(lambda: self._proc.memory_info().rss // 1e6, [],
tf.float32,
name='memory_usage')
if not tf.executing_eagerly():
mem.set_shape(())
tf.summary.scalar(name='memory_usage', data=mem)
def optimizer_update(iterate_collection, iteration_idx, objective_fn,
update_fn, get_params_fn, first_order, clip_grad_norm):
"""Returns the next iterate in the optimization of objective_fn wrt variables.
Args:
iterate_collection: A (potentially structured) container of tf.Tensors
corresponding to the state of the current iterate.
iteration_idx: An int Tensor; the iteration number.
objective_fn: Callable that takes in variables and produces the value of the
objective function.
update_fn: Callable that takes in the gradient of the objective function and
the current iterate and produces the next iterate.
get_params_fn: Callable that takes in the gradient of the objective function
and the current iterate and produces the next iterate.
first_order: If True, prevent the computation of higher order gradients.
clip_grad_norm: If not None, gradient dimensions are independently clipped
to lie in the interval [-clip_grad_norm, clip_grad_norm].
"""
variables = [get_params_fn(iterate) for iterate in iterate_collection]
if tf.executing_eagerly():
with tf.GradientTape(persistent=True) as g:
g.watch(variables)
loss = objective_fn(variables, iteration_idx)
grads = g.gradient(loss, variables)
else:
loss = objective_fn(variables, iteration_idx)
grads = tf.gradients(ys=loss, xs=variables)
if clip_grad_norm:
grads = [
tf.clip_by_value(grad, -1 * clip_grad_norm, clip_grad_norm)
for grad in grads
]
if first_order:
grads = [tf.stop_gradient(dv) for dv in grads]
return [
update_fn(i=iteration_idx, grad=dv, state=s)
for (s, dv) in zip(iterate_collection, grads)
]
def bit_variable_generator(model, imported_model):
"""Pair variables for reparameterizable and BiT backbones.
Args:
model: A valid `reparameterizable_backbones.ReparameterizableBackbone`.
imported_model: A `tf.AutoTrackable`.
Returns:
Pairs of corresponding `tf.Variable`s from `model` and `imported_model`.
Raises:
ValueError: If `model` is not a valid model.
"""
valid_models = (
reparameterizable_backbones.ResNet18V2,
reparameterizable_backbones.GNResNet18V2,
reparameterizable_backbones.ResNet50V2,
reparameterizable_backbones.GNResNet50V2,
)
if all(not isinstance(model, valid_model) for valid_model in valid_models):
raise ValueError(
'BiT variable mapping defined only for models among: {}'.format(
valid_models))
pairs = []
name_to_variable = dict((v.name, v) for v in model.variables)
imported_name_to_variable = dict(
(v.name, v) for v in imported_model.variables)
group_norm = isinstance(model, reparameterizable_backbones.GNResNet)
# Assign root layer variables.
# TODO(eringrant): Determine why 'input_stage' is not prepended to variables
# in tensorflow Eager mode.
input_stage_prefix = 'input_stage/' if not tf.executing_eagerly() else ''
output_stage_prefix = 'output_stage/' if not tf.executing_eagerly() else ''
name = '{}conv2d/kernel:0'.format(input_stage_prefix)
imported_name = 'resnet/root_block/{}conv2d/kernel:0'.format(
'standardized_' if group_norm else '')
pairs += [(name_to_variable[name],
imported_name_to_variable[imported_name])]
# Assign normalization variables.
for n in ('gamma', 'beta') + (() if group_norm else
('moving_mean', 'moving_variance')):
name = '{}{}_normalization_{}/{}:0'.format(
output_stage_prefix, 'group' if group_norm else 'batch', 16 if
(isinstance(model, reparameterizable_backbones.ResNet18V2)
or isinstance(model, reparameterizable_backbones.GNResNet18V2))
else 48, n)
imported_name = 'resnet/{}/{}:0'.format(
'group_norm' if group_norm else 'batch_normalization', n)
pairs += [(name_to_variable[name],
imported_name_to_variable[imported_name])]
# Assign other layer variables.
if (isinstance(model, reparameterizable_backbones.ResNet18V2)
or isinstance(model, reparameterizable_backbones.GNResNet18V2)):
for i, (stage, block, layer) in enumerate(
itertools.product(range(4), range(2), range(2))):
imported_base_path = 'resnet/block{}/unit0{}/{}/'.format(
stage + 1, block + 1, {
0: 'a',
1: 'b'
}[layer])
for n in ('gamma', 'beta') + (() if group_norm else
('moving_mean', 'moving_variance')):
name = 'stage_{}/block_{}/{}_normalization{}/{}:0'.format(
stage, block, 'group' if group_norm else 'batch',
'_{}'.format(i) if i > 0 else '', n)
imported_name = (imported_base_path + '{}/{}:0'.format(
'group_norm' if group_norm else 'batch_normalization', n))
pairs += [(name_to_variable[name],
imported_name_to_variable[imported_name])]
if stage > 0 and block == layer == 0:
name = 'stage_{}/block_{}/shortcut_conv/kernel:0'.format(
stage, block)
imported_name = (imported_base_path +
'proj/{}conv2d/kernel:0'.format(
'standardized_' if group_norm else ''))
pairs += [(name_to_variable[name],
imported_name_to_variable[imported_name])]
name = 'stage_{}/block_{}/conv2d_{}/kernel:0'.format(
stage, block, i + 1)
imported_name = (imported_base_path + '{}conv2d/kernel:0'.format(
'standardized_' if group_norm else ''))
pairs += [(name_to_variable[name],
imported_name_to_variable[imported_name])]
elif (isinstance(model, reparameterizable_backbones.ResNet50V2)
or isinstance(model, reparameterizable_backbones.GNResNet50V2)):
conv_id = 1
batch_norm_id = 0
for stage, num_blocks in enumerate((3, 4, 6, 3)):
for block in range(num_blocks):
for layer in range(3):
imported_base_path = 'resnet/block{}/unit0{}/{}/'.format(
stage + 1, block + 1, {
0: 'a',
1: 'b',
2: 'c'
}[layer])
for n in ('gamma', 'beta') + (
() if group_norm else
('moving_mean', 'moving_variance')):
name = 'stage_{}/block_{}/{}_normalization{}/{}:0'.format(
stage, block, 'group' if group_norm else 'batch',
'_{}'.format(batch_norm_id)
if batch_norm_id > 0 else '', n)
imported_name = (imported_base_path + '{}/{}:0'.format(
'group_norm'
if group_norm else 'batch_normalization', n))
pairs += [(name_to_variable[name],
imported_name_to_variable[imported_name])]
batch_norm_id += 1
name = 'stage_{}/block_{}/conv2d_{}/kernel:0'.format(
stage, block, conv_id)
imported_name = (
imported_base_path + '{}conv2d/kernel:0'.format(
'standardized_' if group_norm else ''))
pairs += [(name_to_variable[name],
imported_name_to_variable[imported_name])]
conv_id += 1
if block == 0:
name = 'stage_{}/block_{}/conv2d{}/kernel:0'.format(
stage, block,
'_{}'.format(conv_id) if conv_id > 0 else '')
imported_name = 'resnet/block{}/unit01/a/proj/{}conv2d/kernel:0'.format(
stage + 1, 'standardized_' if group_norm else '')
pairs += [(name_to_variable[name],
imported_name_to_variable[imported_name])]
conv_id += 1
else:
raise NotImplementedError('BiT variable mappings are not defined for '
'models other than {}.'.format(valid_models))
return pairs
def prepare_lidar_images_and_correspondences(
inputs,
resized_image_height,
resized_image_width,
camera_names=('front', 'front_left', 'front_right', 'side_left',
'side_right'),
lidar_names=('top', 'front', 'side_left', 'side_right', 'rear')):
"""Integrates and returns the lidars, cameras and their correspondences.
Args:
inputs: A dictionary containing the images and point / pixel
correspondences.
resized_image_height: Target height of the images.
resized_image_width: Target width of the images.
camera_names: List of cameras to include images from.
lidar_names: List of lidars to include point clouds from.
Returns:
A tf.float32 tensor of size [num_points, 3] containing point positions.
A tf.float32 tensor of size [num_points, 1] containing point intensities.
A tf.float32 tensor of size [num_points, 1] containing point elongations.
A tf.float32 tensor of size [num_points, 3] containing point normals.
A tf.float32 tensor of size [num_images, resized_image_height,
resized_image_width, 3].
A tf.int32 tensor of size [num_images, num_points, 2].
Raises:
ValueError: If camera_names or lidar_names are empty lists.
"""
if not camera_names:
raise ValueError('camera_names should contain at least one name.')
if not lidar_names:
raise ValueError('lidar_names should contain at least one name.')
(points_position, points_intensity, points_elongation, points_normal,
points_in_image_frame_yx, points_in_image_frame_id) = _prepare_lidar_points(
inputs=inputs, lidar_names=lidar_names)
images = []
points_in_image_frame = []
for camera_name in camera_names:
image_key = ('cameras/%s/image' % camera_name)
image_height = tf.shape(inputs[image_key])[0]
image_width = tf.shape(inputs[image_key])[1]
height_ratio = tf.cast(
resized_image_height, dtype=tf.float32) / tf.cast(
image_height, dtype=tf.float32)
width_ratio = tf.cast(
resized_image_width, dtype=tf.float32) / tf.cast(
image_width, dtype=tf.float32)
if tf.executing_eagerly():
resize_method = tf.image.ResizeMethod.NEAREST_NEIGHBOR
else:
resize_method = tf.image.ResizeMethod.BILINEAR
if inputs[image_key].dtype in [
tf.int8, tf.uint8, tf.int16, tf.uint16, tf.int32, tf.int64
]:
resize_method = tf.image.ResizeMethod.NEAREST_NEIGHBOR
images.append(
tf.image.resize(
images=inputs[image_key],
size=[resized_image_height, resized_image_width],
method=resize_method,
antialias=True))
camera_id = tf.cast(inputs[('cameras/%s/id' % camera_name)], dtype=tf.int32)
valid_points = tf.equal(points_in_image_frame_id, camera_id)
valid_points = tf.tile(valid_points, [1, 2])
point_coords = tf.cast(
tf.cast(points_in_image_frame_yx, dtype=tf.float32) *
tf.stack([height_ratio, width_ratio]),
dtype=tf.int32)
points_in_image_frame_camera = tf.where(
valid_points, point_coords, -tf.ones_like(valid_points, dtype=tf.int32))
points_in_image_frame.append(points_in_image_frame_camera)
num_images = len(images)
images = tf.stack(images, axis=0)
images.set_shape([num_images, resized_image_height, resized_image_width, 3])
points_in_image_frame = tf.stack(points_in_image_frame, axis=0)
return {
'points_position': points_position,
'points_intensity': points_intensity,
'points_elongation': points_elongation,
'points_normal': points_normal,
'view_images': {'rgb_view': images},
'view_indices_2d': {'rgb_view': points_in_image_frame}
}