def add_profile(self, epoch, logs, run_metadata: tf.RunMetadata, profiler: tf.profiler.Profiler, profile_writer: tf.summary.FileWriter, save_path: str): timeline_path = f"{save_path}/timeline" pathlib.Path(timeline_path).mkdir(exist_ok=True, parents=True) profiler.add_step(epoch, run_meta=run_metadata) opts = (option_builder.ProfileOptionBuilder( option_builder.ProfileOptionBuilder.time_and_memory()).with_step( epoch).with_timeline_output(f"{timeline_path}/step").build()) profiler.profile_graph(options=opts) profile_writer.add_run_metadata(run_metadata, f"step{epoch}")
def _train_step(self, obs, states, rewards, masks, actions, values, update, writer: tf.summary.FileWriter = None, features=None, rewards_bonuses=None): """ applies a training step to the model :param obs: ([float]) The input observations :param states: ([float]) The states (used for recurrent policies) :param rewards: ([float]) The rewards from the environment :param masks: ([bool]) Whether or not the episode is over (used for recurrent policies) :param actions: ([float]) The actions taken :param values: ([float]) The logits values :param update: (int) the current step iteration :param writer: (TensorFlow Summary.writer) the writer for tensorboard :return: (float, float, float) policy loss, value loss, policy entropy """ advs = rewards - values cur_lr = None for _ in range(len(obs)): cur_lr = self.learning_rate_schedule.value() assert cur_lr is not None, "Error: the observation input array cannon be empty" rewards_bonuses = rewards_bonuses if self.use_sf else np.zeros_like(rewards_bonuses) td_map = {self.train_model.obs_ph: obs, self.actions_ph: actions, self.advs_ph: advs, self.rewards_ph: rewards + rewards_bonuses, self.learning_rate_ph: cur_lr, self.successor_feature_ph: features} if states is not None: td_map[self.train_model.states_ph] = states td_map[self.train_model.masks_ph] = masks if writer is not None: # run loss backprop with summary, but once every 10 runs save the metadata (memory, compute time, ...) if (1 + update) % 10 == 0: run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE) run_metadata = tf.RunMetadata() summary, policy_loss, value_loss, policy_entropy, _, sf_loss = self.sess.run( [self.summary, self.pg_loss, self.vf_loss, self.entropy, self.apply_backprop, self.sf_loss], td_map, options=run_options, run_metadata=run_metadata) writer.add_run_metadata(run_metadata, 'step%d' % (update * (self.n_batch + 1))) else: summary, policy_loss, value_loss, policy_entropy, _, sf_loss = self.sess.run( [self.summary, self.pg_loss, self.vf_loss, self.entropy, self.apply_backprop, self.sf_loss], td_map) writer.add_summary(summary, update * (self.n_batch + 1)) else: policy_loss, value_loss, policy_entropy, _, sf_loss = self.sess.run( [self.pg_loss, self.vf_loss, self.entropy, self.apply_backprop, self.sf_loss], td_map) return policy_loss, value_loss, policy_entropy, sf_loss
def log_trace(cfg: TrainConfig, step: int, opts, summary_writer: tf.summary.FileWriter): logger.debug("Starting trace logging") fetched_timeline = timeline.Timeline(opts['run_metadata'].step_stats) chrome_trace = fetched_timeline.generate_chrome_trace_format( show_memory=True) with open(os.path.join(cfg.logdir, f'timeline_{step:05}.json'), 'w') as f: f.write(chrome_trace) summary_writer.add_run_metadata(opts['run_metadata'], f"step_{step:05}", global_step=step) logger.info( f"Saved trace metadata both to timeline_{step:05}.json and step_{step:05} in tensorboard" )