|
|
import os |
|
|
import tensorflow as tf |
|
|
import tf_utils as tfu |
|
|
from tqdm import tqdm |
|
|
from abc import ABC, abstractmethod |
|
|
from utils import RANDOM |
|
|
|
|
|
|
|
|
class TfContext: |
|
|
def __init__(self, |
|
|
sess, |
|
|
saver, |
|
|
learn_rate_op): |
|
|
self.sess = sess |
|
|
self.saver = saver |
|
|
self.learn_rate_op = learn_rate_op |
|
|
self.epoch = 0 |
|
|
|
|
|
|
|
|
class GraphPartBase(ABC): |
|
|
def __init__(self, for_usage, global_settings, current_settings, optimiser, reset_optimiser, key, metric_names): |
|
|
self.key = key |
|
|
self.global_settings = global_settings |
|
|
self.filler = global_settings['filler'] |
|
|
self.main_metric_name = current_settings['main_metric_type'] |
|
|
self.settings = current_settings |
|
|
self.for_usage = for_usage |
|
|
self.optimiser = optimiser |
|
|
self.reset_optimiser = reset_optimiser |
|
|
self.metric_names = metric_names |
|
|
self.max_word_size = global_settings['max_word_size'] |
|
|
self.checkpoints_keep = 10000 |
|
|
self.chars_count = len(global_settings['chars']) + 1 |
|
|
self.dataset_path = global_settings['dataset_path'] |
|
|
self.grammemes_count = len(global_settings['grammemes_types']) |
|
|
self.main_classes = global_settings['main_classes'] |
|
|
self.main_classes_count = len(self.main_classes) |
|
|
self.metrics_reset = [] |
|
|
self.metrics_update = [] |
|
|
self.devices_metrics = {metr: [] for metr in self.metric_names} |
|
|
self.main_scope_name = key.title() |
|
|
self.save_path = global_settings['save_path'] |
|
|
self.dev_grads = [] |
|
|
self.losses = [] |
|
|
self.devices = global_settings['train_devices'] |
|
|
self.devices_count = len(self.devices) |
|
|
self.dataset_path = global_settings['dataset_path'] |
|
|
self.xs = [] |
|
|
self.x_seq_lens = [] |
|
|
self.prints = [] |
|
|
self.main_cls_dic = self.global_settings['main_classes'] |
|
|
self.learn_rate_val = self.settings['learn_rate'] |
|
|
self.best_model_metric = None |
|
|
self.best_epoch = None |
|
|
self.init_checkpoint = None |
|
|
|
|
|
def train(self, tc): |
|
|
return_step = 0 |
|
|
trains = self.__load_dataset__('train') |
|
|
valids = self.__load_dataset__('valid') |
|
|
self.best_model_metric = self.__valid_loop__(tc, valids) |
|
|
self.best_epoch = -1 |
|
|
while True: |
|
|
tqdm.write(self.filler) |
|
|
tqdm.write(self.filler) |
|
|
tqdm.write(self.main_scope_name) |
|
|
|
|
|
train_main_metric = self.__train_loop__(tc, trains) |
|
|
valid_main_metric = self.__valid_loop__(tc, valids) |
|
|
|
|
|
tqdm.write(f"Epoch {tc.epoch} Train {self.main_metric_name}: {train_main_metric} Validation {self.main_metric_name}: {valid_main_metric}") |
|
|
need_decay = False |
|
|
delta = self.__calc_metric_delta__(self.best_model_metric, valid_main_metric) |
|
|
|
|
|
if delta > 0: |
|
|
if delta < self.settings['stop_main_metric_delta']: |
|
|
tqdm.write(f"{self.main_metric_name} delta is less then min value") |
|
|
need_decay = True |
|
|
else: |
|
|
return_step = 0 |
|
|
self.best_model_metric = valid_main_metric |
|
|
self.best_epoch = tc.epoch |
|
|
tc.saver.save(tc.sess, self.save_path, tc.epoch) |
|
|
tc.epoch += 1 |
|
|
else: |
|
|
tqdm.write("Best epoch is better then current") |
|
|
tc.sess.run(self.reset_optimiser) |
|
|
need_decay = True |
|
|
self.__restore_best_epoch__(tc) |
|
|
|
|
|
if not need_decay: |
|
|
continue |
|
|
|
|
|
if return_step == self.settings['return_step']: |
|
|
self.__decay_params__() |
|
|
if self.learn_rate_val < self.settings['min_learn_rate']: |
|
|
tqdm.write(f"Learning rate {self.learn_rate_val} is less then min learning rate") |
|
|
finish = self.__before_finish__() |
|
|
if finish: |
|
|
break |
|
|
return_step = 0 |
|
|
else: |
|
|
RANDOM.shuffle(trains) |
|
|
tqdm.write(f"Return step increased") |
|
|
return_step += 1 |
|
|
|
|
|
return self.best_epoch, self.best_model_metric |
|
|
|
|
|
def __train_loop__(self, tc, trains): |
|
|
tc.sess.run(self.metrics_reset) |
|
|
for item in tqdm(trains, desc=f"Train, epoch {tc.epoch}"): |
|
|
launch = [self.optimize] |
|
|
launch.extend(self.metrics_update) |
|
|
if len(self.prints): |
|
|
launch.extend(self.prints) |
|
|
|
|
|
feed_dic = self.__create_feed_dict__('train', item) |
|
|
feed_dic[tc.learn_rate_op] = self.learn_rate_val |
|
|
tc.sess.run(launch, feed_dic) |
|
|
|
|
|
train_main_metric = self.__write_metrics_report__(tc.sess, "Train") |
|
|
return train_main_metric |
|
|
|
|
|
def __valid_loop__(self, tc, valids): |
|
|
tc.sess.run(self.metrics_reset) |
|
|
for item in tqdm(valids, desc=f"Validation, epoch {tc.epoch}"): |
|
|
launch = [] |
|
|
launch.extend(self.metrics_update) |
|
|
if len(self.prints): |
|
|
launch.extend(self.prints) |
|
|
feed_dic = self.__create_feed_dict__('valid', item) |
|
|
tc.sess.run(launch, feed_dic) |
|
|
|
|
|
valid_main_metric = self.__write_metrics_report__(tc.sess, "Valid") |
|
|
return valid_main_metric |
|
|
|
|
|
def __calc_metric_delta__(self, best_model_metric, cur_model_metric): |
|
|
delta = best_model_metric - cur_model_metric |
|
|
if self.main_metric_name != "Loss": |
|
|
delta = -delta |
|
|
return delta |
|
|
|
|
|
def __before_finish__(self): |
|
|
return True |
|
|
|
|
|
def __decay_params__(self): |
|
|
self.learn_rate_val = self.learn_rate_val * self.settings['learn_rate_decay_step'] |
|
|
tqdm.write(f"Learning rate decayed. New value: {self.learn_rate_val}") |
|
|
|
|
|
def __restore_best_epoch__(self, tc): |
|
|
if self.best_epoch == -1 and tc.epoch == 0: |
|
|
tqdm.write(f"Restoring from init_checkpoint {self.best_epoch}") |
|
|
self.restore(tc.sess, self.init_checkpoint) |
|
|
elif self.best_epoch == -1: |
|
|
tqdm.write(f"Restoring best epoch {tc.epoch}") |
|
|
self.restore(tc.sess, os.path.join(self.save_path, f"-{tc.epoch}")) |
|
|
else: |
|
|
tqdm.write(f"Restoring best epoch {self.best_epoch}") |
|
|
self.restore(tc.sess, os.path.join(self.save_path, f"-{self.best_epoch}")) |
|
|
|
|
|
def build_graph_end(self): |
|
|
with tf.variable_scope(self.main_scope_name, reuse=tf.AUTO_REUSE) as scope: |
|
|
self.metrics = { |
|
|
metr: tf.reduce_mean(self.devices_metrics[metr], name=metr) |
|
|
for metr in self.devices_metrics |
|
|
} |
|
|
if not self.for_usage: |
|
|
self.grads = tfu.average_gradients(self.dev_grads) |
|
|
if self.settings['clip_grads']: |
|
|
self.grads = [(tf.clip_by_value(grad, -1., 1.), var) for grad, var in self.grads] |
|
|
|
|
|
self.optimize = self.optimiser.apply_gradients(self.grads, name='Optimize') |
|
|
self.loss = tf.reduce_sum(self.losses, name='GlobalLoss') |
|
|
|
|
|
def build_graph_for_device(self, *args): |
|
|
with tf.variable_scope(self.main_scope_name, reuse=tf.AUTO_REUSE) as scope: |
|
|
self.__build_graph_for_device__(*args) |
|
|
|
|
|
def restore(self, sess, check_point): |
|
|
try: |
|
|
vars = [ |
|
|
var |
|
|
for var in tf.global_variables(f"{self.main_scope_name}/") |
|
|
if "Adam" not in var.name |
|
|
] |
|
|
saver = tf.train.Saver(var_list=vars) |
|
|
saver.restore(sess, check_point) |
|
|
self.init_checkpoint = check_point |
|
|
tqdm.write(f"Restoration for graph part '{self.key}', scope {self.main_scope_name} success") |
|
|
except Exception as ex: |
|
|
tqdm.write(f"Restoration for graph part '{self.key}', scope {self.main_scope_name} failed. Error: {ex}") |
|
|
|
|
|
def __write_metrics_report__(self, sess, step_name): |
|
|
tqdm.write('') |
|
|
launch_results = sess.run(self.metrics) |
|
|
result = [f"{step_name} metrics: "] |
|
|
|
|
|
for index, metr in enumerate(self.metrics): |
|
|
result.append('{:>8}'.format(self.metric_names[index])) |
|
|
result.append("=") |
|
|
result.append("{0:.7f}".format(launch_results[metr])) |
|
|
result.append(" ") |
|
|
|
|
|
result = "".join(result) |
|
|
tqdm.write(result) |
|
|
return launch_results[self.main_metric_name] |
|
|
|
|
|
def create_mean_metric(self, metric_index, values): |
|
|
metr_epoch_loss, metr_update, metr_reset = tfu.create_reset_metric( |
|
|
tf.metrics.mean, |
|
|
self.metric_names[metric_index], |
|
|
values |
|
|
) |
|
|
self.metrics_reset.append(metr_reset) |
|
|
self.metrics_update.append(metr_update) |
|
|
self.devices_metrics[self.metric_names[metric_index]].append(metr_epoch_loss) |
|
|
|
|
|
def create_accuracy_metric(self, metric_index, labels, predictions): |
|
|
metr_epoch_loss, metr_update, metr_reset = tfu.create_reset_metric( |
|
|
tf.metrics.accuracy, |
|
|
self.metric_names[metric_index], |
|
|
labels=labels, |
|
|
predictions=predictions |
|
|
) |
|
|
self.metrics_reset.append(metr_reset) |
|
|
self.metrics_update.append(metr_update) |
|
|
self.devices_metrics[self.metric_names[metric_index]].append(metr_epoch_loss) |
|
|
|
|
|
def __create_feed_dict__(self, op_name, item): |
|
|
feed_dic = {} |
|
|
for dev_num, batch in enumerate(item): |
|
|
feed_dic[self.xs[dev_num]] = batch['x'] |
|
|
feed_dic[self.x_seq_lens[dev_num]] = batch['x_seq_len'] |
|
|
self.__update_feed_dict__(op_name, feed_dic, batch, dev_num) |
|
|
|
|
|
return feed_dic |
|
|
|
|
|
@abstractmethod |
|
|
def __update_feed_dict__(self, op_name, feed_dict, batch, dev_num): |
|
|
pass |
|
|
|
|
|
@abstractmethod |
|
|
def __build_graph_for_device__(self, *args): |
|
|
pass |
|
|
|
|
|
@abstractmethod |
|
|
def __load_dataset__(self, operation_name): |
|
|
return [] |
|
|
|
