def main(argv=None):
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list
# check if checkpoint path exists
if not os.path.exists(FLAGS.checkpoint_path):
os.mkdir(FLAGS.checkpoint_path)
else:
#if not FLAGS.restore:
# shutil.rmtree(FLAGS.checkpoint_path)
# os.mkdir(FLAGS.checkpoint_path)
shutil.rmtree(FLAGS.checkpoint_path)
os.mkdir(FLAGS.checkpoint_path)
train_data_generator = data_processor.generator(FLAGS)
train_samples_count = data_processor.count_samples(FLAGS)
val_data = data_processor.load_data(FLAGS)
if len(gpus) <= 1:
print('Training with 1 GPU')
if FLAGS.drn:
east = EAST_DRN_model(input_size=FLAGS.input_size)
else:
east = EAST_model(FLAGS.input_size)
parallel_model = east.model
else:
print('Training with %d GPUs' % len(gpus))
with tf.device("/cpu:0"):
east = EAST_model(FLAGS.input_size)
if FLAGS.restore_model is not '':
east.model.load_weights(FLAGS.restore_model)
parallel_model = multi_gpu_model(east.model, gpus=len(gpus))
score_map_loss_weight = K.variable(0.01, name='score_map_loss_weight')
small_text_weight = K.variable(0., name='small_text_weight')
lr_scheduler = LearningRateScheduler(lr_decay)
ckpt = CustomModelCheckpoint(model=east.model, path=FLAGS.checkpoint_path + '/model-{epoch:02d}.h5', period=FLAGS.save_checkpoint_epochs, save_weights_only=True)
tb = CustomTensorBoard(log_dir=FLAGS.checkpoint_path + '/train', score_map_loss_weight=score_map_loss_weight, small_text_weight=small_text_weight, data_generator=train_data_generator, write_graph=True)
small_text_weight_callback = SmallTextWeight(small_text_weight)
validation_evaluator = ValidationEvaluator(val_data, validation_log_dir=FLAGS.checkpoint_path + '/val')
callbacks = [lr_scheduler, ckpt, tb, small_text_weight_callback, validation_evaluator]
opt = AdamW(FLAGS.init_learning_rate)
parallel_model.compile(loss=[dice_loss(east.overly_small_text_region_training_mask, east.text_region_boundary_training_mask, score_map_loss_weight, small_text_weight),
rbox_loss(east.overly_small_text_region_training_mask, east.text_region_boundary_training_mask, small_text_weight, east.target_score_map)],
loss_weights=[1., 1.],
optimizer=opt)
east.model.summary()
model_json = east.model.to_json()
with open(FLAGS.checkpoint_path + '/model.json', 'w') as json_file:
json_file.write(model_json)
history = parallel_model.fit_generator(train_data_generator, epochs=FLAGS.max_epochs, steps_per_epoch=train_samples_count/FLAGS.batch_size, workers=FLAGS.nb_workers, use_multiprocessing=True, callbacks=callbacks, verbose=1)
def main(argv=None):
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list
# check if checkpoint path exists
if not os.path.exists(FLAGS.checkpoint_path):
os.mkdir(FLAGS.checkpoint_path)
else:
#if not FLAGS.restore:
# shutil.rmtree(FLAGS.checkpoint_path)
# os.mkdir(FLAGS.checkpoint_path)
shutil.rmtree(FLAGS.checkpoint_path)
os.mkdir(FLAGS.checkpoint_path)
# train_data_generator = data_processor.generator(FLAGS)
# train_samples_count = data_processor.count_samples(FLAGS)
train_data_generator = data_processor.DataGenerator(FLAGS)
# train_data_generator = data_processor.DataGenerator(FLAGS)
# val_data = data_processor.load_data(FLAGS)
val_data = data_processor.val_generator(FLAGS)
val_samples_count = data_processor.count_val_samples(FLAGS)
east = EAST_model(FLAGS.input_size)
if FLAGS.restore_model != '':
east.model.load_weights(FLAGS.restore_model)
score_map_loss_weight = K.variable(0.01, name='score_map_loss_weight')
small_text_weight = K.variable(0., name='small_text_weight')
lr_scheduler = LearningRateScheduler(lr_decay)
ckpt = CustomModelCheckpoint(model=east.model, path=FLAGS.checkpoint_path + '/model-{epoch:02d}.h5', period=FLAGS.save_checkpoint_epochs, save_weights_only=False)
small_text_weight_callback = SmallTextWeight(small_text_weight)
validation_evaluator = ValidationEvaluator(val_data, val_samples_count)
callbacks = [lr_scheduler, ckpt, small_text_weight_callback, validation_evaluator]
opt = tfa.optimizers.AdamW(weight_decay=1e-4, learning_rate=FLAGS.init_learning_rate, name='AdamW')
east.model.compile(
loss=[dice_loss(east.overly_small_text_region_training_mask, east.text_region_boundary_training_mask, score_map_loss_weight, small_text_weight),
rbox_loss(east.overly_small_text_region_training_mask, east.text_region_boundary_training_mask, small_text_weight, east.target_score_map)],
loss_weights=[1., 1.],
optimizer=opt,
# metrics=[['accuracy'], ['accuracy']]
)
east.model.summary()
model_json = east.model.to_json()
with open(FLAGS.checkpoint_path + '/model.json', 'w') as json_file:
json_file.write(model_json)
# history = east.model.fit(train_data_generator, epochs=FLAGS.max_epochs, steps_per_epoch=train_samples_count/FLAGS.batch_size, workers=FLAGS.nb_workers, use_multiprocessing=True, max_queue_size=10, callbacks=callbacks, verbose=1)
history = east.model.fit(train_data_generator, epochs=FLAGS.max_epochs, workers=FLAGS.nb_workers, use_multiprocessing=True, max_queue_size=10, callbacks=callbacks, verbose=1)
east.model.save('east_retrained.h5')
def main():
train_data_generator = DataGenerator(input_size=FLAGS.input_size, batch_size=FLAGS.batch_size,
data_path=FLAGS.training_data_path, FLAGS=FLAGS, is_train=True)
train_samples_count = len(train_data_generator.image_paths)
validation_data_generator = DataGenerator(input_size=FLAGS.input_size, batch_size=FLAGS.batch_size,
data_path=FLAGS.validation_data_path, FLAGS=FLAGS, is_train=False)
east = EastModel(FLAGS.input_size)
if FLAGS.pretrained_weights_path != '':
print(f'Loading pre-trained model at {FLAGS.pretrained_weights_path}')
east.model.load_weights(FLAGS.pretrained_weights_path)
score_map_loss_weight = K.variable(0.01, name='score_map_loss_weight')
small_text_weight = K.variable(0., name='small_text_weight')
opt = AdamW(FLAGS.init_learning_rate)
east.model.compile(
loss=[
dice_loss(east.overly_small_text_region_training_mask, east.text_region_boundary_training_mask,
score_map_loss_weight, small_text_weight),
rbox_loss(east.overly_small_text_region_training_mask, east.text_region_boundary_training_mask,
small_text_weight, east.target_score_map)
],
loss_weights=[1., 1.],
optimizer=opt,
)
tb_callback = tensorboard_callback()
cp_callback = checkpoint_callback()
with open(os.path.join(FLAGS.checkpoint_path, 'model.json'), 'w') as json_file:
json_file.write(east.model.to_json())
east.model.fit_generator(
generator=train_data_generator,
epochs=FLAGS.max_epochs,
steps_per_epoch=train_samples_count // FLAGS.batch_size,
validation_data=validation_data_generator,
callbacks=[cp_callback, tb_callback],
workers=FLAGS.nb_workers,
use_multiprocessing=True,
max_queue_size=10,
verbose=1,
)
def train_step(model,
x,
optimizer,
overly_small_text_region_training_mask,
text_region_boundary_training_mask,
small_text_weight,
target_score_map,
target_geo_maps,
loss_weight):
with tf.GradientTape() as tape:
score_y_pred, geo_y_pred = model(x)
_dice_loss = dice_loss(overly_small_text_region_training_mask, text_region_boundary_training_mask, loss_weight, small_text_weight, target_score_map, score_y_pred)
_rbox_loss = rbox_loss(overly_small_text_region_training_mask, text_region_boundary_training_mask, small_text_weight, target_score_map, target_geo_maps, geo_y_pred)
loss = _dice_loss + _rbox_loss
gradients = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(gradients, model.trainable_variables))
def main(argv=None):
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list
# check if checkpoint path exists
if not os.path.exists(FLAGS.checkpoint_path):
os.mkdir(FLAGS.checkpoint_path)
else:
#if not FLAGS.restore:
# shutil.rmtree(FLAGS.checkpoint_path)
# os.mkdir(FLAGS.checkpoint_path)
shutil.rmtree(FLAGS.checkpoint_path)
os.mkdir(FLAGS.checkpoint_path)
if len(gpus) <= 1:
print('Training with 1 GPU')
east = EAST_model(FLAGS.input_size)
parallel_model = east.model
else:
print('Training with %d GPUs' % len(gpus))
with tf.device("/cpu:0"):
east = EAST_model(FLAGS.input_size)
if FLAGS.restore_model is not '':
east.model.load_weights(FLAGS.restore_model)
parallel_model = multi_gpu_model(east.model, gpus=len(gpus))
score_map_loss_weight = K.variable(0.01, name='score_map_loss_weight')
small_text_weight = K.variable(0., name='small_text_weight')
# opt = AdamW(FLAGS.init_learning_rate)
opt = Adam(learning_rate=0.001, beta_1=0.9, beta_2=0.999, epsilon=1e-4)
parallel_model.compile(loss=[
dice_loss(east.overly_small_text_region_training_mask,
east.text_region_boundary_training_mask,
score_map_loss_weight, small_text_weight),
rbox_loss(east.overly_small_text_region_training_mask,
east.text_region_boundary_training_mask, small_text_weight,
east.target_score_map)
],
loss_weights=[1., 1.],
optimizer=opt)
east.model.summary()
model_json = east.model.to_json()
with open(FLAGS.checkpoint_path + '/model.json', 'w') as json_file:
json_file.write(model_json)
plot_model(east.model,
to_file=FLAGS.checkpoint_path + '/east-train.png',
show_shapes=True,
show_layer_names=True,
expand_nested=True)
plot_model(east.resnet,
to_file=FLAGS.checkpoint_path + '/resnet.png',
show_shapes=True,
show_layer_names=True,
expand_nested=True)
train_data_generator = data_processor.TrainDataSequence(FLAGS)
#train_data_generator = data_processor.generator(FLAGS)
#train_data_x, train_data_y = data_processor.load_train_data(FLAGS)
#print('# of train data: %d' %(len(train_data_x[0])))
train_samples_count = data_processor.count_samples(FLAGS)
print("train_samples_count: %d" % (train_samples_count))
val_data = data_processor.load_data(FLAGS)
lr_scheduler = LearningRateScheduler(lr_decay)
ckpt = CustomModelCheckpoint(model=east.model,
path=FLAGS.checkpoint_path,
period=FLAGS.save_checkpoint_epochs,
save_weights_only=False)
# tb = CustomTensorBoard(log_dir=FLAGS.checkpoint_path + '/train', score_map_loss_weight=score_map_loss_weight, small_text_weight=small_text_weight, data_generator=train_data_generator, write_graph=True)
small_text_weight_callback = SmallTextWeight(small_text_weight)
validation_evaluator = ValidationEvaluator(
val_data, validation_log_dir=FLAGS.checkpoint_path + '/val')
callbacks = [
lr_scheduler, ckpt, small_text_weight_callback, validation_evaluator
]
#callbacks = [lr_scheduler, ckpt, small_text_weight_callback]
#history = parallel_model.fit(x=train_data_x, y=train_data_y, batch_size=FLAGS.batch_size, epochs=FLAGS.max_epochs, verbose=1, callbacks=callbacks, max_queue_size=10, workers=FLAGS.nb_workers, use_multiprocessing=True)
#history = parallel_model.fit(x=train_data_generator, epochs=FLAGS.max_epochs, steps_per_epoch=train_samples_count/FLAGS.batch_size, callbacks=callbacks, max_queue_size=10, verbose=1)
history = parallel_model.fit(x=train_data_generator,
epochs=FLAGS.max_epochs,
callbacks=callbacks,
max_queue_size=10,
workers=FLAGS.nb_workers,
use_multiprocessing=False,
verbose=1)
file_name = FLAGS.checkpoint_path + '/model-train.h5'
east.model.save(file_name, overwrite=True)
plot_model(east.model_core,
to_file=FLAGS.checkpoint_path + '/east.png',
show_shapes=True,
show_layer_names=True,
expand_nested=True)
file_name = FLAGS.checkpoint_path + '/model.h5'
east.model_core.save(file_name, overwrite=True, include_optimizer=False)
tflite_model = lite.TFLiteConverter.from_keras_model_file(
file_name,
input_arrays=["input_image"],
input_shapes={
"input_image": [1, 224, 320, 3]
}).convert()
with open(file_name + '.tflite', 'wb') as tflite_file:
tflite_file.write(tflite_model)
def main(_):
# check if checkpoint path exists
if not os.path.exists(FLAGS.checkpoint_path):
os.mkdir(FLAGS.checkpoint_path)
train_data_generator = data_processor.generator(FLAGS)
train_samples_count = data_processor.count_samples(FLAGS)
print('total batches per epoch : {}'.format(train_samples_count / FLAGS.batch_size))
east = EAST_model(FLAGS.input_size)
east.model.summary()
score_map_loss_weight = tf.Variable(0.01, name='score_map_loss_weight')
small_text_weight = tf.Variable(0., name='small_text_weight')
lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
FLAGS.init_learning_rate,
decay_steps=FLAGS.lr_decay_steps,
decay_rate=FLAGS.lr_decay_rate,
staircase=True)
optimizer = tf.optimizers.Adam(lr_schedule)
# set checkpoint manager
ckpt = tf.train.Checkpoint(step=tf.Variable(0), model=east)
ckpt_manager = tf.train.CheckpointManager(ckpt,
directory=FLAGS.checkpoint_path,
max_to_keep=5)
latest_ckpt = tf.train.latest_checkpoint(FLAGS.checkpoint_path)
# restore latest checkpoint
if latest_ckpt:
ckpt.restore(latest_ckpt)
print('global_step : {}, checkpoint is restored!'.format(int(ckpt.step)))
# set tensorboard summary writer
summary_writer = tf.summary.create_file_writer(FLAGS.checkpoint_path + '/train')
while int(ckpt.step) < (FLAGS.max_steps + 1):
# load data
[input_images, overly_small_text_region_training_masks, text_region_boundary_training_masks, score_maps], \
[target_score_maps, target_geo_maps] = next(train_data_generator)
# update parameter
train_step(east,
input_images,
optimizer,
overly_small_text_region_training_masks,
text_region_boundary_training_masks,
small_text_weight,
target_score_maps,
target_geo_maps,
score_map_loss_weight
)
score_y_pred, geo_y_pred = east(input_images)
_dice_loss = dice_loss(overly_small_text_region_training_masks, text_region_boundary_training_masks, score_map_loss_weight,
small_text_weight, target_score_maps, score_y_pred)
_rbox_loss = rbox_loss(overly_small_text_region_training_masks, text_region_boundary_training_masks,
small_text_weight, target_score_maps, target_geo_maps, geo_y_pred)
loss = _dice_loss + _rbox_loss
print('Step {:06d}, dice_loss {:.4f}, rbox_loss {:.4f}, total_loss {:.4f}'.format(int(ckpt.step), _dice_loss, _rbox_loss, loss))
if ckpt.step % FLAGS.save_checkpoint_steps == 0:
# save checkpoint
ckpt_manager.save(checkpoint_number=ckpt.step)
print('global_step : {}, checkpoint is saved!'.format(int(ckpt.step)))
with summary_writer.as_default():
tf.summary.scalar('loss', loss, step=int(ckpt.step))
tf.summary.scalar('pred_score_map_loss', _dice_loss, step=int(ckpt.step))
tf.summary.scalar('pred_geo_map_loss ', _rbox_loss, step=int(ckpt.step))
tf.summary.scalar('learning_rate ', optimizer.lr(ckpt.step).numpy(), step=int(ckpt.step))
tf.summary.scalar('small_text_weight', small_text_weight, step=int(ckpt.step))
tf.summary.image("input_image", tf.cast((input_images + 1) * 127.5, tf.uint8), step=int(ckpt.step), max_outputs=3)
tf.summary.image("overly_small_text_region_training_mask", tf.cast(overly_small_text_region_training_masks * 255, tf.uint8), step=int(ckpt.step), max_outputs=3)
tf.summary.image("text_region_boundary_training_mask", tf.cast(text_region_boundary_training_masks * 255, tf.uint8), step=int(ckpt.step), max_outputs=3)
tf.summary.image("score_map_target", tf.cast(target_score_maps * 255, tf.uint8), step=int(ckpt.step), max_outputs=3)
tf.summary.image("score_map_pred", tf.cast(score_y_pred * 255, tf.uint8), step=int(ckpt.step), max_outputs=3)
for i in range(4):
tf.summary.image("geo_map_%d_target" % (i), tf.cast(tf.expand_dims(target_geo_maps[:, :, :, i], axis=3) / FLAGS.input_size * 255, tf.uint8), step=int(ckpt.step), max_outputs=3)
tf.summary.image("geo_map_%d_pred" % (i), tf.cast(tf.expand_dims(geo_y_pred[:, :, :, i], axis=3) / FLAGS.input_size * 255, tf.uint8), step=int(ckpt.step), max_outputs=3)
tf.summary.image("geo_map_4_target", tf.cast((tf.expand_dims(target_geo_maps[:, :, :, 4], axis=3) + 1) * 127.5, tf.uint8), step=int(ckpt.step), max_outputs=3)
tf.summary.image("geo_map_4_pred", tf.cast((tf.expand_dims(geo_y_pred[:, :, :, 4], axis=3) + 1) * 127.5, tf.uint8), step=int(ckpt.step), max_outputs=3)
ckpt.step.assign_add(1)
def main():
os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu_list
# check if checkpoint path exists
if not os.path.exists(FLAGS.checkpoint_path):
os.mkdir(FLAGS.checkpoint_path)
else:
# if not FLAGS.restore:
# shutil.rmtree(FLAGS.checkpoint_path)
# os.mkdir(FLAGS.checkpoint_path)
shutil.rmtree(FLAGS.checkpoint_path)
os.mkdir(FLAGS.checkpoint_path)
# train_data_generator = data_processor.generator(FLAGS)
train_data_generator = EastSequence(FLAGS, input_size=FLAGS.input_size)
train_samples_count = data_processor.count_samples(FLAGS)
val_data = data_processor.load_data(FLAGS)
train_graph = tf.Graph()
train_sess = tf.Session(graph=train_graph)
K.set_session(train_sess)
with train_graph.as_default():
K.set_learning_phase(0)
if len(gpus) <= 1:
print('Training with 1 GPU')
east = EASTModel(FLAGS.input_size)
parallel_model = east.model
else:
print('Training with %d GPUs' % len(gpus))
with tf.device("/cpu:0"):
east = EASTModel(FLAGS.input_size)
parallel_model = multi_gpu_model(east.model, gpus=len(gpus))
score_map_loss_weight = K.variable(0.01, name='score_map_loss_weight')
small_text_weight = K.variable(0., name='small_text_weight')
tf.contrib.quantize.create_training_graph(input_graph=train_graph,
quant_delay=250000)
train_sess.run(tf.global_variables_initializer())
lr_scheduler = LearningRateScheduler(lr_decay)
ckpt = CustomModelCheckpoint(model=east.model,
path=FLAGS.checkpoint_path +
'/model-{epoch:02d}.h5',
period=FLAGS.save_checkpoint_epochs,
save_weights_only=True)
tb = TensorBoard(log_dir=FLAGS.checkpoint_path + '/train',
batch_size=FLAGS.batch_size,
write_images=True)
small_text_weight_callback = SmallTextWeight(small_text_weight)
validation_evaluator = ValidationEvaluator(
val_data, validation_log_dir=FLAGS.checkpoint_path + '/val')
callbacks = [
lr_scheduler, ckpt, tb, small_text_weight_callback,
validation_evaluator
]
# opt = AdamW(FLAGS.init_learning_rate)
parallel_model.compile(loss=[
dice_loss(east.overly_small_text_region_training_mask,
east.text_region_boundary_training_mask,
score_map_loss_weight, small_text_weight),
rbox_loss(east.overly_small_text_region_training_mask,
east.text_region_boundary_training_mask,
small_text_weight, east.target_score_map)
],
loss_weights=[1., 1.],
optimizer='adam')
model_json = east.model.to_json()
with open(FLAGS.checkpoint_path + '/model.json', 'w') as json_file:
json_file.write(model_json)
history = parallel_model.fit_generator(
train_data_generator,
epochs=FLAGS.max_epochs,
steps_per_epoch=train_samples_count / FLAGS.batch_size,
workers=FLAGS.nb_workers,
use_multiprocessing=True,
max_queue_size=10,
callbacks=callbacks,
verbose=1)
saver = tf.train.Saver()
saver.save(train_sess, 'checkpoints')