|
|
|
@ -36,7 +36,6 @@ from ppocr.optimizer import build_optimizer
|
|
|
|
|
from ppocr.postprocess import build_post_process
|
|
|
|
|
from ppocr.metrics import build_metric
|
|
|
|
|
from ppocr.utils.save_load import init_model
|
|
|
|
|
from ppocr.utils.utility import print_dict
|
|
|
|
|
import tools.program as program
|
|
|
|
|
|
|
|
|
|
dist.get_world_size()
|
|
|
|
@ -61,7 +60,7 @@ def main(config, device, logger, vdl_writer):
|
|
|
|
|
global_config)
|
|
|
|
|
|
|
|
|
|
# build model
|
|
|
|
|
#for rec algorithm
|
|
|
|
|
# for rec algorithm
|
|
|
|
|
if hasattr(post_process_class, 'character'):
|
|
|
|
|
char_num = len(getattr(post_process_class, 'character'))
|
|
|
|
|
config['Architecture']["Head"]['out_channels'] = char_num
|
|
|
|
@ -81,10 +80,11 @@ def main(config, device, logger, vdl_writer):
|
|
|
|
|
|
|
|
|
|
# build metric
|
|
|
|
|
eval_class = build_metric(config['Metric'])
|
|
|
|
|
|
|
|
|
|
# load pretrain model
|
|
|
|
|
pre_best_model_dict = init_model(config, model, logger, optimizer)
|
|
|
|
|
|
|
|
|
|
logger.info('train dataloader has {} iter, valid dataloader has {} iter'.
|
|
|
|
|
format(len(train_dataloader), len(valid_dataloader)))
|
|
|
|
|
# start train
|
|
|
|
|
program.train(config, train_dataloader, valid_dataloader, device, model,
|
|
|
|
|
loss_class, optimizer, lr_scheduler, post_process_class,
|
|
|
|
@ -92,8 +92,7 @@ def main(config, device, logger, vdl_writer):
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_reader(config, device, logger):
|
|
|
|
|
loader = build_dataloader(config, 'Train', device)
|
|
|
|
|
# loader = build_dataloader(config, 'Eval', device)
|
|
|
|
|
loader = build_dataloader(config, 'Train', device, logger)
|
|
|
|
|
import time
|
|
|
|
|
starttime = time.time()
|
|
|
|
|
count = 0
|
|
|
|
|