1 Star 0 Fork 7

liujialiang970207/LaTeX_OCR_PRO

forked from 大漠/LaTeX_OCR_PRO 
加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
克隆/下载
train_GAN.py 2.86 KB
一键复制 编辑 原始数据 按行查看 历史
ジ兮尘 提交于 2019-08-20 23:46 . remove torch
import click
from model.utils.data_generator import DataGenerator
from model.seqGAN import SeqGAN
from model.utils.lr_schedule import LRSchedule
from model.utils.general import Config
from model.utils.text import Vocab
from model.utils.image import greyscale
@click.command()
@click.option('--data', default="configs/data_small.json",
help='Path to data json config')
@click.option('--vocab', default="configs/vocab_small.json",
help='Path to vocab json config')
@click.option('--training', default="configs/training_small.json",
help='Path to training json config')
@click.option('--model', default="configs/model.json",
help='Path to model json config')
@click.option('--output', default="results/small/",
help='Dir for results and model weights')
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,
# iter_mode='formulas_only',
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,
# iter_mode='formulas_only',
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 = SeqGAN(config, dir_output, vocab)
model.build_train(config)
model.train(config, train_set, val_set, lr_schedule)
if __name__ == "__main__":
main()
Loading...
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
1
https://gitee.com/liujialiang970207/LaTeX_OCR_PRO.git
git@gitee.com:liujialiang970207/LaTeX_OCR_PRO.git
liujialiang970207
LaTeX_OCR_PRO
LaTeX_OCR_PRO
master

搜索帮助

0d507c66 1850385 C8b1a773 1850385