def multi_worker_selfplay(hparams,
identity,
scope=None,
target_session='',
is_chief=True,
ps_tasks=0,
num_workers=1,
jobid=0,
startup_delay_steps=0):
"""This is the multi worker selfplay, mostly used for self play
distributed training.
identity is used.
"""
immutable_model_reload_freq = hparams.immutable_model_reload_freq
# 1. models and summary writer
model_creator = diag_model.Model
extra_args = model_helper.ExtraArgs(
single_cell_fn=None,
model_device_fn=tf.train.replica_device_setter(ps_tasks),
attention_mechanism_fn=None)
mutable_model = model_helper.create_selfplay_model(model_creator,
is_mutable=True,
num_workers=num_workers,
jobid=jobid,
hparams=hparams,
scope=scope,
extra_args=extra_args)
immutable_hparams = copy.deepcopy(hparams)
immutable_hparams.num_gpus = 0
immutable_model = model_helper.create_selfplay_model(
model_creator,
is_mutable=False,
num_workers=num_workers,
jobid=jobid,
hparams=immutable_hparams,
scope=scope)
if hparams.self_play_immutable_gpu:
print('using GPU for immutable')
immutable_sess = tf.Session(
graph=immutable_model.graph,
config=tf.ConfigProto(allow_soft_placement=True))
else:
print('not using GPU for immutable')
immutable_sess = tf.Session(graph=immutable_model.graph,
config=tf.ConfigProto(
allow_soft_placement=True,
device_count={'GPU': 0}))
immutable_model, immutable_sess = load_self_play_model(
immutable_model, immutable_sess, 'immutable',
hparams.self_play_pretrain_dir, hparams.out_dir)
global_step = immutable_model.model.global_step.eval(
session=immutable_sess)
if is_chief:
ckpt = tf.train.latest_checkpoint(hparams.out_dir)
if not ckpt:
print('global_step, saving pretrain model to hparams.out_dir',
global_step, hparams.out_dir)
immutable_model.model.saver.save( # this is the prevent adam error
immutable_sess,
os.path.join(hparams.out_dir, 'dialogue.ckpt'),
global_step=global_step)
print('save finished')
if is_chief:
summary_writer_path = os.path.join(
hparams.out_dir, identity + task_SP_DISTRIBUTED + '_log')
summary_writer = tf.summary.FileWriter(summary_writer_path,
mutable_model.graph)
print('summary writer established at', summary_writer_path)
else:
summary_writer = None
# 2. supervisor and sessions
sv = tf.train.Supervisor(
graph=mutable_model.graph,
is_chief=is_chief,
saver=mutable_model.model.saver,
save_model_secs=0, # disable automatic save checkpoints
summary_op=None,
logdir=hparams.out_dir,
checkpoint_basename='dialogue.ckpt')
mutable_config = utils.get_config_proto(
log_device_placement=hparams.log_device_placement,
allow_soft_placement=True)
mutable_config.device_count['GPU'] = hparams.num_gpus
mutable_sess = sv.prepare_or_wait_for_session(target_session,
config=mutable_config)
# 3. additiona preparations
global_step = mutable_model.model.global_step.eval(session=mutable_sess)
while global_step < (jobid * (jobid + 1) * startup_delay_steps / 2):
time.sleep(1)
global_step = mutable_model.model.global_step.eval(
session=mutable_sess)
# save first model
if is_chief:
print('saveing the first checkpoint to', hparams.out_dir)
mutable_model.model.saver.save(mutable_sess,
os.path.join(hparams.out_dir,
'dialogue.ckpt'),
global_step=global_step)
last_save_step = global_step
# Read data
selfplay_data = dialogue_utils.load_data(hparams.self_play_train_data)
selfplay_kb = dialogue_utils.load_data(hparams.self_play_train_kb)
dialogue = SelfplayDialogue(mutable_model,
immutable_model,
mutable_sess,
immutable_sess,
hparams.max_dialogue_turns,
hparams.train_threadhold,
hparams.start_of_turn1,
hparams.start_of_turn2,
hparams.end_of_dialogue,
summary_writer=summary_writer,
dialogue_mode=task_SP_DISTRIBUTED,
hparams=hparams)
# 4. main loop
last_immmutable_model_reload = global_step
last_save_step = global_step
batch_size = dialogue.batch_size
assert batch_size <= len(selfplay_data)
# this is the start point of the self-play data. force shuffling at the beginning
i = len(selfplay_data)
train_stats = [0, 0]
while global_step < hparams.num_self_play_train_steps:
# a. reload immutable model, muttable will be automated managed by supervisor
if immutable_model_reload_freq > 0 and global_step - last_immmutable_model_reload > immutable_model_reload_freq:
immutable_model, immutable_sess = load_self_play_model(
immutable_model, immutable_sess, 'immutable',
hparams.self_play_pretrain_dir, hparams.out_dir)
last_immmutable_model_reload = global_step
# b. possiblely flip between speakers (or roll out models),
# based on either a random policy or by step counts
agent1, agent2, mutable_agent_index = dialogue.flip_agent(
(mutable_model, mutable_sess, dialogue.mutable_handles),
(immutable_model, immutable_sess, dialogue.immutable_handles))
train_stats[mutable_agent_index] += 1
# read selfplay data
start_time = time.time()
if i * batch_size + batch_size > len(selfplay_data): # reacehd the end
input_data = zip(selfplay_data, selfplay_kb)
random.shuffle(input_data) # random shuffle input data
i = 0
selfplay_data, selfplay_kb = zip(*input_data)
start_ind, end_ind = i * batch_size, i * batch_size + batch_size
batch_data, batch_kb = selfplay_data[start_ind:end_ind], selfplay_kb[
start_ind:end_ind]
train_example, _, _ = dialogue.talk(hparams.max_dialogue_len,
batch_data, batch_kb, agent1,
agent2, batch_size, global_step)
possible_global_step = dialogue.maybe_train(train_example,
mutable_agent_index,
global_step,
force=True)
if possible_global_step:
global_step = possible_global_step
if is_chief and global_step - last_save_step > hparams.self_play_dist_save_freq:
mutable_model.model.saver.save(mutable_sess,
os.path.join(
hparams.out_dir,
'dialogue.ckpt'),
global_step=global_step)
last_save_step = global_step
end_time = time.time()
if is_chief:
utils.add_summary(summary_writer, global_step,
task_SP_DISTRIBUTED + '_' + 'time',
end_time - start_time)
utils.add_summary(summary_writer, global_step,
task_SP_DISTRIBUTED + '_' + 'train_ratio',
train_stats[0] * 1.0 / (train_stats[1] + 0.1))
i += 1
if is_chief:
summary_writer.close()
mutable_sess.close()
immutable_sess.close()
def self_play_eval_fn(hparams,
identity,
num_workers=1,
jobid=0,
scope=None,
target_session=''):
"""This is the single worker self play.
Mostly used for self play
evaluation. identity is used here to distinguish between workers.
"""
model_creator = diag_model.Model
mutable_model = model_helper.create_selfplay_model(
model_creator,
True, # mutable is True
num_workers,
jobid,
hparams=hparams,
scope=scope)
immutable_model = model_helper.create_selfplay_model(
model_creator,
False, # mutable is False
num_workers,
jobid,
hparams=hparams,
scope=scope)
mutable_sess = tf.Session(graph=mutable_model.graph,
config=tf.ConfigProto(
allow_soft_placement=True,
device_count={'GPU': hparams.num_gpus}))
immutable_sess = tf.Session(graph=immutable_model.graph,
config=tf.ConfigProto(
allow_soft_placement=True,
device_count={'GPU': hparams.num_gpus}))
# number of steps per external eval
steps_per_external_eval = 10
# force conducting a self play at the beginning
last_external_eval_step = -1 * steps_per_external_eval
print('hparams.self_play_pretrain_dir=', hparams.self_play_pretrain_dir)
print('steps_per_external_eval=', steps_per_external_eval)
writer_path = os.path.join(hparams.out_dir,
identity + hparams.task_type + '_log')
summary_writer = tf.summary.FileWriter(writer_path, mutable_sess.graph)
print('summary_writer estabilished at', writer_path)
# waiting for checkpoints and loop forever
latest_ckpt = None
while True:
latest_ckpt = tf.contrib.training.wait_for_new_checkpoint(
hparams.out_dir, latest_ckpt)
print('got checkpoint', latest_ckpt)
# get the global_step variable first
with mutable_model.graph.as_default():
# first initialize to avoid encountering missing component for adam optimizer
_, global_step = model_helper.create_or_load_model(
mutable_model.model, hparams.out_dir, mutable_sess,
hparams.task_type)
# valid evaluation step
if (not hparams.eval_forever) or (global_step - last_external_eval_step
>= steps_per_external_eval):
# if eval_forever is disabled, we will do one selfplay evalation
# otherwise, we will wait until certain number of timesteps are elapsed.
last_external_eval_step = global_step
print('do single worker evaluation')
single_worker_selfplay(mutable_model, immutable_model,
mutable_sess, immutable_sess,
hparams.self_play_eval_data,
hparams.self_play_eval_kb, global_step,
hparams, summary_writer)
else:
print('Wait until steps_per_external_eval is reached.',
global_step, last_external_eval_step,
steps_per_external_eval)
if not hparams.eval_forever:
break # if eval_foever is disabled, we only evaluate once
mutable_sess.close()
immutable_sess.close()