def __init__(self,config):
self.config=config['learning_config']
self.am = AM(config)
self.am.load_model(training=True)
if self.am.model_type!='MultiTask':
self.dg = AM_DataLoader(config)
else:
self.dg=MultiTask_DataLoader(config)
self.dg.speech_config['reduction_factor']=self.am.model.time_reduction_factor
self.dg.load_state(self.config['running_config']['outdir'])
if self.am.model_type=='CTC':
self.runner = ctc_runners.CTCTrainer(self.dg.speech_featurizer,self.dg.text_featurizer,self.config['running_config'])
elif self.am.model_type=='LAS':
self.runner=las_runners.LASTrainer(self.dg.speech_featurizer,self.dg.text_featurizer,self.config['running_config'])
self.dg.LAS=True
elif self.am.model_type == 'MultiTask':
self.runner = multi_runners.MultiTaskLASTrainer(self.dg.speech_featurizer, self.dg.token4_featurizer,
self.config['running_config'])
else:
self.runner = transducer_runners.TransducerTrainer(self.dg.speech_featurizer,self.dg.text_featurizer,self.config['running_config'] )
self.STT = self.am.model
if self.dg.augment.available():
factor=2
else:
factor=1
self.opt = tf.keras.optimizers.Adamax(**config['optimizer_config'])
self.runner.set_total_train_steps(self.dg.get_per_epoch_steps() * self.config['running_config']['num_epochs']*factor)
self.runner.compile(self.STT,self.opt)
self.dg.batch=self.runner.global_batch_size
class AM_Tester():
def __init__(self, config):
self.config = config['learning_config']
self.am = AM(config)
self.am.load_model(training=False)
if self.am.model_type != 'MultiTask':
self.dg = AM_DataLoader(config, training=False)
self.runner = am_tester.AMTester(
self.config['running_config'],
self.dg.text_featurizer,
streaming=config['speech_config']['streaming'])
else:
self.dg = MultiTask_DataLoader(config, training=False)
self.runner = multi_task_tester.MultiTaskTester(
self.config['running_config'], self.dg.token3_featurizer)
self.STT = self.am.model
self.runner.set_progbar(self.dg.eval_per_epoch_steps())
self.runner.set_all_steps(self.dg.eval_per_epoch_steps())
self.runner.compile(self.STT)
def make_eval_batch_data(self):
batches = []
for _ in range(
self.config['running_config']['eval_steps_per_batches']):
if self.am.model_type != 'MultiTask':
features, input_length, labels, label_length = self.dg.eval_data_generator(
)
input_length = np.expand_dims(input_length, -1)
batches.append((features, input_length, labels, label_length))
else:
speech_features, input_length, words_label, words_label_length, phone_label, phone_label_length, py_label, py_label_length = self.dg.eval_data_generator(
)
input_length = np.expand_dims(input_length, -1)
batches.append((speech_features, input_length, py_label))
return batches
def test(self):
while 1:
eval_batches = self.make_eval_batch_data()
# print('now',self.dg.offset)
self.runner.run(eval_batches)
if self.dg.offset > len(self.dg.test_list) - 1:
break
def __init__(self, config):
self.config = config['learning_config']
self.config['running_config'].update(
{'streaming': config['speech_config']['streaming']})
self.am = AM(config)
self.am.load_model(training=True)
if self.am.model_type != 'MultiTask':
self.dg = AM_DataLoader(config)
else:
self.dg = MultiTask_DataLoader(config)
self.dg.speech_config[
'reduction_factor'] = self.am.model.time_reduction_factor
self.dg.load_state(self.config['running_config']['outdir'])
if self.am.model_type == 'CTC':
self.runner = ctc_runners.CTCTrainer(self.dg.speech_featurizer,
self.dg.text_featurizer,
self.config['running_config'])
elif self.am.model_type == 'LAS':
self.runner = las_runners.LASTrainer(self.dg.speech_featurizer,
self.dg.text_featurizer,
self.config['running_config'])
self.dg.LAS = True
elif self.am.model_type == 'MultiTask':
self.runner = multi_runners.MultiTaskCTCTrainer(
self.dg.speech_featurizer, self.config['running_config'])
else:
self.runner = transducer_runners.TransducerTrainer(
self.dg.speech_featurizer, self.dg.text_featurizer,
self.config['running_config'])
self.STT = self.am.model
if self.dg.augment.available():
factor = 2
else:
factor = 1
all_train_step = self.dg.get_per_epoch_steps(
) * self.config['running_config']['num_epochs'] * factor
lr = CustomSchedule(config['model_config']['dmodel'],
warmup_steps=int(all_train_step * 0.1))
config['optimizer_config']['learning_rate'] = lr
self.opt = tf.keras.optimizers.Adam(**config['optimizer_config'])
self.runner.set_total_train_steps(all_train_step)
self.runner.compile(self.STT, self.opt)
self.dg.batch = self.runner.global_batch_size
def __init__(self,config):
self.config=config['learning_config']
self.am = AM(config)
self.am.load_model(training=False)
f,c=self.am.speech_feature.compute_feature_dim()
self.am.model.return_pb_function(f,c)
if self.am.model_type!='MultiTask':
self.dg = AM_DataLoader(config,training=False)
self.runner = am_tester.AMTester(self.config['running_config'], self.dg.text_featurizer)
else:
self.dg=MultiTask_DataLoader(config,training=False)
self.runner=multi_task_tester.MultiTaskTester(self.config['running_config'],self.dg.token3_featurizer,self.dg.token4_featurizer)
self.STT = self.am.model
self.runner.set_progbar(self.dg.eval_per_epoch_steps())
self.runner.compile(self.STT)
def __init__(self, config):
self.config = config['learning_config']
self.am = AM(config)
self.am.load_model(training=False)
if self.am.model_type != 'MultiTask':
self.dg = AM_DataLoader(config, training=False)
self.runner = am_tester.AMTester(
self.config['running_config'],
self.dg.text_featurizer,
streaming=config['speech_config']['streaming'])
else:
self.dg = MultiTask_DataLoader(config, training=False)
self.runner = multi_task_tester.MultiTaskTester(
self.config['running_config'], self.dg.token3_featurizer)
self.STT = self.am.model
self.runner.set_progbar(self.dg.eval_per_epoch_steps())
self.runner.set_all_steps(self.dg.eval_per_epoch_steps())
self.runner.compile(self.STT)
class AM_Trainer():
def __init__(self, config):
self.config = config['learning_config']
self.am = AM(config)
self.am.load_model(training=True)
if self.am.model_type != 'MultiTask':
self.dg = AM_DataLoader(config)
else:
self.dg = MultiTask_DataLoader(config)
self.dg.speech_config[
'reduction_factor'] = self.am.model.time_reduction_factor
self.dg.load_state(self.config['running_config']['outdir'])
if self.am.model_type == 'CTC':
self.runner = ctc_runners.CTCTrainer(self.dg.speech_featurizer,
self.dg.text_featurizer,
self.config['running_config'])
elif self.am.model_type == 'LAS':
self.runner = las_runners.LASTrainer(self.dg.speech_featurizer,
self.dg.text_featurizer,
self.config['running_config'])
self.dg.LAS = True
elif self.am.model_type == 'MultiTask':
self.runner = multi_runners.MultiTaskLASTrainer(
self.dg.speech_featurizer, self.dg.token4_featurizer,
self.config['running_config'])
else:
self.runner = transducer_runners.TransducerTrainer(
self.dg.speech_featurizer, self.dg.text_featurizer,
self.config['running_config'])
self.STT = self.am.model
if self.dg.augment.available():
factor = 2
else:
factor = 1
self.opt = tf.keras.optimizers.Adamax(**config['optimizer_config'])
self.runner.set_total_train_steps(
self.dg.get_per_epoch_steps() *
self.config['running_config']['num_epochs'] * factor)
self.runner.compile(self.STT, self.opt)
self.dg.batch = self.runner.global_batch_size
def load_checkpoint(self, config, model):
"""Load checkpoint."""
self.checkpoint_dir = os.path.join(
config['learning_config']['running_config']["outdir"],
"checkpoints")
files = os.listdir(self.checkpoint_dir)
files.sort(key=lambda x: int(x.split('_')[-1].replace('.h5', '')))
model.load_weights(os.path.join(self.checkpoint_dir, files[-1]))
self.init_steps = int(files[-1].split('_')[-1].replace('.h5', ''))
def recevie_data(self, r):
data = r.rpop(self.config['data_name'])
data = eval(data)
trains = []
for key in self.config['data_dict_key']:
x = data[key]
dtype = data['%s_dtype' % key]
shape = data['%s_shape' % key]
x = np.frombuffer(x, dtype)
x = x.reshape(shape)
trains.append(x)
return trains
def train(self):
if self.am.model_type != 'MultiTask':
train_datasets = tf.data.Dataset.from_generator(
self.dg.generator,
self.dg.return_data_types(),
self.dg.return_data_shape(),
args=(True, ))
eval_datasets = tf.data.Dataset.from_generator(
self.dg.generator,
self.dg.return_data_types(),
self.dg.return_data_shape(),
args=(False, ))
self.runner.set_datasets(train_datasets, eval_datasets)
else:
self.runner.set_datasets(self.dg.generator(True),
self.dg.generator(False))
while 1:
self.runner.fit(epoch=self.dg.epochs)
if self.runner._finished():
self.runner.save_checkpoint()
logging.info('Finish training!')
break
if self.runner.steps % self.config['running_config'][
'save_interval_steps'] == 0:
self.dg.save_state(self.config['running_config']['outdir'])
class AM_Trainer():
def __init__(self, config):
self.config = config['learning_config']
self.am = AM(config)
self.am.load_model(training=True)
if self.am.model_type != 'MultiTask':
self.dg = AM_DataLoader(config)
else:
self.dg = MultiTask_DataLoader(config)
self.dg.speech_config[
'reduction_factor'] = self.am.model.time_reduction_factor
self.dg.load_state(self.config['running_config']['outdir'])
if self.am.model_type == 'CTC':
self.runner = ctc_runners.CTCTrainer(self.dg.speech_featurizer,
self.dg.text_featurizer,
self.config['running_config'])
elif self.am.model_type == 'LAS':
self.runner = las_runners.LASTrainer(self.dg.speech_featurizer,
self.dg.text_featurizer,
self.config['running_config'])
self.dg.LAS = True
elif self.am.model_type == 'MultiTask':
self.runner = multi_runners.MultiTaskLASTrainer(
self.dg.speech_featurizer, self.dg.token4_featurizer,
self.config['running_config'])
else:
self.runner = transducer_runners.TransducerTrainer(
self.dg.speech_featurizer, self.dg.text_featurizer,
self.config['running_config'])
self.STT = self.am.model
if self.dg.augment.available():
factor = 2
else:
factor = 1
all_train_step = self.dg.get_per_epoch_steps(
) * self.config['running_config']['num_epochs'] * factor
lr = CustomSchedule(config['model_config']['dmodel'],
warmup_steps=int(all_train_step * 0.1))
config['optimizer_config']['learning_rate'] = lr
self.opt = tf.keras.optimizers.Adamax(**config['optimizer_config'])
self.runner.set_total_train_steps(all_train_step)
self.runner.compile(self.STT, self.opt)
self.dg.batch = self.runner.global_batch_size
def recevie_data(self, r):
data = r.rpop(self.config['data_name'])
data = eval(data)
trains = []
for key in self.config['data_dict_key']:
x = data[key]
dtype = data['%s_dtype' % key]
shape = data['%s_shape' % key]
x = np.frombuffer(x, dtype)
x = x.reshape(shape)
trains.append(x)
return trains
def train(self):
if self.am.model_type != 'MultiTask':
train_datasets = tf.data.Dataset.from_generator(
self.dg.generator,
self.dg.return_data_types(),
self.dg.return_data_shape(),
args=(True, ))
eval_datasets = tf.data.Dataset.from_generator(
self.dg.generator,
self.dg.return_data_types(),
self.dg.return_data_shape(),
args=(False, ))
self.runner.set_datasets(train_datasets, eval_datasets)
else:
self.runner.set_datasets(self.dg.generator(True),
self.dg.generator(False))
while 1:
self.runner.fit(epoch=self.dg.epochs)
if self.runner._finished():
self.runner.save_checkpoint()
logging.info('Finish training!')
break
if self.runner.steps % self.config['running_config'][
'save_interval_steps'] == 0:
self.dg.save_state(self.config['running_config']['outdir'])