def _set_up(self, eval_mode):
"""Sets up the runner by creating and initializing the agent."""
# Reset the tf default graph to avoid name collisions from previous runs
# before doing anything else.
tf.reset_default_graph()
self._summary_writer = tf.summary.FileWriter(self._output_dir)
if self._episode_log_file:
self._episode_writer = tf.io.TFRecordWriter(
os.path.join(self._output_dir, self._episode_log_file))
# Set up a session and initialize variables.
self._sess = tf.Session(config=tf.ConfigProto(
allow_soft_placement=True))
self._agent = self._create_agent_fn(
self._sess,
self._env,
summary_writer=self._summary_writer,
eval_mode=eval_mode)
# type check: env/agent must both be multi- or single-user
if self._agent.multi_user and not isinstance(
self._env.environment, environment.MultiUserEnvironment):
raise ValueError(
'Multi-user agent requires multi-user environment.')
if not self._agent.multi_user and isinstance(
self._env.environment, environment.MultiUserEnvironment):
raise ValueError(
'Single-user agent requires single-user environment.')
self._summary_writer.add_graph(graph=tf.get_default_graph())
self._sess.run(tf.global_variables_initializer())
self._sess.run(tf.local_variables_initializer())
def init_sess(var_list=None, path=None):
sess = tf.Session()
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
saver = None
writer = None
if var_list is not None:
saver = tf.train.Saver(var_list=var_list, filename='model.ckpt')
if path is not None:
writer = tf.summary.FileWriter(logdir=path)
return sess, saver, writer
def _set_up(self, eval_mode):
"""Sets up the runner by creating and initializing the agent."""
# Reset the tf default graph to avoid name collisions from previous runs
# before doing anything else.
tf.reset_default_graph()
self._summary_writer = tf.summary.FileWriter(self._output_dir)
if self._episode_log_file:
self._episode_writer = tf.python_io.TFRecordWriter(
os.path.join(self._output_dir, self._episode_log_file))
# Set up a session and initialize variables.
self._sess = tf.Session(config=tf.ConfigProto(
allow_soft_placement=True))
self._agent = self._create_agent_fn(
self._sess,
self._env,
summary_writer=self._summary_writer,
eval_mode=eval_mode)
self._summary_writer.add_graph(graph=tf.get_default_graph())
self._sess.run(tf.global_variables_initializer())
self._sess.run(tf.local_variables_initializer())
def initialize_session(self):
"""Initializes a tf Session."""
if ENABLE_TF_OPTIMIZATIONS:
self.sess = tf.Session()
else:
rewriter_config = rewriter_config_pb2.RewriterConfig(
disable_model_pruning=True,
constant_folding=rewriter_config_pb2.RewriterConfig.OFF,
arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF,
remapping=rewriter_config_pb2.RewriterConfig.OFF,
shape_optimization=rewriter_config_pb2.RewriterConfig.OFF,
dependency_optimization=rewriter_config_pb2.RewriterConfig.OFF,
function_optimization=rewriter_config_pb2.RewriterConfig.OFF,
layout_optimizer=rewriter_config_pb2.RewriterConfig.OFF,
loop_optimization=rewriter_config_pb2.RewriterConfig.OFF,
memory_optimization=rewriter_config_pb2.RewriterConfig.
NO_MEM_OPT)
graph_options = tf.GraphOptions(rewrite_options=rewriter_config)
session_config = tf.ConfigProto(graph_options=graph_options)
self.sess = tf.Session(config=session_config)
# Restore or initialize the variables.
self.sess.run(tf.global_variables_initializer())
self.sess.run(tf.local_variables_initializer())
if self.learner_config.checkpoint_for_eval:
# Requested a specific checkpoint.
self.saver.restore(self.sess,
self.learner_config.checkpoint_for_eval)
tf.logging.info('Restored checkpoint: %s' %
self.learner_config.checkpoint_for_eval)
else:
# Continue from the latest checkpoint if one exists.
# This handles fault-tolerance.
latest_checkpoint = None
if self.checkpoint_dir is not None:
latest_checkpoint = tf.train.latest_checkpoint(
self.checkpoint_dir)
if latest_checkpoint:
self.saver.restore(self.sess, latest_checkpoint)
tf.logging.info('Restored checkpoint: %s' % latest_checkpoint)
else:
tf.logging.info('No previous checkpoint.')
self.sess.run(tf.global_variables_initializer())
self.sess.run(tf.local_variables_initializer())
# For episodic models, potentially use pretrained weights at the start of
# training. If this happens it will overwrite the embedding weights, but
# taking care to not restore the Adam parameters.
if self.learner_config.pretrained_checkpoint and not self.sess.run(
tf.train.get_global_step()):
self.saver.restore(self.sess,
self.learner_config.pretrained_checkpoint)
tf.logging.info('Restored checkpoint: %s' %
self.learner_config.pretrained_checkpoint)
# We only want the embedding weights of the checkpoint we just restored.
# So we re-initialize everything that's not an embedding weight. Also,
# since this episodic finetuning procedure is a different optimization
# problem than the original training of the baseline whose embedding
# weights are re-used, we do not reload ADAM's variables and instead learn
# them from scratch.
vars_to_reinit, embedding_var_names, vars_to_reinit_names = [], [], []
for var in tf.global_variables():
if (any(keyword in var.name for keyword in EMBEDDING_KEYWORDS)
and 'adam' not in var.name.lower()):
embedding_var_names.append(var.name)
continue
vars_to_reinit.append(var)
vars_to_reinit_names.append(var.name)
tf.logging.info('Initializing all variables except for %s.' %
embedding_var_names)
self.sess.run(tf.variables_initializer(vars_to_reinit))
tf.logging.info('Re-initialized vars %s.' % vars_to_reinit_names)