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} }