def main(data, vocab, training, model, output): # Load configs dir_output = output config = Config([data, vocab, training, model]) config.save(dir_output) vocab = Vocab(config) # Load datasets train_set = DataGenerator(path_formulas=config.path_formulas_train, dir_images=config.dir_images_train, img_prepro=greyscale, max_iter=config.max_iter, bucket=config.bucket_train, path_matching=config.path_matching_train, max_len=config.max_length_formula, form_prepro=vocab.form_prepro) val_set = DataGenerator(path_formulas=config.path_formulas_val, dir_images=config.dir_images_val, img_prepro=greyscale, max_iter=config.max_iter, bucket=config.bucket_val, path_matching=config.path_matching_val, max_len=config.max_length_formula, form_prepro=vocab.form_prepro) # Define learning rate schedule n_batches_epoch = ((len(train_set) + config.batch_size - 1) // config.batch_size) lr_schedule = LRSchedule(lr_init=config.lr_init, start_decay=config.start_decay*n_batches_epoch, end_decay=config.end_decay*n_batches_epoch, end_warm=config.end_warm*n_batches_epoch, lr_warm=config.lr_warm, lr_min=config.lr_min) # Build model and train model = Img2SeqModel(config, dir_output, vocab) model.build_train(config) model.train(config, train_set, val_set, lr_schedule)
def main(results): # restore config and model dir_output = results config_data = Config(dir_output + "data.json") config_vocab = Config(dir_output + "vocab.json") config_model = Config(dir_output + "model.json") vocab = Vocab(config_vocab) model = Img2SeqModel(config_model, dir_output, vocab) model.build_pred() model.restore_session(dir_output + "model.weights/") # load dataset test_set = DataGenerator(path_formulas=config_data.path_formulas_test, dir_images=config_data.dir_images_test, img_prepro=greyscale, max_iter=config_data.max_iter, bucket=config_data.bucket_test, path_matching=config_data.path_matching_test, max_len=config_data.max_length_formula, form_prepro=vocab.form_prepro,) # build images from formulas formula_ref = path.join(dir_output, "formulas_test/ref.txt") formula_hyp = path.join(dir_output, "formulas_test/hyp_0.txt") images_ref = path.join(dir_output, "images_test/ref/") images_test = path.join(dir_output, "images_test/hyp_0/") build_images(load_formulas(formula_ref), images_ref) build_images(load_formulas(formula_hyp), images_test) # score the repositories scores = score_dirs(images_ref, images_test, greyscale) msg = " - ".join(["{} {:04.2f}".format(k, v) for k, v in scores.items()]) model.logger.info("- Eval Img: {}".format(msg))
def main(data, vocab): data_config = Config(data) # datasets train_set = DataGenerator(path_formulas=data_config.path_formulas_train, dir_images=data_config.dir_images_train, path_matching=data_config.path_matching_train) test_set = DataGenerator(path_formulas=data_config.path_formulas_test, dir_images=data_config.dir_images_test, path_matching=data_config.path_matching_test) val_set = DataGenerator(path_formulas=data_config.path_formulas_val, dir_images=data_config.dir_images_val, path_matching=data_config.path_matching_val) # produce images and matching files # train_set.build(buckets=None, n_threads=1) train_set.build(buckets=None) test_set.build(buckets=None) val_set.build(buckets=None) # train_set.build(buckets=data_config.buckets) # test_set.build(buckets=data_config.buckets) # val_set.build(buckets=data_config.buckets) # p = Augmentor.Pipeline(data_config.dir_images_train) # p.zoom(probability=1, min_factor=1.1, max_factor=1.5) # # p.process() # augmented_images, labels = p.sample(3) # print labels # vocab vocab_config = Config(vocab) vocab = build_vocab([train_set], min_count=vocab_config.min_count_tok) write_vocab(vocab, vocab_config.path_vocab)
def main(data, vocab): data_config = Config(data) # datasets train_set = DataGenerator( path_formulas=data_config.path_formulas_train, dir_images=data_config.dir_images_train, path_matching=data_config.path_matching_train) test_set = DataGenerator( path_formulas=data_config.path_formulas_test, dir_images=data_config.dir_images_test, path_matching=data_config.path_matching_test) val_set = DataGenerator( path_formulas=data_config.path_formulas_val, dir_images=data_config.dir_images_val, path_matching=data_config.path_matching_val) # produce images and matching files train_set.build(buckets=data_config.buckets) test_set.build(buckets=data_config.buckets) val_set.build(buckets=data_config.buckets) # vocab vocab_config = Config(vocab) vocab = build_vocab([train_set], min_count=vocab_config.min_count_tok) write_vocab(vocab, vocab_config.path_vocab)
def main(results): # restore config and model dir_output = results config_data = Config(dir_output + "data.json") config_vocab = Config(dir_output + "vocab.json") config_model = Config(dir_output + "model.json") vocab = Vocab(config_vocab) model = Img2SeqModel(config_model, dir_output, vocab) model.build_pred() model.restore_session(dir_output + "model.weights/") # load dataset test_set = DataGenerator( index_file=config_data.index_test, path_formulas=config_data.path_formulas_test, dir_images=config_data.dir_images_test, max_iter=config_data.max_iter, path_matching=config_data.path_matching_test, form_prepro=vocab.form_prepro) # use model to write predictions in files config_eval = Config({"dir_answers": dir_output + "formulas_test/", "batch_size": 20}) files, perplexity = model.write_prediction(config_eval, test_set) formula_ref, formula_hyp = files[0], files[1] # score the ref and prediction files scores = score_files(formula_ref, formula_hyp) scores["perplexity"] = perplexity msg = " - ".join(["{} {:04.2f}".format(k, v) for k, v in scores.items()]) model.logger.info("- Test Txt: {}".format(msg))
def main(results): tf.enable_resource_variables() # restore config and model dir_output = results weights_dir = os.path.join(dir_output, 'model.weights/') t = datetime.datetime.today().strftime('%Y-%m-%d-%H-%M-%S') # saved_path = 'saved_' + t saved_path = 'saved_word' saved_path = os.path.join(dir_output, saved_path) config_data = Config(dir_output + "data.json") config_vocab = Config(dir_output + "vocab.json") config_model = Config(dir_output + "model.json") vocab = Vocab(config_vocab) if not os.path.isdir(saved_path): model = Img2SeqModel(config_model, dir_output, vocab) model.build_pred() model.restore_session(weights_dir) model.save_savedmodel(saved_path) # chkp.print_tensors_in_checkpoint_file(weights_dir, tensor_name='', all_tensors=True) SAMPLE_DIR = 'tools/data/hand/raw_word' def representative_dataset_gen(): num_calibration_steps = 10 if not os.path.isdir(SAMPLE_DIR): print 'Failed to read representative_dataset' return for f in os.listdir(SAMPLE_DIR): img_path = os.path.join(SAMPLE_DIR, f) img = Image.open(img_path) img = img.resize((80, 100), Image.BILINEAR) img.show() img = np.array(img) yield [img] num_calibration_steps -= 1 if num_calibration_steps == 0: break converter = tf.lite.TFLiteConverter.from_saved_model(saved_path) converter.target_ops = [ # tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS ] # Following has "Segmentation fault" # converter.optimizations = [tf.lite.Optimize.DEFAULT] # converter.representative_dataset = representative_dataset_gen tflite_model = converter.convert() open("converted_model_word.tflite", "wb").write(tflite_model)
def getModelForPrediction(): # restore config and model dir_output = "./results/full/" config_vocab = Config(dir_output + "vocab.json") config_model = Config(dir_output + "model.json") vocab = Vocab(config_vocab) model = Img2SeqModel(config_model, dir_output, vocab) model.build_pred() # model.restore_session(dir_output + "model_weights/model.cpkt") return model
def main(data, vocab, training, model, output): # Load configs dir_output = output config = Config([data, vocab, training, model]) config.save(dir_output) vocab = Vocab(config) # Load datasets train_set = DataGenerator(path_formulas=config.path_formulas_train, dir_images=config.dir_images_train, img_prepro=greyscale, max_iter=config.max_iter, bucket=config.bucket_train, path_matching=config.path_matching_train, max_len=config.max_length_formula, form_prepro=vocab.form_prepro) all_img = [] all_formula = [] for i, (_img, _formula) in enumerate(minibatches(train_set, batch_size)): all_img.append(_img) if _formula is not None: _formula, _formula_length = pad_batch_formulas( _formula, vocab.id_pad, vocab.id_end ) all_formula.append(_formula) np.save('np_formula', np.array(all_formula)) np.save('np_img', np.array(all_img)) print("DONE EXPORTING NUMPY FILES") return None val_set = DataGenerator(path_formulas=config.path_formulas_val, dir_images=config.dir_images_val, img_prepro=greyscale, max_iter=config.max_iter, bucket=config.bucket_val, path_matching=config.path_matching_val, max_len=config.max_length_formula, form_prepro=vocab.form_prepro) # Define learning rate schedule n_batches_epoch = ((len(train_set) + config.batch_size - 1) // config.batch_size) lr_schedule = LRSchedule(lr_init=config.lr_init, start_decay=config.start_decay*n_batches_epoch, end_decay=config.end_decay*n_batches_epoch, end_warm=config.end_warm*n_batches_epoch, lr_warm=config.lr_warm, lr_min=config.lr_min) # Build model and train model = Img2SeqModel(config, dir_output, vocab) model.build_train(config) model.train(config, train_set, val_set, lr_schedule)
def load_model(dir_output="results/full/", vocab_config='vocab.json', model_config='model.json', model_path='model.weights/'): config_vocab = Config(os.path.join(dir_output, vocab_config)) config_model = Config(os.path.join(dir_output, model_config)) vocab = Vocab(config_vocab) model = Img2SeqModel(config_model, dir_output, vocab) model.build_pred() model.restore_session(os.path.join(dir_output, model_path)) return model
def main(data, vocab, training, model, output): # Load configs dir_output = output config = Config([data, vocab, training, model]) config.save(dir_output) vocab = Vocab(config) # Load datasets train_set = DataGenerator( index_file=config.index_train, path_formulas=config.path_formulas_train, dir_images=config.dir_images_train, max_iter=config.max_iter, path_matching=config.path_matching_train, max_len=config.max_length_formula, form_prepro=vocab.form_prepro) val_set = DataGenerator( index_file=config.index_val, path_formulas=config.path_formulas_val, dir_images=config.dir_images_val, max_iter=config.max_iter, path_matching=config.path_matching_val, max_len=config.max_length_formula, form_prepro=vocab.form_prepro) # Define learning rate schedule n_batches_epoch = ((len(train_set) + config.batch_size - 1) // config.batch_size) print len(train_set) print config.batch_size print n_batches_epoch lr_schedule = LRSchedule(lr_init=config.lr_init, start_decay=config.start_decay * n_batches_epoch, end_decay=config.end_decay * n_batches_epoch, end_warm=config.end_warm * n_batches_epoch, lr_warm=config.lr_warm, lr_min=config.lr_min) transfer_model = config.transfer_model # Build model and train model = Img2SeqModel(config, dir_output, vocab) model.build_train(config) if transfer_model and os.path.isdir(transfer_model): model.restore_session(transfer_model) model.train(config, train_set, val_set, lr_schedule)
def _run_epoch(self, config, train_set, val_set, epoch, lr_schedule): """Performs an epoch of training Args: config: Config instance train_set: Dataset instance val_set: Dataset instance epoch: (int) id of the epoch, starting at 0 lr_schedule: LRSchedule instance that takes care of learning proc Returns: score: (float) model will select weights that achieve the highest score """ # logging batch_size = config.batch_size nbatches = (len(train_set) + batch_size - 1) // batch_size prog = Progbar(nbatches) # iterate over dataset for i, (img, formula) in enumerate(minibatches(train_set, batch_size)): # get feed dict fd = self._get_feed_dict(img, training=True, formula=formula, lr=lr_schedule.lr, dropout=config.dropout) # update step _, loss_eval = self.sess.run([self.train_op, self.loss], feed_dict=fd) prog.update(i + 1, [("loss", loss_eval), ("perplexity", np.exp(loss_eval)), ("lr", lr_schedule.lr)]) # update learning rate lr_schedule.update(batch_no=epoch * nbatches + i) # 生成summary summary_str = self.sess.run(self.merged, feed_dict=fd) self.file_writer.add_summary(summary_str, epoch) # 将summary 写入文件 # if (i+1) % 100 == 0: # # 太慢了,读了 100 批次后就保存先,保存的权重要用于调试 attention # self.save_debug_session(epoch, i) # logging self.logger.info("- Training: {}".format(prog.info)) # evaluation config_eval = Config({ "dir_answers": self._dir_output + "formulas_val/", "batch_size": config.batch_size }) scores = self.evaluate(config_eval, val_set) score = scores[config.metric_val] lr_schedule.update(score=score) return score
def get_im2latex_model(weight_dir): """ Load up model from the given weight location :param weight_dir: weight location :return: trained model """ os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" tf.logging.set_verbosity(tf.logging.ERROR) config_vocab = Config(weight_dir + "vocab.json") config_model = Config(weight_dir + "model.json") vocab = Vocab(config_vocab) model = Img2SeqModel(config_model, weight_dir, vocab) model.build_pred() model.restore_session(weight_dir + "model.weights/") return model
def _run_train_epoch(self, config, train_set, val_set, epoch, lr_schedule): """Performs an epoch of training Args: config: Config instance train_set: Dataset instance val_set: Dataset instance epoch: (int) id of the epoch, starting at 0 lr_schedule: LRSchedule instance that takes care of learning proc Returns: score: (float) model will select weights that achieve the highest score """ # logging batch_size = config.batch_size nbatches = (len(train_set) + batch_size - 1) // batch_size prog = Progbar(nbatches) self.model.train() self.encoder.train() self.decoder.train() train_loader = torch.utils.data.DataLoader( ImgFormulaDataset(train_set), batch_size=batch_size, shuffle=True, num_workers=3, pin_memory=True) # for i, (img, formula) in enumerate(train_loader): for i, (img, formula) in enumerate(minibatches(train_set, batch_size)): img = pad_batch_images_2(img) img = torch.FloatTensor(img) # (N, W, H, C) formula, formula_length = pad_batch_formulas( formula, self._vocab.id_pad, self._vocab.id_end) img = img.permute(0, 3, 1, 2) # (N, C, W, H) formula = torch.LongTensor(formula) # (N,) loss_eval = self.getLoss(img, formula=formula, lr=lr_schedule.lr, dropout=config.dropout, training=True) prog.update(i + 1, [("loss", loss_eval), ("lr", lr_schedule.lr)]) # update learning rate lr_schedule.update(batch_no=epoch * nbatches + i) self.logger.info("- Training: {}".format(prog.info)) self.logger.info("- Config: (before evaluate, we need to see config)") config.show(fun=self.logger.info) # evaluation config_eval = Config({ "dir_answers": self._dir_output + "formulas_val/", "batch_size": config.batch_size }) scores = self.evaluate(config_eval, val_set) score = scores["perplexity"] lr_schedule.update(score=score) return score
def main(image, vocab, model, output): dir_output = output img_path = image # 加载配置,根据配置初始化字典和模型 config = Config([vocab, model]) vocab = Vocab(config) img2SeqModel = Img2SeqModel(config, dir_output, vocab) img2SeqModel.build_pred() vis_img_with_attention(img2SeqModel, img_path, dir_output)
def img2latex_api(weight_dir, img_path, downsample_image_ratio, cropping, padding, gray_scale): os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' tf.logging.set_verbosity(tf.logging.ERROR) config_vocab = Config(weight_dir + "vocab.json") config_model = Config(weight_dir + "model.json") vocab = Vocab(config_vocab) model = Img2SeqModel(config_model, weight_dir, vocab) model.build_pred() model.restore_session(weight_dir + "model.weights/") seq = iaa.Sequential([iaa.GammaContrast(2)]) latex, _, _ = img2latex(model, img_path, downsample_image_ratio=downsample_image_ratio, cropping=cropping, padding=padding, img_augment=seq, gray_scale=gray_scale) processed_latex = postprocess(latex) return processed_latex
def main(results): # restore config and model dir_output = results weights_dir = os.path.join(dir_output, 'model.weights/') saved_path = os.path.join(dir_output, 'saved') config_data = Config(dir_output + "data.json") config_vocab = Config(dir_output + "vocab.json") config_model = Config(dir_output + "model.json") vocab = Vocab(config_vocab) if not os.path.isdir(saved_path): model = Img2SeqModel(config_model, dir_output, vocab) model.build_pred() model.restore_session(weights_dir) model.save_savedmodel(saved_path) # chkp.print_tensors_in_checkpoint_file(weights_dir, tensor_name='', all_tensors=True) converter = tf.lite.TFLiteConverter.from_saved_model(saved_path) tflite_model = converter.convert() open("converted_model.tflite", "wb").write(tflite_model)
def predict(trained_model): """ Args: trained_model (Img2SeqModel): """ if len(sys.argv) <= 1: print "Usage" print "python predict.py abs/path/to/image.png" return image_path = sys.argv[1] img = imread(image_path) img = greyscale(img) res = trained_model.predict(img) trained_model.logger.info(res[0]) if __name__ == "__main__": # restore config and model dir_output = "results/full/" config_vocab = Config(dir_output + "vocab.json") config_model = Config(dir_output + "model.json") vocab = Vocab(config_vocab) model = Img2SeqModel(config_model, dir_output, vocab) model.build_pred() model.restore_session(dir_output + "model.weights/") # interactive_shell(model) predict(model)
from model.utils.general import Config, Progbar, minibatches # from model.evaluation.text import score_files, write_answers, truncate_end # _attn_cell_config = { # 'cell_type': 'lstm', # 'num_units': 12, # 'dim_e' : 14, # 'dim_o' : 16, # 'dim_embeddings': 32, # } dir_output = "results/small/" config = Config([ "configs/data_small.json", "configs/vocab_small.json", "configs/training_small.json", "configs/model.json", ]) config.save(dir_output) vocab = Vocab(config) train_set = DataGenerator( path_formulas=config.path_formulas_train, dir_images=config.dir_images_train, img_prepro=greyscale, max_iter=config.max_iter, bucket=config.bucket_train, path_matching=config.path_matching_train, max_len=config.max_length_formula, form_prepro=vocab.form_prepro )
def _run_train(self, config, train_set, val_set, epoch, lr_schedule): """Performs an epoch of training Args: config: Config instance train_set: Dataset instance val_set: Dataset instance epoch: (int) id of the epoch, starting at 0 lr_schedule: LRSchedule instance that takes care of learning proc Returns: score: (float) model will select weights that achieve the highest score """ # logging batch_size = config.batch_size train_set.shuffle() nbatches = (len(train_set) + batch_size - 1) // batch_size prog = Progbar(nbatches) # iterate over dataset for i, (img, formula) in enumerate(minibatches(train_set, batch_size)): # get feed dict fd = self._get_feed_dict(img, formula=formula, lr=lr_schedule.lr, dropout=config.dropout) # 来试试随机的 dropout # random_dropout = 0.5 + random.random() * 0.5 # fd = self._get_feed_dict(img, formula=formula, lr=lr_schedule.lr, dropout=random_dropout) # update step _, G_loss_eval = self.sess.run([self.train_op, self.loss], feed_dict=fd) if (G_loss_eval <= 1): # 等 Generator 的 loss 下降到 1 以下再训练判别器,不然判别器训练了感觉都没啥意义 _, D_loss_eval = self.sess.run([self.D_optimizer, self.D_loss], feed_dict=fd) else: D_loss_eval = 0 prog.update(i + 1, [("D_loss", D_loss_eval), ("G_loss", G_loss_eval), ("G_perplexity", np.exp(G_loss_eval)), ("lr", lr_schedule.lr)]) # update learning rate lr_schedule.update(batch_no=epoch*nbatches + i) # 生成summary if (i+1) % 10 == 0: summary_str = self.sess.run(self.merged, feed_dict=fd) self.file_writer.add_summary(summary_str, epoch) # 将summary 写入文件 # if (i+1) % 100 == 0: # # 太慢了,读了 100 批次后就保存先,保存的权重要用于调试 attention # self.save_debug_session(epoch, i) # logging self.logger.info("- Training: {}".format(prog.info)) # evaluation config_eval = Config({ "dir_answers": self._dir_output + "formulas_val/", "batch_size": config.batch_size }) scores = self.evaluate(config_eval, val_set) score = scores["perplexity"] + (scores["ExactMatchScore"] + scores["BLEU-4"] + scores["EditDistance"]) / 10 lr_schedule.update(score=score) return score
import os from model.img2seq import Img2SeqModel from model.utils.general import Config, run from model.utils.text import Vocab APP_ROOT = os.path.dirname( os.path.abspath(__file__)) # refers to application_top MODEL_FOLDER = os.path.join(APP_ROOT, 'full') print("loading model ...") config_vocab = Config(os.path.join(MODEL_FOLDER, "vocab.json")) config_model = Config(os.path.join(MODEL_FOLDER, "model.json")) vocab = Vocab(config_vocab) model = Img2SeqModel(config_model, "/tmp/mxlatexfull", vocab) model.build_pred() model.restore_session(os.path.join(MODEL_FOLDER, "model.weights/"))