|
|
import os |
|
|
import shutil |
|
|
import tensorflow as tf |
|
|
from graph.gram_cls import GramCls |
|
|
from graph.main_cls import MainCls |
|
|
from graph.lemm import Lemm |
|
|
from graph.inflect import Inflect |
|
|
from graph.base import TfContext |
|
|
from utils import MyDefaultDict, CONFIG |
|
|
from tensorflow.python.tools import freeze_graph |
|
|
|
|
|
|
|
|
class RNN: |
|
|
def __init__(self, for_usage): |
|
|
self.config = CONFIG |
|
|
self.filler = self.config['filler'] |
|
|
self.checkpoints_keep = 200000 |
|
|
self.for_usage = for_usage |
|
|
self.default_config = self.config['graph_part_configs']['default'] |
|
|
self.key_configs = MyDefaultDict( |
|
|
lambda key: self.default_config, |
|
|
{ |
|
|
key: MyDefaultDict(lambda prop_key: self.default_config[prop_key], self.config['graph_part_configs'][key]) |
|
|
for key in self.config['graph_part_configs'] |
|
|
if key != 'default' |
|
|
} |
|
|
) |
|
|
self.export_path = self.config['export_path'] |
|
|
self.save_path = self.config['save_path'] |
|
|
self.model_key = self.config['model_key'] |
|
|
self.miss_steps = self.config['miss_steps'] if 'miss_steps' in self.config else [] |
|
|
self.start_char = self.config['start_token'] |
|
|
self.end_char = self.config['end_token'] |
|
|
self.gram_keys = [ |
|
|
key |
|
|
for key in sorted(self.config['grammemes_types'], key=lambda x: self.config['grammemes_types'][x]['index']) |
|
|
] |
|
|
self.main_class_k = self.config['main_class_k'] |
|
|
self.train_steps = self.config['train_steps'] |
|
|
|
|
|
if for_usage: |
|
|
self.devices = ['/cpu:0'] |
|
|
else: |
|
|
self.devices = self.config['train_devices'] |
|
|
|
|
|
if not os.path.exists(self.save_path): |
|
|
os.mkdir(self.save_path) |
|
|
|
|
|
self._build_graph() |
|
|
|
|
|
def _build_graph(self): |
|
|
self.graph = tf.Graph() |
|
|
self.checks = [] |
|
|
self.xs = [] |
|
|
self.x_seq_lens = [] |
|
|
|
|
|
self.x_inds = [] |
|
|
self.x_vals = [] |
|
|
self.x_shape = [] |
|
|
|
|
|
self.prints = [] |
|
|
|
|
|
with self.graph.as_default(), tf.device('/cpu:0'): |
|
|
self.is_training = tf.placeholder(tf.bool, name="IsTraining") |
|
|
self.learn_rate = tf.placeholder(tf.float32, name="LearningRate") |
|
|
self.batch_size = tf.placeholder(tf.int32, [], name="BatchSize") if self.for_usage else None |
|
|
self.optimiser = tf.train.AdamOptimizer(self.learn_rate) |
|
|
self.reset_optimizer = tf.variables_initializer(self.optimiser.variables()) |
|
|
self.gram_graph_parts = { |
|
|
gram: GramCls(gram, self.for_usage, self.config, self.key_configs[gram], self.optimiser, self.reset_optimizer) |
|
|
for gram in self.gram_keys |
|
|
} |
|
|
self.lem_graph_part = Lemm(self.for_usage, self.config, self.key_configs["lemm"], self.optimiser, self.reset_optimizer) |
|
|
self.main_graph_part = MainCls(self.for_usage, self.config, self.key_configs["main"], self.optimiser, self.reset_optimizer) |
|
|
self.inflect_graph_part = Inflect(self.for_usage, self.config, self.key_configs['inflect'], self.optimiser, self.reset_optimizer) |
|
|
|
|
|
for device_index, device_name in enumerate(self.devices): |
|
|
with tf.device(device_name): |
|
|
if self.for_usage: |
|
|
x_ind_pl = tf.placeholder(dtype=tf.int32, shape=(None, None), name='XIndexes') |
|
|
x_val_pl = tf.placeholder(dtype=tf.int32, shape=(None,), name='XValues') |
|
|
x_shape_pl = tf.placeholder(dtype=tf.int32, shape=(2,), name='XShape') |
|
|
x_ind = tf.dtypes.cast(x_ind_pl, dtype=tf.int64) |
|
|
x_val = tf.dtypes.cast(x_val_pl, dtype=tf.int64) |
|
|
x_shape = tf.dtypes.cast(x_shape_pl, dtype=tf.int64) |
|
|
|
|
|
x_sparse = tf.sparse.SparseTensor(x_ind, x_val, x_shape) |
|
|
x = tf.sparse.to_dense(x_sparse, default_value=self.end_char) |
|
|
self.x_inds.append(x_ind_pl) |
|
|
self.x_vals.append(x_val_pl) |
|
|
self.x_shape.append(x_shape_pl) |
|
|
else: |
|
|
x = tf.placeholder(dtype=tf.int32, shape=(None, None), name='X') |
|
|
self.xs.append(x) |
|
|
|
|
|
x_seq_len = tf.placeholder(dtype=tf.int32, shape=(None,), name='SeqLen') |
|
|
self.x_seq_lens.append(x_seq_len) |
|
|
|
|
|
for gram in self.gram_keys: |
|
|
self.gram_graph_parts[gram].build_graph_for_device(x, x_seq_len) |
|
|
|
|
|
gram_probs = [self.gram_graph_parts[gram].probs[-1] for gram in self.gram_keys] |
|
|
gram_keep_drops = [self.gram_graph_parts[gram].keep_drops[-1] for gram in self.gram_keys] |
|
|
self.main_graph_part.build_graph_for_device(x, x_seq_len, gram_probs, gram_keep_drops) |
|
|
self.prints.append(tf.print("main_result", self.main_graph_part.results[0].indices)) |
|
|
if self.for_usage: |
|
|
x_tiled = tf.contrib.seq2seq.tile_batch(x, multiplier=self.main_class_k) |
|
|
seq_len_tiled = tf.contrib.seq2seq.tile_batch(x_seq_len, multiplier=self.main_class_k) |
|
|
cls = tf.reshape(self.main_graph_part.results[0].indices, (-1,)) |
|
|
batch_size_tiled = self.batch_size * self.main_class_k |
|
|
self.lem_graph_part.build_graph_for_device(x_tiled, |
|
|
seq_len_tiled, |
|
|
batch_size_tiled, |
|
|
cls) |
|
|
self.lem_cls_result = tf.reshape(self.lem_graph_part.results[0], |
|
|
(self.batch_size, self.main_class_k, -1)) |
|
|
self.lem_graph_part.build_graph_for_device(x, |
|
|
x_seq_len, |
|
|
self.batch_size) |
|
|
self.lem_result = tf.expand_dims(self.lem_graph_part.results[1], 1) |
|
|
self.lem_class_pl = self.lem_graph_part.cls[0] |
|
|
else: |
|
|
self.lem_graph_part.build_graph_for_device(x, |
|
|
x_seq_len, |
|
|
self.batch_size) |
|
|
|
|
|
self.inflect_graph_part.build_graph_for_device(x, x_seq_len, self.batch_size) |
|
|
if self.for_usage: |
|
|
self.inflect_result = self.inflect_graph_part.results[0] |
|
|
self.inflect_x_class_pl = self.inflect_graph_part.x_cls[0] |
|
|
self.inflect_y_class_pl = self.inflect_graph_part.y_cls[0] |
|
|
|
|
|
for gram in self.gram_keys: |
|
|
self.gram_graph_parts[gram].build_graph_end() |
|
|
self.main_graph_part.build_graph_end() |
|
|
self.lem_graph_part.build_graph_end() |
|
|
self.inflect_graph_part.build_graph_end() |
|
|
self.saver = tf.train.Saver(tf.global_variables(), max_to_keep=self.checkpoints_keep) |
|
|
|
|
|
def restore(self, sess): |
|
|
latest_checkpoint = tf.train.latest_checkpoint(self.save_path) |
|
|
for gram in self.gram_keys: |
|
|
if gram not in self.config['ignore_restore']: |
|
|
self.gram_graph_parts[gram].restore(sess, latest_checkpoint) |
|
|
if self.main_graph_part.key not in self.config['ignore_restore']: |
|
|
self.main_graph_part.restore(sess, latest_checkpoint) |
|
|
if self.lem_graph_part.key not in self.config['ignore_restore']: |
|
|
self.lem_graph_part.restore(sess, latest_checkpoint) |
|
|
if self.inflect_graph_part not in self.config['ignore_restore']: |
|
|
self.inflect_graph_part.restore(sess, latest_checkpoint) |
|
|
|
|
|
if self.inflect_graph_part.settings['transfer_init']: |
|
|
self.inflect_graph_part.transfer_learning_init(sess) |
|
|
|
|
|
def train(self): |
|
|
config = tf.ConfigProto(allow_soft_placement=True) |
|
|
if not os.path.isdir(self.save_path): |
|
|
os.mkdir(self.save_path) |
|
|
|
|
|
with tf.Session(config = config, graph=self.graph) as sess: |
|
|
sess.run(tf.global_variables_initializer()) |
|
|
sess.run(tf.local_variables_initializer()) |
|
|
|
|
|
tc = TfContext(sess, self.saver, self.learn_rate) |
|
|
self.restore(sess) |
|
|
sess.run(self.reset_optimizer) |
|
|
|
|
|
for gram in self.gram_keys: |
|
|
if gram in self.train_steps: |
|
|
self.gram_graph_parts[gram].train(tc) |
|
|
|
|
|
if self.main_graph_part.key in self.train_steps: |
|
|
self.main_graph_part.train(tc) |
|
|
|
|
|
if self.lem_graph_part.key in self.train_steps: |
|
|
self.lem_graph_part.train(tc) |
|
|
|
|
|
if self.inflect_graph_part.key in self.train_steps: |
|
|
self.inflect_graph_part.train(tc) |
|
|
|
|
|
def release(self): |
|
|
with tf.Session(graph=self.graph) as sess: |
|
|
sess.run(tf.global_variables_initializer()) |
|
|
sess.run(tf.local_variables_initializer()) |
|
|
|
|
|
|
|
|
latest_checkpoint = tf.train.latest_checkpoint(self.save_path) |
|
|
if latest_checkpoint: |
|
|
self.restore(sess) |
|
|
|
|
|
if os.path.isdir(self.export_path): |
|
|
shutil.rmtree(self.export_path) |
|
|
|
|
|
output_dic = {} |
|
|
gram_op_dic = {} |
|
|
for gram in self.gram_keys: |
|
|
res = self.gram_graph_parts[gram].results[0] |
|
|
prob = self.gram_graph_parts[gram].probs[0] |
|
|
output_dic[f'res_{gram}'] = res |
|
|
output_dic[f'prob_{gram}'] = prob |
|
|
gram_op_dic[gram] = { |
|
|
'res': res.op.name, |
|
|
'prob': prob.op.name |
|
|
} |
|
|
|
|
|
output_dic['res_values'] = self.main_graph_part.results[0].values |
|
|
output_dic['res_indexes'] = self.main_graph_part.results[0].indices |
|
|
|
|
|
output_dic['lem_cls_result'] = self.lem_cls_result |
|
|
output_dic['lem_result'] = self.lem_result |
|
|
|
|
|
output_dic['inflect_result'] = self.inflect_graph_part.results[0] |
|
|
|
|
|
|
|
|
tf.saved_model.simple_save(sess, |
|
|
self.export_path, |
|
|
inputs={ |
|
|
'x_ind': self.x_inds[0], |
|
|
'x_val': self.x_vals[0], |
|
|
'x_shape': self.x_shape[0], |
|
|
'seq_len': self.x_seq_lens[0], |
|
|
'batch_size': self.batch_size, |
|
|
'lem_x_cls': self.lem_class_pl, |
|
|
'inflect_x_cls': self.inflect_x_class_pl, |
|
|
'inflect_y_cls': self.inflect_y_class_pl |
|
|
}, |
|
|
outputs=output_dic) |
|
|
|
|
|
|
|
|
input_graph = 'graph.pbtxt' |
|
|
tf.train.write_graph(sess.graph.as_graph_def(), self.export_path, input_graph, as_text=True) |
|
|
input_graph = os.path.join(self.export_path, input_graph) |
|
|
frozen_path = os.path.join(self.export_path, 'frozen_model.pb') |
|
|
output_ops = [output_dic[key].op.name for key in output_dic] |
|
|
output_ops = ",".join(output_ops) |
|
|
freeze_graph.freeze_graph(input_graph, |
|
|
"", |
|
|
False, |
|
|
latest_checkpoint, |
|
|
output_ops, |
|
|
"", |
|
|
"", |
|
|
frozen_path, |
|
|
True, |
|
|
"", |
|
|
input_saved_model_dir=self.export_path) |
|
|
|
|
|
op_dic = { |
|
|
key: output_dic[key].op.name |
|
|
for key in output_dic |
|
|
} |
|
|
|
|
|
op_dic['x_ind'] = self.x_inds[0].op.name |
|
|
op_dic['x_val'] = self.x_vals[0].op.name |
|
|
op_dic['x_shape'] = self.x_shape[0].op.name |
|
|
op_dic['seq_len'] = self.main_graph_part.x_seq_lens[0].op.name |
|
|
op_dic['batch_size'] = self.batch_size.op.name |
|
|
op_dic['lem_x_cls'] = self.lem_class_pl.op.name |
|
|
op_dic['inflect_x_cls'] = self.inflect_x_class_pl.op.name |
|
|
op_dic['inflect_y_cls'] = self.inflect_y_class_pl.op.name |
|
|
|
|
|
return frozen_path, \ |
|
|
gram_op_dic , \ |
|
|
op_dic |
|
|
|
