-
Notifications
You must be signed in to change notification settings - Fork 8
/
train.py
60 lines (48 loc) · 1.95 KB
/
train.py
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
57
58
59
60
import importlib
import gin
import os
from tensorflow.python.keras.callbacks import CSVLogger, TensorBoard, ModelCheckpoint
@gin.configurable
class TrainConfig(object):
def __init__(self, args,
optimizer=None,
loss=None,
metrics=None,
batch_size=None,
epochs=None):
self.optimizer = optimizer
self.loss = loss
self.metrics = metrics
self.batch_size = batch_size
self.epochs = epochs
def train_model(data, args):
train_config = TrainConfig(args)
mode_module = importlib.import_module("modes." + args.mode)
train_generator = mode_module.DataGenerator(data.train_data)
val_generator = mode_module.DataGenerator(data.validation_data)
model = mode_module.build_model()
model.compile(
optimizer=train_config.optimizer,
loss=train_config.loss,
metrics=train_config.metrics
)
results_csv_file = os.path.join(args.results_dir, "results.csv")
ckpt_filename = "Epoch-{epoch:02d}-Val-Acc-{val_accuracy:.4f}.hdf5"
weight_file = os.path.join(args.checkpoints_dir, ckpt_filename)
results_callback = CSVLogger(results_csv_file, append=True, separator=',')
checkpoints_callback = ModelCheckpoint(weight_file,
save_best_only=True,
save_weights_only=True)
tensorboard_callback = TensorBoard(log_dir=args.results_dir,
histogram_freq=0, write_graph=True,
write_images=True)
model.fit_generator(
generator=train_generator,
validation_data=val_generator,
verbose=2,
epochs=train_config.epochs,
shuffle=True,
callbacks=[results_callback,
tensorboard_callback, checkpoints_callback]
)
return model