def train(ps_device):
    (train_input_fn, record_info_dict) = data_utils.get_input_fn(tfrecord_dir=FLAGS.record_info_dir, split='train', bsz_per_host=FLAGS.train_batch_size, seq_len=FLAGS.seq_len, reuse_len=FLAGS.reuse_len, bi_data=FLAGS.bi_data, num_hosts=1, num_core_per_host=1, perm_size=FLAGS.perm_size, mask_alpha=FLAGS.mask_alpha, mask_beta=FLAGS.mask_beta, uncased=FLAGS.uncased, num_passes=FLAGS.num_passes, use_bfloat16=FLAGS.use_bfloat16, num_predict=FLAGS.num_predict)
    tf.logging.info('num of batches {}'.format(record_info_dict['num_batch']))
    bsz_per_core = (FLAGS.train_batch_size // FLAGS.num_core_per_host)
    params = {'batch_size': FLAGS.train_batch_size}
    train_set = train_input_fn(params)
    example = train_set.make_one_shot_iterator().get_next()
    if (FLAGS.num_core_per_host > 1):
        examples = [{} for _ in range(FLAGS.num_core_per_host)]
        for key in example.keys():
            vals = tf.split(example[key], FLAGS.num_core_per_host, 0)
            for device_id in range(FLAGS.num_core_per_host):
                examples[device_id][key] = vals[device_id]
    else:
        examples = [example]
    (tower_mems, tower_losses, tower_new_mems, tower_grads_and_vars) = ([], [], [], [])
    for i in range(FLAGS.num_core_per_host):
        reuse = (True if (i > 0) else None)
        with tf.device(assign_to_gpu(i, ps_device)), tf.variable_scope(tf.get_variable_scope(), reuse=reuse):
            mems_i = {}
            if FLAGS.mem_len:
                mems_i['mems'] = create_mems_tf(bsz_per_core)
            (loss_i, new_mems_i, grads_and_vars_i) = single_core_graph(is_training=True, features=examples[i], mems=mems_i)
            tower_mems.append(mems_i)
            tower_losses.append(loss_i)
            tower_new_mems.append(new_mems_i)
            tower_grads_and_vars.append(grads_and_vars_i)
    if (len(tower_losses) > 1):
        loss = (tf.add_n(tower_losses) / len(tower_losses))
        grads_and_vars = average_grads_and_vars(tower_grads_and_vars)
    else:
        loss = tower_losses[0]
        grads_and_vars = tower_grads_and_vars[0]
    (train_op, learning_rate, gnorm) = model_utils.get_train_op(FLAGS, None, grads_and_vars=grads_and_vars)
    global_step = tf.train.get_global_step()
    tower_mems_np = []
    for i in range(FLAGS.num_core_per_host):
        mems_i_np = {}
        for key in tower_mems[i].keys():
            mems_i_np[key] = initialize_mems_np(bsz_per_core)
        tower_mems_np.append(mems_i_np)
    saver = tf.train.Saver()
    gpu_options = tf.GPUOptions(allow_growth=True)
    model_utils.init_from_checkpoint(FLAGS, global_vars=True)
    with tf.Session(config=tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options)) as sess:
        sess.run(tf.global_variables_initializer())
        fetches = [loss, tower_new_mems, global_step, gnorm, learning_rate, train_op]
        (total_loss, prev_step) = (0.0, (- 1))
        while True:
            feed_dict = {}
            for i in range(FLAGS.num_core_per_host):
                for key in tower_mems_np[i].keys():
                    for (m, m_np) in zip(tower_mems[i][key], tower_mems_np[i][key]):
                        feed_dict[m] = m_np
            fetched = sess.run(fetches, feed_dict=feed_dict)
            (loss_np, tower_mems_np, curr_step) = fetched[:3]
            total_loss += loss_np
            if ((curr_step > 0) and ((curr_step % FLAGS.iterations) == 0)):
                curr_loss = (total_loss / (curr_step - prev_step))
                tf.logging.info('[{}] | gnorm {:.2f} lr {:8.6f} | loss {:.2f} | pplx {:>7.2f}, bpc {:>7.4f}'.format(curr_step, fetched[(- 3)], fetched[(- 2)], curr_loss, math.exp(curr_loss), (curr_loss / math.log(2))))
                (total_loss, prev_step) = (0.0, curr_step)
            if ((curr_step > 0) and ((curr_step % FLAGS.save_steps) == 0)):
                save_path = os.path.join(FLAGS.model_dir, 'model.ckpt')
                saver.save(sess, save_path)
                tf.logging.info('Model saved in path: {}'.format(save_path))
            if (curr_step >= FLAGS.train_steps):
                break
def train(ps_device):
    (train_input_fn, record_info_dict) = data_utils.get_input_fn(tfrecord_dir=FLAGS.record_info_dir, split='train', bsz_per_host=FLAGS.train_batch_size, seq_len=FLAGS.seq_len, reuse_len=FLAGS.reuse_len, bi_data=FLAGS.bi_data, num_hosts=1, num_core_per_host=1, perm_size=FLAGS.perm_size, mask_alpha=FLAGS.mask_alpha, mask_beta=FLAGS.mask_beta, uncased=FLAGS.uncased, num_passes=FLAGS.num_passes, use_bfloat16=FLAGS.use_bfloat16, num_predict=FLAGS.num_predict)
    tf.logging.info('num of batches {}'.format(record_info_dict['num_batch']))
    bsz_per_core = (FLAGS.train_batch_size // FLAGS.num_core_per_host)
    params = {'batch_size': FLAGS.train_batch_size}
    train_set = train_input_fn(params)
    example = train_set.make_one_shot_iterator().get_next()
    if (FLAGS.num_core_per_host > 1):
        examples = [{} for _ in range(FLAGS.num_core_per_host)]
        for key in example.keys():
            vals = tf.split(example[key], FLAGS.num_core_per_host, 0)
            for device_id in range(FLAGS.num_core_per_host):
                examples[device_id][key] = vals[device_id]
    else:
        examples = [example]
    (tower_mems, tower_losses, tower_new_mems, tower_grads_and_vars) = ([], [], [], [])
    for i in range(FLAGS.num_core_per_host):
        reuse = (True if (i > 0) else None)
        with tf.device(assign_to_gpu(i, ps_device)), tf.variable_scope(tf.get_variable_scope(), reuse=reuse):
            mems_i = {}
            if FLAGS.mem_len:
                mems_i['mems'] = create_mems_tf(bsz_per_core)
            (loss_i, new_mems_i, grads_and_vars_i) = single_core_graph(is_training=True, features=examples[i], mems=mems_i)
            tower_mems.append(mems_i)
            tower_losses.append(loss_i)
            tower_new_mems.append(new_mems_i)
            tower_grads_and_vars.append(grads_and_vars_i)
    if (len(tower_losses) > 1):
        loss = (tf.add_n(tower_losses) / len(tower_losses))
        grads_and_vars = average_grads_and_vars(tower_grads_and_vars)
    else:
        loss = tower_losses[0]
        grads_and_vars = tower_grads_and_vars[0]
    (train_op, learning_rate, gnorm) = model_utils.get_train_op(FLAGS, None, grads_and_vars=grads_and_vars)
    global_step = tf.train.get_global_step()
    tower_mems_np = []
    for i in range(FLAGS.num_core_per_host):
        mems_i_np = {}
        for key in tower_mems[i].keys():
            mems_i_np[key] = initialize_mems_np(bsz_per_core)
        tower_mems_np.append(mems_i_np)
    saver = tf.train.Saver()
    gpu_options = tf.GPUOptions(allow_growth=True)
    model_utils.init_from_checkpoint(FLAGS, global_vars=True)
    with tf.Session(config=tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options)) as sess:
        sess.run(tf.global_variables_initializer())
        fetches = [loss, tower_new_mems, global_step, gnorm, learning_rate, train_op]
        (total_loss, prev_step) = (0.0, (- 1))
        while True:
            feed_dict = {}
            for i in range(FLAGS.num_core_per_host):
                for key in tower_mems_np[i].keys():
                    for (m, m_np) in zip(tower_mems[i][key], tower_mems_np[i][key]):
                        feed_dict[m] = m_np
            fetched = sess.run(fetches, feed_dict=feed_dict)
            (loss_np, tower_mems_np, curr_step) = fetched[:3]
            total_loss += loss_np
            if ((curr_step > 0) and ((curr_step % FLAGS.iterations) == 0)):
                curr_loss = (total_loss / (curr_step - prev_step))
                tf.logging.info('[{}] | gnorm {:.2f} lr {:8.6f} | loss {:.2f} | pplx {:>7.2f}, bpc {:>7.4f}'.format(curr_step, fetched[(- 3)], fetched[(- 2)], curr_loss, math.exp(curr_loss), (curr_loss / math.log(2))))
                (total_loss, prev_step) = (0.0, curr_step)
            if ((curr_step > 0) and ((curr_step % FLAGS.save_steps) == 0)):
                save_path = os.path.join(FLAGS.model_dir, 'model.ckpt')
                saver.save(sess, save_path)
                tf.logging.info('Model saved in path: {}'.format(save_path))
            if (curr_step >= FLAGS.train_steps):
                break