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')