def train(self, dataset_train, dataset_val, learning_rate, epochs):
print("START TRAINING!")
#for l in self.rcnn_model.layers:
# print(l.name)
train_generator = data_generator.gen(dataset_train,
self.config,
shuffle=True,
batch_size=self.config.BATCH_SIZE)
val_generator = data_generator.gen(dataset_val,
self.config,
shuffle=True,
batch_size=self.config.BATCH_SIZE)
#display_data.view(dataset_train, self.config, shuffle=True, batch_size=self.config.BATCH_SIZE)
if not os.path.exists(self.log_dir):
os.makedirs(self.log_dir)
callbacks = [
keras.callbacks.TensorBoard(log_dir=self.log_dir,
histogram_freq=0,
write_graph=True,
write_images=False),
keras.callbacks.ModelCheckpoint(self.checkpoint_path,
verbose=0,
save_weights_only=True)
]
print("\nStarting at epoch {}. LR={}\n".format(self.epoch,
learning_rate))
print("Checkpoint Path: {}".format(self.checkpoint_path))
#set layers trainable
#self.set_trainable()
#since already load pretrained model, so we don't need train backbone
layers = r"(mrcnn\_.*)|(rpn\_.*)|(fpn\_.*)"
#layers=r"(res5.*)|(bn5.*)|(mrcnn\_.*)|(rpn\_.*)|(fpn\_.*)"
#layers = ".*"
self.set_trainable(layers)
self.compile(learning_rate, self.config.LEARNING_MOMENTUM)
workers = 0
self.rcnn_model.fit(train_generator,
initial_epoch=self.epoch,
epochs=epochs,
steps_per_epoch=self.config.STEPS_PER_EPOCH,
callbacks=callbacks,
validation_data=val_generator,
validation_steps=self.config.VALIDATION_STEPS,
max_queue_size=100,
workers=workers,
use_multiprocessing=True)
self.epoch = max(self.epoch, epochs)
def train_model():
east = East()
east_network = east.east_network()
east_network.summary()
east_network.compile(
loss=quad_loss,
optimizer=Nadam(
lr=cfg.lr,
# clipvalue=cfg.clipvalue,
schedule_decay=cfg.decay))
# load pre-trained model
if cfg.load_weights and os.path.exists(cfg.saved_model_weights_file_path):
east_network.load_weights(cfg.saved_model_weights_file_path,
by_name=True,
skip_mismatch=True)
print('start training task:' + cfg.train_task_id + '....................')
# train on current scale data
east_network.fit_generator(
generator=gen(),
steps_per_epoch=cfg.steps_per_epoch,
epochs=cfg.epoch_num,
validation_data=gen(is_val=True),
validation_steps=cfg.validation_steps,
verbose=2,
initial_epoch=cfg.initial_epoch,
callbacks=[
EarlyStopping(patience=cfg.patience, verbose=2),
ModelCheckpoint(filepath=cfg.saved_model_weights_file_path,
save_best_only=True,
save_weights_only=True,
verbose=1)
])
del east_network, east
import cfg
from network import East
from losses import quad_loss
from data_generator import gen
east = East()
east_network = east.east_network()
east_network.summary()
east_network.compile(loss=quad_loss, optimizer=Adam(lr=cfg.lr,
# clipvalue=cfg.clipvalue,
decay=cfg.decay))
if cfg.load_weights and os.path.exists(cfg.saved_model_weights_file_path):
east_network.load_weights(cfg.saved_model_weights_file_path)
east_network.fit_generator(generator=gen(),
steps_per_epoch=cfg.steps_per_epoch,
epochs=cfg.epoch_num,
validation_data=gen(is_val=True),
validation_steps=cfg.validation_steps,
verbose=1,
initial_epoch=cfg.initial_epoch,
callbacks=[
EarlyStopping(patience=cfg.patience, verbose=1),
ModelCheckpoint(filepath=cfg.model_weights_path,
save_best_only=True,
save_weights_only=True,
verbose=1)])
east_network.save(cfg.saved_model_file_path)
east_network.save(cfg.saved_model_weights_file_path)
print("=======Training using multiple GPUs..======")
except ValueError:
# parallel_model = east_network
print("=======Training using single GPU or CPU====")
east_network.summary()
east_network.compile(
loss=quad_loss,
optimizer=Adam(
lr=cfg.lr,
# clipvalue=cfg.clipvalue,
decay=cfg.decay))
if cfg.load_weights and os.path.exists(cfg.saved_model_weights_file_path):
east_network.load_weights(cfg.saved_model_weights_file_path)
east_network.fit_generator(
generator=gen(),
steps_per_epoch=cfg.steps_per_epoch,
epochs=cfg.epoch_num,
validation_data=gen(is_val=True),
validation_steps=cfg.validation_steps,
verbose=1,
initial_epoch=cfg.initial_epoch,
callbacks=[
EarlyStopping(patience=cfg.patience, verbose=1),
ModelCheckpoint(filepath=cfg.model_weights_path,
save_best_only=True,
save_weights_only=True,
verbose=1),
TensorBoard(
log_dir=cfg.log_dir, #使用tensorboard进行记录
histogram_freq=0,
east_network.compile(
loss=quad_loss,
optimizer=Adam(
lr=cfg.lr,
# clipvalue=cfg.clipvalue,
decay=cfg.decay))
# fine-tuning
if cfg.load_weights and os.path.exists(cfg.last_saved_model_weights_file_path):
print("ggggggggggggggggggggggggggggggggggggggggggggg")
print("load : " + cfg.last_saved_model_weights_file_path)
east_network.load_weights(
cfg.last_saved_model_weights_file_path) # 上一次加载权重参数
east_network.fit_generator(
generator=gen(), # 产生训练集
steps_per_epoch=cfg.steps_per_epoch,
epochs=cfg.epoch_num,
validation_data=gen(is_val=True), # 产生验证集
validation_steps=cfg.validation_steps,
verbose=1,
initial_epoch=cfg.initial_epoch, # 设置初始epoch数
callbacks=[
EarlyStopping(patience=cfg.patience, verbose=1),
ModelCheckpoint(filepath=cfg.model_weights_path,
save_best_only=True,
save_weights_only=True,
verbose=1)
])
east_network.save(cfg.saved_model_file_path)
east_network.save_weights(cfg.saved_model_weights_file_path)
