def train(self, model, agent_model, model_name, ckpt_dir='./checkpoint', summary_dir='./summary', test_result_dir='./test_result', learning_rate=0.5, max_epoch=60, pretrain_agent_epoch=1, pretrain_model=None, test_epoch=1, optimizer=tf.train.GradientDescentOptimizer):
    print('Start training...')
    self.model = model(self.train_data_loader, self.train_data_loader.batch_size, self.train_data_loader.max_length)
    model_optimizer = optimizer(learning_rate)
    grads = model_optimizer.compute_gradients(self.model.loss())
    self.train_op = model_optimizer.apply_gradients(grads)
    self.agent_model = agent_model(self.train_data_loader, self.train_data_loader.batch_size, self.train_data_loader.max_length)
    agent_optimizer = optimizer(learning_rate)
    agent_grads = agent_optimizer.compute_gradients(self.agent_model.loss())
    self.agent_train_op = agent_optimizer.apply_gradients(agent_grads)
    self.sess = tf.Session()
    summary_writer = tf.summary.FileWriter(summary_dir, self.sess.graph)
    saver = tf.train.Saver(max_to_keep=None)
    if (pretrain_model is None):
        self.sess.run(tf.global_variables_initializer())
    else:
        saver.restore(self.sess, pretrain_model)
    self.pretrain_main_model(max_epoch=5)
    self.pretrain_agent_model(max_epoch=1)
    tot_delete = 0
    batch_count = 0
    instance_count = 0
    reward = 0.0
    best_metric = 0
    best_prec = None
    best_recall = None
    not_best_count = 0
    for epoch in range(max_epoch):
        print((('###### Epoch ' + str(epoch)) + ' ######'))
        tot_correct = 0
        tot_not_na_correct = 0
        tot = 0
        tot_not_na = 0
        i = 0
        time_sum = 0
        batch_stack = []
        for (i, batch_data) in enumerate(self.train_data_loader):
            batch_data['agent_label'] = (batch_data['ins_rel'] + 0)
            batch_data['agent_label'][(batch_data['agent_label'] > 0)] = 1
            batch_stack.append(batch_data)
            iter_logit = self.agent_one_step(self.sess, self.agent_model, batch_data, [self.agent_model.train_logit()])[0]
            action_result = iter_logit.argmax((- 1))
            batch_delete = np.sum(np.logical_and((batch_data['ins_rel'] != 0), (action_result == 0)))
            batch_data['ins_rel'][(action_result == 0)] = 0
            iter_loss = self.one_step(self.sess, self.model, batch_data, [self.model.loss()])[0]
            reward += iter_loss
            tot_delete += batch_delete
            batch_count += 1
            alpha = 0.1
            if (batch_count == 100):
                reward = (reward / float(batch_count))
                average_loss = reward
                reward = (- math.log((1 - (math.e ** (- reward)))))
                sys.stdout.write(('tot delete : %f | reward : %f | average loss : %f\r' % (tot_delete, reward, average_loss)))
                sys.stdout.flush()
                for batch_data in batch_stack:
                    self.agent_one_step(self.sess, self.agent_model, batch_data, [self.agent_train_op], weights=(reward * alpha))
                batch_count = 0
                reward = 0
                tot_delete = 0
                batch_stack = []
            i += 1
        for (i, batch_data) in enumerate(self.train_data_loader):
            batch_data['agent_label'] = (batch_data['ins_rel'] + 0)
            batch_data['agent_label'][(batch_data['agent_label'] > 0)] = 1
            time_start = time.time()
            iter_logit = self.agent_one_step(self.sess, self.agent_model, batch_data, [self.agent_model.train_logit()])[0]
            action_result = iter_logit.argmax((- 1))
            batch_data['ins_rel'][(action_result == 0)] = 0
            (iter_loss, iter_logit, _train_op) = self.agent_one_step(self.sess, self.agent_model, batch_data, [self.agent_model.loss(), self.agent_model.train_logit(), self.agent_train_op])
            time_end = time.time()
            t = (time_end - time_start)
            time_sum += t
            iter_output = iter_logit.argmax((- 1))
            if (tot_not_na > 0):
                sys.stdout.write(('epoch %d step %d time %.2f | loss: %f, not NA accuracy: %f, accuracy: %f\r' % (epoch, i, t, iter_loss, (float(tot_not_na_correct) / tot_not_na), (float(tot_correct) / tot))))
                sys.stdout.flush()
            i += 1
        print(('\nAverage iteration time: %f' % (time_sum / i)))
        if (((epoch + 1) % test_epoch) == 0):
            metric = self.test(model)
            if (metric > best_metric):
                best_metric = metric
                best_prec = self.cur_prec
                best_recall = self.cur_recall
                print('Best model, storing...')
                if (not os.path.isdir(ckpt_dir)):
                    os.mkdir(ckpt_dir)
                path = saver.save(self.sess, os.path.join(ckpt_dir, model_name))
                print('Finish storing')
                not_best_count = 0
            else:
                not_best_count += 1
        if (not_best_count >= 20):
            break
    print('######')
    print(('Finish training ' + model_name))
    print(('Best epoch auc = %f' % best_metric))
    if ((not (best_prec is None)) and (not (best_recall is None))):
        if (not os.path.isdir(test_result_dir)):
            os.mkdir(test_result_dir)
        np.save(os.path.join(test_result_dir, (model_name + '_x.npy')), best_recall)
        np.save(os.path.join(test_result_dir, (model_name + '_y.npy')), best_prec)
def train(self, model, agent_model, model_name, ckpt_dir='./checkpoint', summary_dir='./summary', test_result_dir='./test_result', learning_rate=0.5, max_epoch=60, pretrain_agent_epoch=1, pretrain_model=None, test_epoch=1, optimizer=tf.train.GradientDescentOptimizer):
    print('Start training...')
    self.model = model(self.train_data_loader, self.train_data_loader.batch_size, self.train_data_loader.max_length)
    model_optimizer = optimizer(learning_rate)
    grads = model_optimizer.compute_gradients(self.model.loss())
    self.train_op = model_optimizer.apply_gradients(grads)
    self.agent_model = agent_model(self.train_data_loader, self.train_data_loader.batch_size, self.train_data_loader.max_length)
    agent_optimizer = optimizer(learning_rate)
    agent_grads = agent_optimizer.compute_gradients(self.agent_model.loss())
    self.agent_train_op = agent_optimizer.apply_gradients(agent_grads)
    self.sess = tf.Session()
    summary_writer = tf.summary.FileWriter(summary_dir, self.sess.graph)
    saver = tf.train.Saver(max_to_keep=None)
    if (pretrain_model is None):
        self.sess.run(tf.global_variables_initializer())
    else:
        saver.restore(self.sess, pretrain_model)
    self.pretrain_main_model(max_epoch=5)
    self.pretrain_agent_model(max_epoch=1)
    tot_delete = 0
    batch_count = 0
    instance_count = 0
    reward = 0.0
    best_metric = 0
    best_prec = None
    best_recall = None
    not_best_count = 0
    for epoch in range(max_epoch):
        print((('###### Epoch ' + str(epoch)) + ' ######'))
        tot_correct = 0
        tot_not_na_correct = 0
        tot = 0
        tot_not_na = 0
        i = 0
        time_sum = 0
        batch_stack = []
        for (i, batch_data) in enumerate(self.train_data_loader):
            batch_data['agent_label'] = (batch_data['ins_rel'] + 0)
            batch_data['agent_label'][(batch_data['agent_label'] > 0)] = 1
            batch_stack.append(batch_data)
            iter_logit = self.agent_one_step(self.sess, self.agent_model, batch_data, [self.agent_model.train_logit()])[0]
            action_result = iter_logit.argmax((- 1))
            batch_delete = np.sum(np.logical_and((batch_data['ins_rel'] != 0), (action_result == 0)))
            batch_data['ins_rel'][(action_result == 0)] = 0
            iter_loss = self.one_step(self.sess, self.model, batch_data, [self.model.loss()])[0]
            reward += iter_loss
            tot_delete += batch_delete
            batch_count += 1
            alpha = 0.1
            if (batch_count == 100):
                reward = (reward / float(batch_count))
                average_loss = reward
                reward = (- math.log((1 - (math.e ** (- reward)))))
                sys.stdout.write(('tot delete : %f | reward : %f | average loss : %f\r' % (tot_delete, reward, average_loss)))
                sys.stdout.flush()
                for batch_data in batch_stack:
                    self.agent_one_step(self.sess, self.agent_model, batch_data, [self.agent_train_op], weights=(reward * alpha))
                batch_count = 0
                reward = 0
                tot_delete = 0
                batch_stack = []
            i += 1
        for (i, batch_data) in enumerate(self.train_data_loader):
            batch_data['agent_label'] = (batch_data['ins_rel'] + 0)
            batch_data['agent_label'][(batch_data['agent_label'] > 0)] = 1
            time_start = time.time()
            iter_logit = self.agent_one_step(self.sess, self.agent_model, batch_data, [self.agent_model.train_logit()])[0]
            action_result = iter_logit.argmax((- 1))
            batch_data['ins_rel'][(action_result == 0)] = 0
            (iter_loss, iter_logit, _train_op) = self.agent_one_step(self.sess, self.agent_model, batch_data, [self.agent_model.loss(), self.agent_model.train_logit(), self.agent_train_op])
            time_end = time.time()
            t = (time_end - time_start)
            time_sum += t
            iter_output = iter_logit.argmax((- 1))
            if (tot_not_na > 0):
                sys.stdout.write(('epoch %d step %d time %.2f | loss: %f, not NA accuracy: %f, accuracy: %f\r' % (epoch, i, t, iter_loss, (float(tot_not_na_correct) / tot_not_na), (float(tot_correct) / tot))))
                sys.stdout.flush()
            i += 1
        print(('\nAverage iteration time: %f' % (time_sum / i)))
        if (((epoch + 1) % test_epoch) == 0):
            metric = self.test(model)
            if (metric > best_metric):
                best_metric = metric
                best_prec = self.cur_prec
                best_recall = self.cur_recall
                print('Best model, storing...')
                if (not os.path.isdir(ckpt_dir)):
                    os.mkdir(ckpt_dir)
                path = saver.save(self.sess, os.path.join(ckpt_dir, model_name))
                print('Finish storing')
                not_best_count = 0
            else:
                not_best_count += 1
        if (not_best_count >= 20):
            break
    print('######')
    print(('Finish training ' + model_name))
    print(('Best epoch auc = %f' % best_metric))
    if ((not (best_prec is None)) and (not (best_recall is None))):
        if (not os.path.isdir(test_result_dir)):
            os.mkdir(test_result_dir)
        np.save(os.path.join(test_result_dir, (model_name + '_x.npy')), best_recall)
        np.save(os.path.join(test_result_dir, (model_name + '_y.npy')), best_prec)