1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
|
from glob import glob
import os
from tensorflow.keras.callbacks import ModelCheckpoint, TensorBoard
class LSTMChemTrainer(object):
def __init__(self, modeler, train_data_loader, valid_data_loader):
self.model = modeler.model
self.config = modeler.config
self.train_data_loader = train_data_loader
self.valid_data_loader = valid_data_loader
self.callbacks = []
self.init_callbacks()
def init_callbacks(self):
self.callbacks.append(
ModelCheckpoint(
filepath=os.path.join(
self.config.checkpoint_dir,
'%s-{epoch:02d}-{val_loss:.2f}.hdf5' %
self.config.exp_name),
monitor=self.config.checkpoint_monitor,
mode=self.config.checkpoint_mode,
save_best_only=self.config.checkpoint_save_best_only,
save_weights_only=self.config.checkpoint_save_weights_only,
verbose=self.config.checkpoint_verbose,
))
self.callbacks.append(
TensorBoard(
log_dir=self.config.tensorboard_log_dir,
write_graph=self.config.tensorboard_write_graph,
))
def train(self):
history = self.model.fit_generator(
self.train_data_loader,
steps_per_epoch=self.train_data_loader.__len__(),
epochs=self.config.num_epochs,
verbose=self.config.verbose_training,
validation_data=self.valid_data_loader,
validation_steps=self.valid_data_loader.__len__(),
use_multiprocessing=True,
shuffle=True,
callbacks=self.callbacks)
last_weight_file = glob(
os.path.join(
f'{self.config.checkpoint_dir}',
f'{self.config.exp_name}-{self.config.num_epochs:02}*.hdf5')
)[0]
assert os.path.exists(last_weight_file)
self.config.model_weight_filename = last_weight_file
with open(os.path.join(self.config.exp_dir, 'config.json'), 'w') as f:
f.write(self.config.toJSON(indent=2))
|
